19166: Close connections to container gateway when finished.

[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 d1d512856ac9744f091c2152985a7ebc5fd81651..ec16ee2be9d9f345810274d252acec579eaf7ddd 100644 (file)
--- a/sdk/go/arvados/container_gateway.go
+++ b/sdk/go/arvados/container_gateway.go
@@ -15,13 +15,16 @@ import (
 )
 
 func (cresp ConnectionResponse) ServeHTTP(w http.ResponseWriter, req *http.Request) {
+       defer cresp.Conn.Close()
        hj, ok := w.(http.Hijacker)
        if !ok {
                http.Error(w, "ResponseWriter does not support connection upgrade", http.StatusInternalServerError)
                return
        }
        w.Header().Set("Connection", "upgrade")
-       w.Header().Set("Upgrade", cresp.UpgradeHeader)
+       for k, v := range cresp.Header {
+               w.Header()[k] = v
+       }
        w.WriteHeader(http.StatusSwitchingProtocols)
        conn, bufrw, err := hj.Hijack()
        if err != nil {
@@ -32,7 +35,7 @@ func (cresp ConnectionResponse) ServeHTTP(w http.ResponseWriter, req *http.Reque
 
        var bytesIn, bytesOut int64
        var wg sync.WaitGroup
-       ctx, cancel := context.WithCancel(context.Background())
+       ctx, cancel := context.WithCancel(req.Context())
        wg.Add(1)
        go func() {
                defer wg.Done()
@@ -44,8 +47,9 @@ func (cresp ConnectionResponse) ServeHTTP(w http.ResponseWriter, req *http.Reque
                        bytesOut += n
                }
                if err != nil {
-                       ctxlog.FromContext(req.Context()).WithError(err).Error("error copying downstream")
+                       ctxlog.FromContext(ctx).WithError(err).Error("error copying downstream")
                }
+               conn.Close()
        }()
        wg.Add(1)
        go func() {
@@ -58,17 +62,15 @@ func (cresp ConnectionResponse) ServeHTTP(w http.ResponseWriter, req *http.Reque
                        bytesIn += n
                }
                if err != nil {
-                       ctxlog.FromContext(req.Context()).WithError(err).Error("error copying upstream")
+                       ctxlog.FromContext(ctx).WithError(err).Error("error copying upstream")
                }
+               cresp.Conn.Close()
        }()
-       <-ctx.Done()
+       wg.Wait()
        if cresp.Logger != nil {
-               go func() {
-                       wg.Wait()
-                       cresp.Logger.WithFields(logrus.Fields{
-                               "bytesIn":  bytesIn,
-                               "bytesOut": bytesOut,
-                       }).Info("closed connection")
-               }()
+               cresp.Logger.WithFields(logrus.Fields{
+                       "bytesIn":  bytesIn,
+                       "bytesOut": bytesOut,
+               }).Info("closed connection")
        }
 }