apiserver: store httplog in context

Kubernetes-commit: e848eaa566b750ad5e5fd26a64eecfc883f936f3

pkg/endpoints/handlers/watch.go

@@ -165,7 +165,7 @@ func (s *WatchServer) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 	metrics.RegisteredWatchers.WithLabelValues(kind.Group, kind.Version, kind.Kind).Inc()
 	defer metrics.RegisteredWatchers.WithLabelValues(kind.Group, kind.Version, kind.Kind).Dec()
 
-	w = httplog.Unlogged(w)
+	w = httplog.Unlogged(req, w)
 
 	if wsstream.IsWebSocketRequest(req) {
 		w.Header().Set("Content-Type", s.MediaType)

pkg/server/filters/wrap.go

@@ -34,14 +34,12 @@ func WithPanicRecovery(handler http.Handler) http.Handler {
 }
 
 func withPanicRecovery(handler http.Handler, crashHandler func(http.ResponseWriter, *http.Request, interface{})) http.Handler {
+	handler = httplog.WithLogging(handler, httplog.DefaultStacktracePred)
 	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
 		defer runtime.HandleCrash(func(err interface{}) {
 			crashHandler(w, req, err)
 		})
 
-		logger := httplog.NewLogged(req, &w)
-		defer logger.Log()
-
 		// Dispatch to the internal handler
 		handler.ServeHTTP(w, req)
 	})

pkg/server/healthz/healthz.go

