Make HTTP middleware easier to use

appengine/tsmon/middleware.go

Index: appengine/tsmon/middleware.go
diff --git a/appengine/tsmon/middleware.go b/appengine/tsmon/middleware.go
index f8f3ae0bed4fcc6254aaaa2b7b5504d91f89142e..a3f5dcb506a0cfc6cd09b78c9815b7c85c146be5 100644
--- a/appengine/tsmon/middleware.go
+++ b/appengine/tsmon/middleware.go
@@ -5,35 +5,34 @@
 package tsmon
 
 import (
 	"fmt"
 	"net/http"
 	"strings"
 	"sync"
 	"time"
 
 	"github.com/golang/protobuf/proto"
-	"github.com/julienschmidt/httprouter"
 	"golang.org/x/net/context"
 
 	"github.com/luci/gae/service/info"
 	gaeauth "github.com/luci/luci-go/appengine/gaeauth/client"
 	"github.com/luci/luci-go/common/clock"
 	gcps "github.com/luci/luci-go/common/gcloud/pubsub"
 	"github.com/luci/luci-go/common/iotools"
 	"github.com/luci/luci-go/common/logging"
 	"github.com/luci/luci-go/common/tsmon"
 	"github.com/luci/luci-go/common/tsmon/metric"
 	"github.com/luci/luci-go/common/tsmon/monitor"
 	"github.com/luci/luci-go/common/tsmon/store"
 	"github.com/luci/luci-go/common/tsmon/target"
-	"github.com/luci/luci-go/server/middleware"
+	"github.com/luci/luci-go/server/router"
 )
 
 // State holds the configuration of the tsmon library for GAE.
 //
 // Define it as a global variable and inject it in the request contexts using
 // State.Middleware().
 //
 // It will initialize itself from the tsmon state in the passed context on
 // a first use, mutating it along the way. Assumes caller is consistently using
 // contexts configured with exact same tsmon state (in a vast majority of cases
@@ -77,43 +76,44 @@ func (rw responseWriter) WriteHeader(code int) {
 }
 
 func newResponseWriter(rw http.ResponseWriter) responseWriter {
 	return responseWriter{
 		ResponseWriter: rw,
 		writer:         iotools.CountingWriter{Writer: rw},
 		status:         http.StatusOK,
 	}
 }
 
-// Middleware returns a middleware that must be inserted into the chain to
-// enable tsmon metrics to be sent on App Engine.
-func (s *State) Middleware(h middleware.Handler) middleware.Handler {
-	return func(c context.Context, rw http.ResponseWriter, r *http.Request, p httprouter.Params) {
-		state, settings := s.checkSettings(c)
-		if settings.Enabled {
-			started := clock.Now(c)
-			userAgent, ok := r.Header["User-Agent"]
-			if !ok || len(userAgent) == 0 {
-				userAgent = []string{"Unknown"}
-			}
-			nrw := newResponseWriter(rw)
-			defer func() {
-				dur := clock.Now(c).Sub(started)
-				metric.UpdatePresenceMetrics(c)
-				metric.UpdateServerMetrics(c, "/", nrw.Status(), dur,
-					r.ContentLength, nrw.Size(), userAgent[0])
-			}()
-			h(c, nrw, r, p)
-			s.flushIfNeeded(c, state, settings)
-		} else {
-			h(c, rw, r, p)
+// Middleware is a middleware that must be inserted into the middleware
+// chain to enable tsmon metrics to be send on AppEngine.
+func (s *State) Middleware(c *router.Context, next router.Handler) {
+	state, settings := s.checkSettings(c.Context)
+	if settings.Enabled {
+		started := clock.Now(c.Context)
+		userAgent, ok := c.Request.Header["User-Agent"]
+		if !ok || len(userAgent) == 0 {
+			userAgent = []string{"Unknown"}
 		}
+		ctx := c.Context
+		contentLength := c.Request.ContentLength
+		nrw := newResponseWriter(c.Writer)
+		c.Writer = nrw
+		defer func() {
+			dur := clock.Now(ctx).Sub(started)
+			metric.UpdatePresenceMetrics(ctx)
+			metric.UpdateServerMetrics(ctx, "/", nrw.Status(), dur,
+				contentLength, nrw.Size(), userAgent[0])
+		}()
+		next(c)
+		s.flushIfNeeded(ctx, state, settings)
+	} else {
+		next(c)
 	}
 }
 
 // checkSettings fetches tsmon settings and initializes, reinitializes or
 // deinitializes tsmon, as needed.
 //
 // Returns current tsmon state and settings. Panics if the context is using
 // unexpected tsmon state.
 func (s *State) checkSettings(c context.Context) (*tsmon.State, *tsmonSettings) {
 	state := tsmon.GetState(c)