X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/fdd3dd304439a7912db1c6ef257cf512ec7b3b52..2e03d03bc55b5a612c2bf04d878a72f2ee420d99:/sdk/go/httpserver/logger.go

diff --git a/sdk/go/httpserver/logger.go b/sdk/go/httpserver/logger.go
index 7eb7f0f03d..b71adf7118 100644
--- 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 {
@@ -44,7 +47,13 @@ func (hn hijackNotifier) Hijack() (net.Conn, *bufio.ReadWriter, error) {
 // HandlerWithDeadline cancels the request context if the request
 // 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.WithCancel(r.Context())
 		defer cancel()
@@ -64,6 +73,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 +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)
 
@@ -124,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