Make HTTP middleware easier to use

appengine/tsmon/middleware.go

 // Copyright 2016 The LUCI Authors. All rights reserved.
 // Use of this source code is governed under the Apache License, Version 2.0
 // that can be found in the LICENSE file.
 
 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
@@ skipping 37 matching lines @@
 }
 
 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)
@@ skipping 207 matching lines @@
                 }
         }
 
         if err := state.Flush(c, mon); err != nil {
                 return err
         }
 
         state.ResetGlobalCallbackMetrics(c)
         return nil
 }