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..a4397ae7c4655f83d0468b047d929b85e58f53b5 100644
--- a/appengine/tsmon/middleware.go
+++ b/appengine/tsmon/middleware.go
@@ -12,7 +12,6 @@ import (
 	"time"
 
 	"github.com/golang/protobuf/proto"
-	"github.com/julienschmidt/httprouter"
 	"golang.org/x/net/context"
 
 	"github.com/luci/gae/service/info"
@@ -26,7 +25,7 @@ import (
 	"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.
@@ -84,29 +83,28 @@ func newResponseWriter(rw http.ResponseWriter) responseWriter {
 	}
 }
 
-// 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 returns a middleware that must be inserted into the middleware

[Vadim Sh., 2016/06/18 16:57:20]

nit: "is a middleware"

[nishanths, 2016/06/19 02:47:02]

Done.

+// 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"}
 		}
+		nrw := newResponseWriter(c.Writer)
+		c.Writer = nrw
+		defer func() {
+			dur := clock.Now(c.Context).Sub(started)
+			metric.UpdatePresenceMetrics(c.Context)
+			metric.UpdateServerMetrics(c.Context, "/", nrw.Status(), dur,
+				c.Request.ContentLength, nrw.Size(), userAgent[0])
+		}()
+		next(c)
+		s.flushIfNeeded(c.Context, state, settings)

[Vadim Sh., 2016/06/18 16:57:20]

I think we need to preserve the original c.Context and use it here.

next(c) can mutate c.Context now, replacing it with something unusable.

Same for defer: keep a copy of c.Context and grab c.Request.ContentLength before
calling next().

[nishanths, 2016/06/19 02:47:02]

Done. With Patch Set 39 applied, it matches behavior in origin/master, except
for:

- In Patch Set 29 we are grabbing c.Request.ContentLength (int64). 
- In origin/master, we only grab r (pointer to http.Request). 

What should we do here? (I'm unsure because I don't understand the expected
behavior of the monitoring code completely yet.)

[nishanths, 2016/06/20 15:14:29]

typo: meant Patch Set 30 at both spots above

+	} else {
+		next(c)
 	}
 }