Instantiate http context before dispatch

Ensure that the status is logged in the context by instantiating before the
request is routed to handlers. While this requires some level of hacking to
acheive, the result is that the context value of "http.request.status" is as
accurate as possible for each request.

Signed-off-by: Stephen J Day <stephen.day@docker.com>

registry/handlers/app.go

@@ -273,6 +273,21 @@ func (app *App) configureRedis(configuration *configuration.Configuration) {
 func (app *App) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	defer r.Body.Close() // ensure that request body is always closed.
 
+	// Instantiate an http context here so we can track the error codes
+	// returned by the request router.
+	ctx := defaultContextManager.context(app, w, r)
+	defer func() {
+		ctxu.GetResponseLogger(ctx).Infof("response completed")
+	}()
+	defer defaultContextManager.release(ctx)
+
+	// NOTE(stevvooe): Total hack to get instrumented responsewriter from context.
+	var err error
+	w, err = ctxu.GetResponseWriter(ctx)
+	if err != nil {
+		ctxu.GetLogger(ctx).Warnf("response writer not found in context")
+	}
+
 	// Set a header with the Docker Distribution API Version for all responses.
 	w.Header().Add("Docker-Distribution-API-Version", "registry/2.0")
 	app.router.ServeHTTP(w, r)
@@ -287,38 +302,12 @@ type dispatchFunc func(ctx *Context, r *http.Request) http.Handler
 // TODO(stevvooe): dispatchers should probably have some validation error
 // chain with proper error reporting.
 
-// singleStatusResponseWriter only allows the first status to be written to be
-// the valid request status. The current use case of this class should be
-// factored out.
-type singleStatusResponseWriter struct {
-	http.ResponseWriter
-	status int
-}
-
-func (ssrw *singleStatusResponseWriter) WriteHeader(status int) {
-	if ssrw.status != 0 {
-		return
-	}
-	ssrw.status = status
-	ssrw.ResponseWriter.WriteHeader(status)
-}
-
-func (ssrw *singleStatusResponseWriter) Flush() {
-	if flusher, ok := ssrw.ResponseWriter.(http.Flusher); ok {
-		flusher.Flush()
-	}
-}
-
 // dispatcher returns a handler that constructs a request specific context and
 // handler, using the dispatch factory function.
 func (app *App) dispatcher(dispatch dispatchFunc) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		context := app.context(w, r)
 
-		defer func() {
-			ctxu.GetResponseLogger(context).Infof("response completed")
-		}()
-
 		if err := app.authorized(w, r, context); err != nil {
 			ctxu.GetLogger(context).Errorf("error authorizing context: %v", err)
 			return
@@ -360,16 +349,16 @@ func (app *App) dispatcher(dispatch dispatchFunc) http.Handler {
 			}
 		}
 
-		handler := dispatch(context, r)
+		dispatch(context, r).ServeHTTP(w, r)
-
-		ssrw := &singleStatusResponseWriter{ResponseWriter: w}
-		handler.ServeHTTP(ssrw, r)
 
 		// Automated error response handling here. Handlers may return their
 		// own errors if they need different behavior (such as range errors
 		// for layer upload).
 		if context.Errors.Len() > 0 {
-			if ssrw.status == 0 {
+			if context.Value("http.response.status") == 0 {
+				// TODO(stevvooe): Getting this value from the context is a
+				// bit of a hack. We can further address with some of our
+				// future refactoring.
 				w.WriteHeader(http.StatusBadRequest)
 			}
 			serveJSON(w, context.Errors)
@@ -380,10 +369,8 @@ func (app *App) dispatcher(dispatch dispatchFunc) http.Handler {
 // context constructs the context object for the application. This only be
 // called once per request.
 func (app *App) context(w http.ResponseWriter, r *http.Request) *Context {
-	ctx := ctxu.WithRequest(app, r)
+	ctx := defaultContextManager.context(app, w, r)
-	ctx, w = ctxu.WithResponseWriter(ctx, w)
 	ctx = ctxu.WithVars(ctx, r)
-	ctx = ctxu.WithLogger(ctx, ctxu.GetRequestLogger(ctx))
 	ctx = ctxu.WithLogger(ctx, ctxu.GetLogger(ctx,
 		"vars.name",
 		"vars.reference",

registry/handlers/context.go

@@ -3,6 +3,7 @@ package handlers
 import (
 	"fmt"
 	"net/http"
+	"sync"
 
 	"github.com/docker/distribution"
 	ctxu "github.com/docker/distribution/context"
@@ -88,3 +89,62 @@ func getUserName(ctx context.Context, r *http.Request) string {
 
 	return username
 }
+
+// contextManager allows us to associate net/context.Context instances with a
+// request, based on the memory identity of http.Request. This prepares http-
+// level context, which is not application specific. If this is called,
+// (*contextManager).release must be called on the context when the request is
+// completed.
+//
+// Providing this circumvents a lot of necessity for dispatchers with the
+// benefit of instantiating the request context much earlier.
+//
+// TODO(stevvooe): Consider making this facility a part of the context package.
+type contextManager struct {
+	contexts map[*http.Request]context.Context
+	mu       sync.Mutex
+}
+
+// defaultContextManager is just a global instance to register request contexts.
+var defaultContextManager = newContextManager()
+
+func newContextManager() *contextManager {
+	return &contextManager{
+		contexts: make(map[*http.Request]context.Context),
+	}
+}
+
+// context either returns a new context or looks it up in the manager.
+func (cm *contextManager) context(parent context.Context, w http.ResponseWriter, r *http.Request) context.Context {
+	cm.mu.Lock()
+	defer cm.mu.Unlock()
+
+	ctx, ok := cm.contexts[r]
+	if ok {
+		return ctx
+	}
+
+	if parent == nil {
+		parent = ctxu.Background()
+	}
+
+	ctx = ctxu.WithRequest(parent, r)
+	ctx, w = ctxu.WithResponseWriter(ctx, w)
+	ctx = ctxu.WithLogger(ctx, ctxu.GetRequestLogger(ctx))
+	cm.contexts[r] = ctx
+
+	return ctx
+}
+
+// releases frees any associated with resources from request.
+func (cm *contextManager) release(ctx context.Context) {
+	cm.mu.Lock()
+	defer cm.mu.Unlock()
+
+	r, err := ctxu.GetRequest(ctx)
+	if err != nil {
+		ctxu.GetLogger(ctx).Errorf("no request found in context during release")
+		return
+	}
+	delete(cm.contexts, r)
+}