sdk/go/arvados/container_gateway.go

   1 // Copyright (C) The Arvados Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: Apache-2.0
   4
   5 package arvados
   6
   7 import (
   8         "context"
   9         "io"
  10         "net/http"
  11         "sync"
  12
  13         "git.arvados.org/arvados.git/sdk/go/ctxlog"
  14         "git.arvados.org/arvados.git/sdk/go/httpserver"
  15         "github.com/sirupsen/logrus"
  16 )
  17
  18 func (cresp ConnectionResponse) ServeHTTP(w http.ResponseWriter, req *http.Request) {
  19         defer cresp.Conn.Close()
  20         conn, bufrw, err := http.NewResponseController(w).Hijack()
  21         if err != nil {
  22                 http.Error(w, "connection upgrade failed: "+err.Error(), http.StatusInternalServerError)
  23                 return
  24         }
  25         defer conn.Close()
  26         conn.Write([]byte("HTTP/1.1 101 Switching Protocols\r\n"))
  27         w.Header().Set("Connection", "upgrade")
  28         for k, v := range cresp.Header {
  29                 w.Header()[k] = v
  30         }
  31         w.Header().Write(conn)
  32         conn.Write([]byte("\r\n"))
  33         httpserver.ExemptFromDeadline(req)
  34
  35         var bytesIn, bytesOut int64
  36         ctx, cancel := context.WithCancel(req.Context())
  37         var wg sync.WaitGroup
  38         wg.Add(1)
  39         go func() {
  40                 defer wg.Done()
  41                 defer cancel()
  42                 n, err := io.CopyN(conn, cresp.Bufrw, int64(cresp.Bufrw.Reader.Buffered()))
  43                 bytesOut += n
  44                 if err == nil {
  45                         n, err = io.Copy(conn, cresp.Conn)
  46                         bytesOut += n
  47                 }
  48                 if err != nil {
  49                         ctxlog.FromContext(ctx).WithError(err).Error("error copying downstream")
  50                 }
  51         }()
  52         wg.Add(1)
  53         go func() {
  54                 defer wg.Done()
  55                 defer cancel()
  56                 n, err := io.CopyN(cresp.Conn, bufrw, int64(bufrw.Reader.Buffered()))
  57                 bytesIn += n
  58                 if err == nil {
  59                         n, err = io.Copy(cresp.Conn, conn)
  60                         bytesIn += n
  61                 }
  62                 if err != nil {
  63                         ctxlog.FromContext(ctx).WithError(err).Error("error copying upstream")
  64                 }
  65         }()
  66         <-ctx.Done()
  67         go func() {
  68                 // Wait for both io.Copy goroutines to finish and increment
  69                 // their byte counters.
  70                 wg.Wait()
  71                 if cresp.Logger != nil {
  72                         cresp.Logger.WithFields(logrus.Fields{
  73                                 "bytesIn":  bytesIn,
  74                                 "bytesOut": bytesOut,
  75                         }).Info("closed connection")
  76                 }
  77         }()
  78 }