task engine for all proxies

proxy/http/server.go

@@ -17,9 +17,9 @@ import (
 	"v2ray.com/core/common/log"
 	v2net "v2ray.com/core/common/net"
 	"v2ray.com/core/common/serial"
+	"v2ray.com/core/common/signal"
 	"v2ray.com/core/proxy"
 	"v2ray.com/core/transport/internet"
-	"v2ray.com/core/transport/ray"
 )
 
 // Server is a HTTP proxy server.
@@ -155,35 +155,32 @@ func (v *Server) handleConnect(request *http.Request, session *proxy.SessionInfo
 	}
 
 	ray := v.packetDispatcher.DispatchToOutbound(session)
-	v.transport(reader, writer, ray)
-}
 
-func (v *Server) transport(input io.Reader, output io.Writer, ray ray.InboundRay) {
-	var wg sync.WaitGroup
-	wg.Add(2)
-	defer wg.Wait()
+	requestDone := signal.ExecuteAsync(func() error {
+		defer ray.InboundInput().Close()
 
-	go func() {
-		v2reader := buf.NewReader(input)
+		v2reader := buf.NewReader(reader)
 		defer v2reader.Release()
 
 		if err := buf.PipeUntilEOF(v2reader, ray.InboundInput()); err != nil {
-			log.Info("HTTP: Failed to transport all TCP request: ", err)
+			return err
 		}
-		ray.InboundInput().Close()
-		wg.Done()
-	}()
+		return nil
+	})
 
-	go func() {
-		v2writer := buf.NewWriter(output)
+	responseDone := signal.ExecuteAsync(func() error {
+		defer ray.InboundOutput().Release()
+
+		v2writer := buf.NewWriter(writer)
 		defer v2writer.Release()
 
 		if err := buf.PipeUntilEOF(ray.InboundOutput(), v2writer); err != nil {
-			log.Info("HTTP: Failed to transport all TCP response: ", err)
+			return err
 		}
-		ray.InboundOutput().Release()
-		wg.Done()
-	}()
+		return nil
+	})
+
+	signal.ErrorOrFinish2(requestDone, responseDone)
 }
 
 // @VisibleForTesting
@@ -239,27 +236,26 @@ func (v *Server) handlePlainHTTP(request *http.Request, session *proxy.SessionIn
 	StripHopByHopHeaders(request)
 
 	ray := v.packetDispatcher.DispatchToOutbound(session)
-	defer ray.InboundInput().Close()
-	defer ray.InboundOutput().Release()
 
-	var finish sync.WaitGroup
-	finish.Add(1)
-	go func() {
-		defer finish.Done()
+	requestDone := signal.ExecuteAsync(func() error {
+		defer ray.InboundInput().Close()
+
 		requestWriter := bufio.NewWriter(buf.NewBytesWriter(ray.InboundInput()))
 		defer requestWriter.Release()
 
 		err := request.Write(requestWriter)
 		if err != nil {
-			log.Warning("HTTP: Failed to write request: ", err)
-			return
+			return err
 		}
-		requestWriter.Flush()
-	}()
+		if err := requestWriter.Flush(); err != nil {
+			return err
+		}
+		return nil
+	})
+
+	responseDone := signal.ExecuteAsync(func() error {
+		defer ray.InboundOutput().Release()
 
-	finish.Add(1)
-	go func() {
-		defer finish.Done()
 		responseReader := bufio.OriginalReader(buf.NewBytesReader(ray.InboundOutput()))
 		response, err := http.ReadResponse(responseReader, request)
 		if err != nil {
@@ -267,14 +263,19 @@ func (v *Server) handlePlainHTTP(request *http.Request, session *proxy.SessionIn
 			response = v.GenerateResponse(503, "Service Unavailable")
 		}
 		responseWriter := bufio.NewWriter(writer)
-		err = response.Write(responseWriter)
-		if err != nil {
-			log.Warning("HTTP: Failed to write response: ", err)
-			return
+		if err := response.Write(responseWriter); err != nil {
+			return err
 		}
-		responseWriter.Flush()
-	}()
-	finish.Wait()
+
+		if err := responseWriter.Flush(); err != nil {
+			return err
+		}
+		return nil
+	})
+
+	if err := signal.ErrorOrFinish2(requestDone, responseDone); err != nil {
+		log.Info("HTTP|Server: Connecton ending with ", err)
+	}
 }
 
 // ServerFactory is a InboundHandlerFactory.