18896: review feedback, un-pyramid the logrus.Fields update.

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

diff --git a/sdk/go/httpserver/logger.go b/sdk/go/httpserver/logger.go
index 7eb7f0f03d57b571e314f8d87ca6714cf7d6563f..5a46635e9102365bbfd01c9c9c120bd8e23a7026 100644 (file)
--- a/sdk/go/httpserver/logger.go
+++ b/sdk/go/httpserver/logger.go
@@ -9,6 +9,7 @@ import (
        "context"
        "net"
        "net/http"
+       "sync"
        "time"
 
        "git.arvados.org/arvados.git/sdk/go/ctxlog"
@@ -21,7 +22,9 @@ type contextKey struct {
 }
 
 var (
-       requestTimeContextKey = contextKey{"requestTime"}
+       requestTimeContextKey       = contextKey{"requestTime"}
+       responseLogFieldsContextKey = contextKey{"responseLogFields"}
+       mutexContextKey             = contextKey{"mutex"}
 )
 
 type hijacker interface {
@@ -64,6 +67,19 @@ func HandlerWithDeadline(timeout time.Duration, next http.Handler) http.Handler
        })
 }
 
+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 {
@@ -81,6 +97,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)
 
@@ -124,6 +142,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