Merge pull request weaveworks#61 from prometheus/context-injection

route: Allow per-request custom context injection

route/route.go

@@ -32,67 +32,79 @@ func WithParam(ctx context.Context, p, v string) context.Context {
 	return context.WithValue(ctx, param(p), v)
 }
 
-// handle turns a Handle into httprouter.Handle
-func handle(h http.HandlerFunc) httprouter.Handle {
-	return func(w http.ResponseWriter, r *http.Request, params httprouter.Params) {
-		ctx, cancel := context.WithCancel(context.Background())
+type contextFn func(r *http.Request) context.Context
+
+// Router wraps httprouter.Router and adds support for prefixed sub-routers
+// and per-request context injections.
+type Router struct {
+	rtr    *httprouter.Router
+	prefix string
+	ctxFn  contextFn
+}
+
+// New returns a new Router.
+func New(ctxFn contextFn) *Router {
+	if ctxFn == nil {
+		ctxFn = func(r *http.Request) context.Context {
+			return context.Background()
+		}
+	}
+	return &Router{
+		rtr:   httprouter.New(),
+		ctxFn: ctxFn,
+	}
+}
+
+// WithPrefix returns a router that prefixes all registered routes with prefix.
+func (r *Router) WithPrefix(prefix string) *Router {
+	return &Router{rtr: r.rtr, prefix: r.prefix + prefix, ctxFn: r.ctxFn}
+}
+
+// handle turns a HandlerFunc into an httprouter.Handle.
+func (r *Router) handle(h http.HandlerFunc) httprouter.Handle {
+	return func(w http.ResponseWriter, req *http.Request, params httprouter.Params) {
+		ctx, cancel := context.WithCancel(r.ctxFn(req))
 		defer cancel()
 
 		for _, p := range params {
 			ctx = context.WithValue(ctx, param(p.Key), p.Value)
 		}
 
 		mtx.Lock()
-		ctxts[r] = ctx
+		ctxts[req] = ctx
 		mtx.Unlock()
 
-		h(w, r)
+		h(w, req)
 
 		mtx.Lock()
-		delete(ctxts, r)
+		delete(ctxts, req)
 		mtx.Unlock()
 	}
 }
 
-// Router wraps httprouter.Router and adds support for prefixed sub-routers.
-type Router struct {
-	rtr    *httprouter.Router
-	prefix string
-}
-
-// New returns a new Router.
-func New() *Router {
-	return &Router{rtr: httprouter.New()}
-}
-
-// WithPrefix returns a router that prefixes all registered routes with prefix.
-func (r *Router) WithPrefix(prefix string) *Router {
-	return &Router{rtr: r.rtr, prefix: r.prefix + prefix}
-}
-
 // Get registers a new GET route.
 func (r *Router) Get(path string, h http.HandlerFunc) {
-	r.rtr.GET(r.prefix+path, handle(h))
+	r.rtr.GET(r.prefix+path, r.handle(h))
 }
 
 // Options registers a new OPTIONS route.
 func (r *Router) Options(path string, h http.HandlerFunc) {
-	r.rtr.OPTIONS(r.prefix+path, handle(h))
+	r.rtr.OPTIONS(r.prefix+path, r.handle(h))
 }
 
 // Del registers a new DELETE route.
 func (r *Router) Del(path string, h http.HandlerFunc) {
-	r.rtr.DELETE(r.prefix+path, handle(h))
+	r.rtr.DELETE(r.prefix+path, r.handle(h))
 }
 
 // Put registers a new PUT route.
 func (r *Router) Put(path string, h http.HandlerFunc) {
-	r.rtr.PUT(r.prefix+path, handle(h))
+	r.rtr.PUT(r.prefix+path, r.handle(h))
 }
 
 // Post registers a new POST route.
 func (r *Router) Post(path string, h http.HandlerFunc) {
-	r.rtr.POST(r.prefix+path, handle(h))
+	r.rtr.POST(r.prefix+path, r.handle(h))
 }
 
 // Redirect takes an absolute path and sends an internal HTTP redirect for it,

route/route_test.go

@@ -4,10 +4,12 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"testing"
+
+	"golang.org/x/net/context"
 )
 
 func TestRedirect(t *testing.T) {
-	router := New().WithPrefix("/test/prefix")
+	router := New(nil).WithPrefix("/test/prefix")
 	w := httptest.NewRecorder()
 	r, err := http.NewRequest("GET", "http://localhost:9090/foo", nil)
 	if err != nil {
@@ -25,3 +27,23 @@ func TestRedirect(t *testing.T) {
 		t.Fatalf("Unexpected redirect location: got %s, want %s", got, want)
 	}
 }
+
+func TestContextFn(t *testing.T) {
+	router := New(func(r *http.Request) context.Context {
+		return context.WithValue(context.Background(), "testkey", "testvalue")
+	})
+
+	router.Get("/test", func(w http.ResponseWriter, r *http.Request) {
+		want := "testvalue"
+		got := Context(r).Value("testkey")
+		if want != got {
+			t.Fatalf("Unexpected context value: want %q, got %q", want, got)
+		}
+	})
+
+	r, err := http.NewRequest("GET", "http://localhost:9090/test", nil)
+	if err != nil {
+		t.Fatalf("Error building test request: %s", err)
+	}
+	router.ServeHTTP(nil, r)
+}