X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/56e130608f8977d20b21c54f6ab8973d71e045a0..79870ba994f0606c8ed13806f00cb8b23d9b2c83:/sdk/go/arvados/container_gateway.go

diff --git a/sdk/go/arvados/container_gateway.go b/sdk/go/arvados/container_gateway.go
index 07f8c0793c..00c98d572e 100644
--- a/sdk/go/arvados/container_gateway.go
+++ b/sdk/go/arvados/container_gateway.go
@@ -8,8 +8,10 @@ import (
 	"context"
 	"io"
 	"net/http"
+	"sync"
 
 	"git.arvados.org/arvados.git/sdk/go/ctxlog"
+	"github.com/sirupsen/logrus"
 )
 
 func (sshconn ContainerSSHConnection) ServeHTTP(w http.ResponseWriter, req *http.Request) {
@@ -28,26 +30,45 @@ func (sshconn ContainerSSHConnection) ServeHTTP(w http.ResponseWriter, req *http
 	}
 	defer conn.Close()
 
+	var bytesIn, bytesOut int64
+	var wg sync.WaitGroup
 	ctx, cancel := context.WithCancel(context.Background())
+	wg.Add(1)
 	go func() {
+		defer wg.Done()
 		defer cancel()
-		_, err := io.CopyN(conn, sshconn.Bufrw, int64(sshconn.Bufrw.Reader.Buffered()))
+		n, err := io.CopyN(conn, sshconn.Bufrw, int64(sshconn.Bufrw.Reader.Buffered()))
+		bytesOut += n
 		if err == nil {
-			_, err = io.Copy(conn, sshconn.Conn)
+			n, err = io.Copy(conn, sshconn.Conn)
+			bytesOut += n
 		}
 		if err != nil {
 			ctxlog.FromContext(req.Context()).WithError(err).Error("error copying downstream")
 		}
 	}()
+	wg.Add(1)
 	go func() {
+		defer wg.Done()
 		defer cancel()
-		_, err := io.CopyN(sshconn.Conn, bufrw, int64(bufrw.Reader.Buffered()))
+		n, err := io.CopyN(sshconn.Conn, bufrw, int64(bufrw.Reader.Buffered()))
+		bytesIn += n
 		if err == nil {
-			_, err = io.Copy(sshconn.Conn, conn)
+			n, err = io.Copy(sshconn.Conn, conn)
+			bytesIn += n
 		}
 		if err != nil {
 			ctxlog.FromContext(req.Context()).WithError(err).Error("error copying upstream")
 		}
 	}()
 	<-ctx.Done()
+	if sshconn.Logger != nil {
+		go func() {
+			wg.Wait()
+			sshconn.Logger.WithFields(logrus.Fields{
+				"bytesIn":  bytesIn,
+				"bytesOut": bytesOut,
+			}).Info("closed connection")
+		}()
+	}
 }