From 1f0a9dbca0fa5dd106adc3d190575d9dc6d8ac8d Mon Sep 17 00:00:00 2001
From: Stephen J Day <stephen.day@docker.com>
Date: Fri, 7 Apr 2017 15:19:53 -0700
Subject: [PATCH] register/handlers: remove context manager

Signed-off-by: Stephen J Day <stephen.day@docker.com>
---
 registry/handlers/app.go     | 24 +++++++--------
 registry/handlers/context.go | 60 ------------------------------------
 2 files changed, 11 insertions(+), 73 deletions(-)

diff --git a/registry/handlers/app.go b/registry/handlers/app.go
index 04e346cc..5393e88d 100644
--- a/registry/handlers/app.go
+++ b/registry/handlers/app.go
@@ -596,24 +596,19 @@ func (app *App) configureSecret(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)
+	// Prepare the context with our own little decorations.
+	ctx := r.Context()
+	ctx = ctxu.WithRequest(ctx, r)
+	ctx, w = ctxu.WithResponseWriter(ctx, w)
+	ctx = ctxu.WithLogger(ctx, ctxu.GetRequestLogger(ctx))
+	r = r.WithContext(ctx)
 
 	defer func() {
 		status, ok := ctx.Value("http.response.status").(int)
 		if ok && status >= 200 && status <= 399 {
-			ctxu.GetResponseLogger(ctx).Infof("response completed")
+			ctxu.GetResponseLogger(r.Context()).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")
@@ -649,6 +644,9 @@ func (app *App) dispatcher(dispatch dispatchFunc) http.Handler {
 		// Add username to request logging
 		context.Context = ctxu.WithLogger(context.Context, ctxu.GetLogger(context.Context, auth.UserNameKey))
 
+		// sync up context on the request.
+		r = r.WithContext(context)
+
 		if app.nameRequired(r) {
 			nameRef, err := reference.WithName(getName(context))
 			if err != nil {
@@ -756,7 +754,7 @@ func (app *App) logError(context context.Context, errors errcode.Errors) {
 // 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 := defaultContextManager.context(app, w, r)
+	ctx := r.Context()
 	ctx = ctxu.WithVars(ctx, r)
 	ctx = ctxu.WithLogger(ctx, ctxu.GetLogger(ctx,
 		"vars.name",
diff --git a/registry/handlers/context.go b/registry/handlers/context.go
index a49ad10e..6c1be5b3 100644
--- a/registry/handlers/context.go
+++ b/registry/handlers/context.go
@@ -3,7 +3,6 @@ package handlers
 import (
 	"fmt"
 	"net/http"
-	"sync"
 
 	"github.com/docker/distribution"
 	ctxu "github.com/docker/distribution/context"
@@ -91,62 +90,3 @@ 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)
-}