Merge pull request #369 from stevvooe/http-request-status-context-man…

…ager

context, registry/handlers: instantiate http context before dispatch

context/http.go

@@ -15,7 +15,8 @@ import (
 
 // Common errors used with this package.
 var (
-	ErrNoRequestContext = errors.New("no http request in context")
+	ErrNoRequestContext        = errors.New("no http request in context")
+	ErrNoResponseWriterContext = errors.New("no http response in context")
 )
 
 func parseIP(ipStr string) net.IP {
@@ -110,6 +111,20 @@ func WithResponseWriter(ctx Context, w http.ResponseWriter) (Context, http.Respo
 	return irw, irw
 }
 
+// GetResponseWriter returns the http.ResponseWriter from the provided
+// context. If not present, ErrNoResponseWriterContext is returned. The
+// returned instance provides instrumentation in the context.
+func GetResponseWriter(ctx Context) (http.ResponseWriter, error) {
+	v := ctx.Value("http.response")
+
+	rw, ok := v.(http.ResponseWriter)
+	if !ok || rw == nil {
+		return nil, ErrNoResponseWriterContext
+	}
+
+	return rw, nil
+}
+
 // getVarsFromRequest let's us change request vars implementation for testing
 // and maybe future changes.
 var getVarsFromRequest = mux.Vars

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)
-
-		ssrw := &singleStatusResponseWriter{ResponseWriter: w}
-		handler.ServeHTTP(ssrw, r)
+		dispatch(context, r).ServeHTTP(w, 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, w = ctxu.WithResponseWriter(ctx, w)
+	ctx := defaultContextManager.context(app, w, r)
 	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)
+}