feat: log long-lived connections acceptance (#17219)

Closes #16904

coderd/httpmw/logger.go

@@ -35,42 +35,93 @@ func Logger(log slog.Logger) func(next http.Handler) http.Handler {
 				slog.F("start", start),
 			)
 
-			next.ServeHTTP(sw, r)
+			logContext := NewRequestLogger(httplog, r.Method, start)
 
-			end := time.Now()
+			ctx := WithRequestLogger(r.Context(), logContext)
+
+			next.ServeHTTP(sw, r.WithContext(ctx))
 
 			// Don't log successful health check requests.
 			if r.URL.Path == "/api/v2" && sw.Status == http.StatusOK {
 				return
 			}
 
-			httplog = httplog.With(
-				slog.F("took", end.Sub(start)),
-				slog.F("status_code", sw.Status),
-				slog.F("latency_ms", float64(end.Sub(start)/time.Millisecond)),
-			)
-
-			// For status codes 400 and higher we
+			// For status codes 500 and higher we
 			// want to log the response body.
 			if sw.Status >= http.StatusInternalServerError {
-				httplog = httplog.With(
+				logContext.WithFields(
 					slog.F("response_body", string(sw.ResponseBody())),
 				)
 			}
 
-			// We should not log at level ERROR for 5xx status codes because 5xx
-			// includes proxy errors etc. It also causes slogtest to fail
-			// instantly without an error message by default.
-			logLevelFn := httplog.Debug
-			if sw.Status >= http.StatusInternalServerError {
-				logLevelFn = httplog.Warn
-			}
-
-			// We already capture most of this information in the span (minus
-			// the response body which we don't want to capture anyways).
-			tracing.RunWithoutSpan(r.Context(), func(ctx context.Context) {
-				logLevelFn(ctx, r.Method)
-			})
+			logContext.WriteLog(r.Context(), sw.Status)
 		})
 	}
 }
+
+type RequestLogger interface {
+	WithFields(fields ...slog.Field)
+	WriteLog(ctx context.Context, status int)
+}
+
+type SlogRequestLogger struct {
+	log     slog.Logger
+	written bool
+	message string
+	start   time.Time
+}
+
+var _ RequestLogger = &SlogRequestLogger{}
+
+func NewRequestLogger(log slog.Logger, message string, start time.Time) RequestLogger {
+	return &SlogRequestLogger{
+		log:     log,
+		written: false,
+		message: message,
+		start:   start,
+	}
+}
+
+func (c *SlogRequestLogger) WithFields(fields ...slog.Field) {
+	c.log = c.log.With(fields...)
+}
+
+func (c *SlogRequestLogger) WriteLog(ctx context.Context, status int) {
+	if c.written {
+		return
+	}
+	c.written = true
+	end := time.Now()
+
+	logger := c.log.With(
+		slog.F("took", end.Sub(c.start)),
+		slog.F("status_code", status),
+		slog.F("latency_ms", float64(end.Sub(c.start)/time.Millisecond)),
+	)
+	// We already capture most of this information in the span (minus
+	// the response body which we don't want to capture anyways).
+	tracing.RunWithoutSpan(ctx, func(ctx context.Context) {
+		// We should not log at level ERROR for 5xx status codes because 5xx
+		// includes proxy errors etc. It also causes slogtest to fail
+		// instantly without an error message by default.
+		if status >= http.StatusInternalServerError {
+			logger.Warn(ctx, c.message)
+		} else {
+			logger.Debug(ctx, c.message)
+		}
+	})
+}
+
+type logContextKey struct{}
+
+func WithRequestLogger(ctx context.Context, rl RequestLogger) context.Context {
+	return context.WithValue(ctx, logContextKey{}, rl)
+}
+
+func RequestLoggerFromContext(ctx context.Context) RequestLogger {
+	val := ctx.Value(logContextKey{})
+	if logCtx, ok := val.(RequestLogger); ok {
+		return logCtx
+	}
+	return nil
+}