19166: Merge branch 'main'

[arvados.git] / sdk / go / arvados / container_gateway.go

diff --git a/sdk/go/arvados/container_gateway.go b/sdk/go/arvados/container_gateway.go
index ec16ee2be9d9f345810274d252acec579eaf7ddd..897ae434e14bd0a0392d041a125a598b2c1d8b34 100644 (file)
--- a/sdk/go/arvados/container_gateway.go
+++ b/sdk/go/arvados/container_gateway.go
@@ -34,8 +34,8 @@ func (cresp ConnectionResponse) ServeHTTP(w http.ResponseWriter, req *http.Reque
        defer conn.Close()
 
        var bytesIn, bytesOut int64
-       var wg sync.WaitGroup
        ctx, cancel := context.WithCancel(req.Context())
+       var wg sync.WaitGroup
        wg.Add(1)
        go func() {
                defer wg.Done()
@@ -49,7 +49,6 @@ func (cresp ConnectionResponse) ServeHTTP(w http.ResponseWriter, req *http.Reque
                if err != nil {
                        ctxlog.FromContext(ctx).WithError(err).Error("error copying downstream")
                }
-               conn.Close()
        }()
        wg.Add(1)
        go func() {
@@ -64,13 +63,17 @@ func (cresp ConnectionResponse) ServeHTTP(w http.ResponseWriter, req *http.Reque
                if err != nil {
                        ctxlog.FromContext(ctx).WithError(err).Error("error copying upstream")
                }
-               cresp.Conn.Close()
        }()
-       wg.Wait()
-       if cresp.Logger != nil {
-               cresp.Logger.WithFields(logrus.Fields{
-                       "bytesIn":  bytesIn,
-                       "bytesOut": bytesOut,
-               }).Info("closed connection")
-       }
+       <-ctx.Done()
+       go func() {
+               // Wait for both io.Copy goroutines to finish and increment
+               // their byte counters.
+               wg.Wait()
+               if cresp.Logger != nil {
+                       cresp.Logger.WithFields(logrus.Fields{
+                               "bytesIn":  bytesIn,
+                               "bytesOut": bytesOut,
+                       }).Info("closed connection")
+               }
+       }()
 }