@@ -181,7 +181,7 @@ func handleRootHealthz(checks ...HealthzChecker) http.HandlerFunc {
 		// always be verbose on failure
 		if failed {
 			klog.V(2).Infof("%vhealthz check failed", verboseOut.String())
-			http.Error(httplog.Unlogged(w), fmt.Sprintf("%vhealthz check failed", verboseOut.String()), http.StatusInternalServerError)
+			http.Error(httplog.Unlogged(r, w), fmt.Sprintf("%vhealthz check failed", verboseOut.String()), http.StatusInternalServerError)
 			return
 		}
 

pkg/server/httplog/httplog.go

@@ -18,6 +18,7 @@ package httplog
 
 import (
 	"bufio"
+	"context"
 	"fmt"
 	"net"
 	"net/http"
@@ -34,6 +35,11 @@ type logger interface {
 	Addf(format string, data ...interface{})
 }
 
+type respLoggerContextKeyType int
+
+// respLoggerContextKey is used to store the respLogger pointer in the request context.
+const respLoggerContextKey respLoggerContextKeyType = iota
+
 // Add a layer on top of ResponseWriter, so we can track latency and error
 // message sources.
 //
@@ -69,47 +75,54 @@ func DefaultStacktracePred(status int) bool {
 	return (status < http.StatusOK || status >= http.StatusInternalServerError) && status != http.StatusSwitchingProtocols
 }
 
-// NewLogged turns a normal response writer into a logged response writer.
-//
-// Usage:
-//
-// defer NewLogged(req, &w).StacktraceWhen(StatusIsNot(200, 202)).Log()
-//
-// (Only the call to Log() is deferred, so you can set everything up in one line!)
-//
-// Note that this *changes* your writer, to route response writing actions
-// through the logger.
-//
-// Use LogOf(w).Addf(...) to log something along with the response result.
-func NewLogged(req *http.Request, w *http.ResponseWriter) *respLogger {
-	if _, ok := (*w).(*respLogger); ok {
-		// Don't double-wrap!
-		panic("multiple NewLogged calls!")
+// WithLogging wraps the handler with logging.
+func WithLogging(handler http.Handler, pred StacktracePred) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+		ctx := req.Context()
+		if old := respLoggerFromContext(req); old != nil {
+			panic("multiple WithLogging calls!")
+		}
+		rl := newLogged(req, w).StacktraceWhen(pred)
+		req = req.WithContext(context.WithValue(ctx, respLoggerContextKey, rl))
+
+		defer rl.Log()
+		handler.ServeHTTP(rl, req)
+	})
+}
+
+// respLoggerFromContext returns the respLogger or nil.
+func respLoggerFromContext(req *http.Request) *respLogger {
+	ctx := req.Context()
+	val := ctx.Value(respLoggerContextKey)
+	if rl, ok := val.(*respLogger); ok {
+		return rl
 	}
-	rl := &respLogger{
+	return nil
+}
+
+// newLogged turns a normal response writer into a logged response writer.
+func newLogged(req *http.Request, w http.ResponseWriter) *respLogger {
+	return &respLogger{
 		startTime:         time.Now(),
 		req:               req,
-		w:                 *w,
+		w:                 w,
 		logStacktracePred: DefaultStacktracePred,
 	}
-	*w = rl // hijack caller's writer!
-	return rl
 }
 
 // LogOf returns the logger hiding in w. If there is not an existing logger
 // then a passthroughLogger will be created which will log to stdout immediately
 // when Addf is called.
 func LogOf(req *http.Request, w http.ResponseWriter) logger {
-	if rl, ok := w.(*respLogger); ok {
+	if rl := respLoggerFromContext(req); rl != nil {
 		return rl
 	}
-
 	return &passthroughLogger{}
 }
 
 // Unlogged returns the original ResponseWriter, or w if it is not our inserted logger.
-func Unlogged(w http.ResponseWriter) http.ResponseWriter {
-	if rl, ok := w.(*respLogger); ok {
+func Unlogged(req *http.Request, w http.ResponseWriter) http.ResponseWriter {
+	if rl := respLoggerFromContext(req); rl != nil {
 		return rl.w
 	}
 	return w

pkg/server/httplog/httplog_test.go

@@ -56,22 +56,24 @@ func TestStatusIsNot(t *testing.T) {
 	}
 }
 
-func TestNewLogged(t *testing.T) {
+func TestWithLogging(t *testing.T) {
 	req, err := http.NewRequest("GET", "http://example.com", nil)
 	if err != nil {
 		t.Errorf("Unexpected error: %v", err)
 	}
-	handler := func(w http.ResponseWriter, r *http.Request) {
-		NewLogged(req, &w)
+	var handler http.Handler
+	handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
+	handler = WithLogging(WithLogging(handler, DefaultStacktracePred), DefaultStacktracePred)
+
+	func() {
 		defer func() {
 			if r := recover(); r == nil {
-				t.Errorf("Expected NewLogged to panic")
+				t.Errorf("Expected newLogged to panic")
 			}
 		}()
-		NewLogged(req, &w)
-	}
-	w := httptest.NewRecorder()
-	handler(w, req)
+		w := httptest.NewRecorder()
+		handler.ServeHTTP(w, req)
+	}()
 }
 
 func TestLogOf(t *testing.T) {
@@ -81,21 +83,23 @@ func TestLogOf(t *testing.T) {
 		if err != nil {
 			t.Errorf("Unexpected error: %v", err)
 		}
-		handler := func(w http.ResponseWriter, r *http.Request) {
-			var want string
-			if makeLogger {
-				NewLogged(req, &w)
-				want = "*httplog.respLogger"
-			} else {
-				want = "*httplog.passthroughLogger"
-			}
+		var want string
+		var handler http.Handler
+		handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			got := reflect.TypeOf(LogOf(r, w)).String()
 			if want != got {
 				t.Errorf("Expected %v, got %v", want, got)
 			}
+		})
+		if makeLogger {
+			handler = WithLogging(handler, DefaultStacktracePred)
+			want = "*httplog.respLogger"
+		} else {
+			want = "*httplog.passthroughLogger"
 		}
+
 		w := httptest.NewRecorder()
-		handler(w, req)
+		handler.ServeHTTP(w, req)
 	}
 }
 
@@ -106,18 +110,20 @@ func TestUnlogged(t *testing.T) {
 		if err != nil {
 			t.Errorf("Unexpected error: %v", err)
 		}
-		handler := func(w http.ResponseWriter, r *http.Request) {
-			want := w
-			if makeLogger {
-				NewLogged(req, &w)
-			}
-			got := Unlogged(w)
-			if want != got {
-				t.Errorf("Expected %v, got %v", want, got)
+
+		origWriter := httptest.NewRecorder()
+		var handler http.Handler
+		handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			got := Unlogged(r, w)
+			if origWriter != got {
+				t.Errorf("Expected origin writer, got %#v", got)
 			}
+		})
+		if makeLogger {
+			handler = WithLogging(handler, DefaultStacktracePred)
 		}
-		w := httptest.NewRecorder()
-		handler(w, req)
+
+		handler.ServeHTTP(origWriter, req)
 	}
 }
 
@@ -134,7 +140,7 @@ func TestLoggedStatus(t *testing.T) {
 	}
 
 	var tw http.ResponseWriter = new(testResponseWriter)
-	logger := NewLogged(req, &tw)
+	logger := newLogged(req, tw)
 	logger.Write(nil)
 
 	if logger.status != http.StatusOK {
@@ -142,7 +148,7 @@ func TestLoggedStatus(t *testing.T) {
 	}
 
 	tw = new(testResponseWriter)
-	logger = NewLogged(req, &tw)
+	logger = newLogged(req, tw)
 	logger.WriteHeader(http.StatusForbidden)
 	logger.Write(nil)