Merge branch '19205-monitor-handler-threads'

[arvados.git] / sdk / go / httpserver / logger.go

diff --git a/sdk/go/httpserver/logger.go b/sdk/go/httpserver/logger.go
index a0ca6bf28d8d37daf2f8a0406d9ce15e7d4571e6..b71adf71181a9eb6b550093f73dbe4ab884038ce 100644 (file)
--- a/sdk/go/httpserver/logger.go
+++ b/sdk/go/httpserver/logger.go
@@ -5,8 +5,11 @@
 package httpserver
 
 import (
+       "bufio"
        "context"
+       "net"
        "net/http"
+       "sync"
        "time"
 
        "git.arvados.org/arvados.git/sdk/go/ctxlog"
@@ -19,19 +22,70 @@ type contextKey struct {
 }
 
 var (
-       requestTimeContextKey = contextKey{"requestTime"}
+       requestTimeContextKey       = contextKey{"requestTime"}
+       responseLogFieldsContextKey = contextKey{"responseLogFields"}
+       mutexContextKey             = contextKey{"mutex"}
 )
 
+type hijacker interface {
+       http.ResponseWriter
+       http.Hijacker
+}
+
+// hijackNotifier wraps a ResponseWriter, calling the provided
+// Notify() func if/when the wrapped Hijacker is hijacked.
+type hijackNotifier struct {
+       hijacker
+       hijacked chan<- bool
+}
+
+func (hn hijackNotifier) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+       close(hn.hijacked)
+       return hn.hijacker.Hijack()
+}
+
 // HandlerWithDeadline cancels the request context if the request
-// takes longer than the specified timeout.
+// takes longer than the specified timeout without having its
+// connection hijacked.
+//
+// If timeout is 0, there is no deadline: HandlerWithDeadline is a
+// no-op.
 func HandlerWithDeadline(timeout time.Duration, next http.Handler) http.Handler {
+       if timeout == 0 {
+               return next
+       }
        return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-               ctx, cancel := context.WithDeadline(r.Context(), time.Now().Add(timeout))
+               ctx, cancel := context.WithCancel(r.Context())
                defer cancel()
+               nodeadline := make(chan bool)
+               go func() {
+                       select {
+                       case <-nodeadline:
+                       case <-ctx.Done():
+                       case <-time.After(timeout):
+                               cancel()
+                       }
+               }()
+               if hj, ok := w.(hijacker); ok {
+                       w = hijackNotifier{hj, nodeadline}
+               }
                next.ServeHTTP(w, r.WithContext(ctx))
        })
 }
 
+func SetResponseLogFields(ctx context.Context, fields logrus.Fields) {
+       m, _ := ctx.Value(&mutexContextKey).(*sync.Mutex)
+       c, _ := ctx.Value(&responseLogFieldsContextKey).(logrus.Fields)
+       if m == nil || c == nil {
+               return
+       }
+       m.Lock()
+       defer m.Unlock()
+       for k, v := range fields {
+               c[k] = v
+       }
+}
+
 // LogRequests wraps an http.Handler, logging each request and
 // response.
 func LogRequests(h http.Handler) http.Handler {
@@ -49,6 +103,8 @@ func LogRequests(h http.Handler) http.Handler {
                })
                ctx := req.Context()
                ctx = context.WithValue(ctx, &requestTimeContextKey, time.Now())
+               ctx = context.WithValue(ctx, &responseLogFieldsContextKey, logrus.Fields{})
+               ctx = context.WithValue(ctx, &mutexContextKey, &sync.Mutex{})
                ctx = ctxlog.Context(ctx, lgr)
                req = req.WithContext(ctx)
 
@@ -92,6 +148,9 @@ func logResponse(w *responseTimer, req *http.Request, lgr *logrus.Entry) {
                        "timeWriteBody": stats.Duration(tDone.Sub(writeTime)),
                })
        }
+       if responseLogFields, ok := req.Context().Value(&responseLogFieldsContextKey).(logrus.Fields); ok {
+               lgr = lgr.WithFields(responseLogFields)
+       }
        respCode := w.WroteStatus()
        if respCode == 0 {
                respCode = http.StatusOK