diff --git a/route/route.go b/route/route.go index bb4688173..dbec638e5 100644 --- a/route/route.go +++ b/route/route.go @@ -19,11 +19,12 @@ func WithParam(ctx context.Context, p, v string) context.Context { return context.WithValue(ctx, param(p), v) } -// Router wraps httprouter.Router and adds support for prefixed sub-routers -// and per-request context injections. +// Router wraps httprouter.Router and adds support for prefixed sub-routers, +// per-request context injections and instrumentation. type Router struct { rtr *httprouter.Router prefix string + instrh func(handlerName string, handler http.HandlerFunc) http.HandlerFunc } // New returns a new Router. @@ -33,13 +34,18 @@ func New() *Router { } } +// WithInstrumentation returns a router with instrumentation support. +func (r *Router) WithInstrumentation(instrh func(handlerName string, handler http.HandlerFunc) http.HandlerFunc) *Router { + return &Router{rtr: r.rtr, prefix: r.prefix, instrh: instrh} +} + // 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} + return &Router{rtr: r.rtr, prefix: r.prefix + prefix, instrh: r.instrh} } // handle turns a HandlerFunc into an httprouter.Handle. -func (r *Router) handle(h http.HandlerFunc) httprouter.Handle { +func (r *Router) handle(handlerName string, h http.HandlerFunc) httprouter.Handle { return func(w http.ResponseWriter, req *http.Request, params httprouter.Params) { ctx, cancel := context.WithCancel(req.Context()) defer cancel() @@ -47,33 +53,36 @@ func (r *Router) handle(h http.HandlerFunc) httprouter.Handle { for _, p := range params { ctx = context.WithValue(ctx, param(p.Key), p.Value) } + if r.instrh != nil { + h = r.instrh(handlerName, h) + } h(w, req.WithContext(ctx)) } } // Get registers a new GET route. func (r *Router) Get(path string, h http.HandlerFunc) { - r.rtr.GET(r.prefix+path, r.handle(h)) + r.rtr.GET(r.prefix+path, r.handle(path, h)) } // Options registers a new OPTIONS route. func (r *Router) Options(path string, h http.HandlerFunc) { - r.rtr.OPTIONS(r.prefix+path, r.handle(h)) + r.rtr.OPTIONS(r.prefix+path, r.handle(path, h)) } // Del registers a new DELETE route. func (r *Router) Del(path string, h http.HandlerFunc) { - r.rtr.DELETE(r.prefix+path, r.handle(h)) + r.rtr.DELETE(r.prefix+path, r.handle(path, h)) } // Put registers a new PUT route. func (r *Router) Put(path string, h http.HandlerFunc) { - r.rtr.PUT(r.prefix+path, r.handle(h)) + r.rtr.PUT(r.prefix+path, r.handle(path, h)) } // Post registers a new POST route. func (r *Router) Post(path string, h http.HandlerFunc) { - r.rtr.POST(r.prefix+path, r.handle(h)) + r.rtr.POST(r.prefix+path, r.handle(path, h)) } // Redirect takes an absolute path and sends an internal HTTP redirect for it, diff --git a/route/route_test.go b/route/route_test.go index a9bb20996..d491cad66 100644 --- a/route/route_test.go +++ b/route/route_test.go @@ -42,3 +42,35 @@ func TestContext(t *testing.T) { } router.ServeHTTP(nil, r) } + +func TestInstrumentation(t *testing.T) { + var got string + cases := []struct { + router *Router + want string + }{ + { + router: New(), + want: "", + }, { + router: New().WithInstrumentation(func(handlerName string, handler http.HandlerFunc) http.HandlerFunc { + got = handlerName + return handler + }), + want: "/foo", + }, + } + + for _, c := range cases { + c.router.Get("/foo", func(w http.ResponseWriter, r *http.Request) {}) + + r, err := http.NewRequest("GET", "http://localhost:9090/foo", nil) + if err != nil { + t.Fatalf("Error building test request: %s", err) + } + c.router.ServeHTTP(nil, r) + if c.want != got { + t.Fatalf("Unexpected value: want %q, got %q", c.want, got) + } + } +}