Make HTTP middleware easier to use

appengine/tumble/service.go

Index: appengine/tumble/service.go
diff --git a/appengine/tumble/service.go b/appengine/tumble/service.go
index fa1efbb6f4cbaee980b41c6fe53a377a06fb502e..9f408f67c505965ad4e4142449e2fc6f6d8957d3 100644
--- a/appengine/tumble/service.go
+++ b/appengine/tumble/service.go
@@ -11,13 +11,13 @@ import (
 	"sync"
 	"time"
 
-	"github.com/julienschmidt/httprouter"
 	"github.com/luci/gae/service/datastore"
 	"github.com/luci/gae/service/info"
 	"github.com/luci/luci-go/appengine/gaemiddleware"
 	"github.com/luci/luci-go/common/errors"
 	"github.com/luci/luci-go/common/logging"
 	"github.com/luci/luci-go/common/parallel"
+	"github.com/luci/luci-go/server/router"
 	"golang.org/x/net/context"
 )
 
@@ -41,13 +41,12 @@ type Service struct {
 }
 
 // InstallHandlers installs http handlers.
-func (s *Service) InstallHandlers(r *httprouter.Router) {
+func (s *Service) InstallHandlers(r *router.Router) {
 	// GET so that this can be invoked from cron
 	r.GET(fireAllTasksURL,
-		gaemiddleware.BaseProd(gaemiddleware.RequireCron(s.FireAllTasksHandler)))
+		append(gaemiddleware.BaseProd(), gaemiddleware.RequireCron(), s.FireAllTasksHandler)...)
 
-	r.POST(processShardPattern,
-		gaemiddleware.BaseProd(gaemiddleware.RequireTaskQueue(baseName, s.ProcessShardHandler)))
+	r.POST(processShardPattern, append(gaemiddleware.BaseProd(), gaemiddleware.RequireTaskQueue(baseName), s.ProcessShardHandler)...)
 }
 
 // FireAllTasksHandler is a http handler suitable for installation into
@@ -56,12 +55,12 @@ func (s *Service) InstallHandlers(r *httprouter.Router) {
 //
 // FireAllTasksHandler verifies that it was called within an Appengine Cron
 // request, and then invokes the FireAllTasks function.
-func (s *Service) FireAllTasksHandler(c context.Context, rw http.ResponseWriter, r *http.Request, _ httprouter.Params) {
-	if err := s.FireAllTasks(c); err != nil {
-		rw.WriteHeader(http.StatusInternalServerError)
-		fmt.Fprintf(rw, "fire_all_tasks failed: %s", err)
+func (s *Service) FireAllTasksHandler(c *router.Context) {
+	if err := s.FireAllTasks(c.Context); err != nil {
+		c.Writer.WriteHeader(http.StatusInternalServerError)
+		fmt.Fprintf(c.Writer, "fire_all_tasks failed: %s", err)
 	} else {
-		rw.Write([]byte("ok"))
+		c.Writer.Write([]byte("ok"))
 	}
 }
 
@@ -187,50 +186,50 @@ func (s *Service) getNamespaces(c context.Context, cfg *Config) (namespaces []st
 //   * shard_id: decimal-encoded shard identifier.
 //
 // ProcessShardHandler then invokes ProcessShard with the parsed parameters.
-func (s *Service) ProcessShardHandler(c context.Context, rw http.ResponseWriter, r *http.Request, p httprouter.Params) {
+func (s *Service) ProcessShardHandler(c *router.Context) {
 	if s.Middleware != nil {
-		c = s.Middleware(c)
+		c.Context = s.Middleware(c.Context)
 	}
 
-	tstampStr := p.ByName("timestamp")
-	sidStr := p.ByName("shard_id")
+	tstampStr := c.Params.ByName("timestamp")
+	sidStr := c.Params.ByName("shard_id")
 
 	tstamp, err := strconv.ParseInt(tstampStr, 10, 64)
 	if err != nil {
-		logging.Errorf(c, "bad timestamp %q", tstampStr)
-		rw.WriteHeader(http.StatusNotFound)
-		fmt.Fprintf(rw, "bad timestamp")
+		logging.Errorf(c.Context, "bad timestamp %q", tstampStr)
+		c.Writer.WriteHeader(http.StatusNotFound)
+		fmt.Fprintf(c.Writer, "bad timestamp")
 		return
 	}
 
 	sid, err := strconv.ParseUint(sidStr, 10, 64)
 	if err != nil {
-		logging.Errorf(c, "bad shardID %q", tstampStr)
-		rw.WriteHeader(http.StatusNotFound)
-		fmt.Fprintf(rw, "bad shardID")
+		logging.Errorf(c.Context, "bad shardID %q", tstampStr)
+		c.Writer.WriteHeader(http.StatusNotFound)
+		fmt.Fprintf(c.Writer, "bad shardID")
 		return
 	}
 
-	cfg := getConfig(c)
+	cfg := getConfig(c.Context)
 
 	// Get the set of namespaces to handle.
-	namespaces, err := s.getNamespaces(c, cfg)
+	namespaces, err := s.getNamespaces(c.Context, cfg)
 	if err != nil {
-		rw.WriteHeader(http.StatusInternalServerError)
+		c.Writer.WriteHeader(http.StatusInternalServerError)
 		return
 	}
 
-	err = processShard(c, cfg, namespaces, time.Unix(tstamp, 0).UTC(), sid)
+	err = processShard(c.Context, cfg, namespaces, time.Unix(tstamp, 0).UTC(), sid)
 	if err != nil {
-		logging.Errorf(c, "failure! %s", err)
+		logging.Errorf(c.Context, "failure! %s", err)
 
 		if errors.IsTransient(err) {
-			rw.Header().Add(transientHTTPHeader, "true")
+			c.Writer.Header().Add(transientHTTPHeader, "true")
 		}
-		rw.WriteHeader(http.StatusInternalServerError)
-		fmt.Fprintf(rw, "error: %s", err)
+		c.Writer.WriteHeader(http.StatusInternalServerError)
+		fmt.Fprintf(c.Writer, "error: %s", err)
 	} else {
-		rw.Write([]byte("ok"))
+		c.Writer.Write([]byte("ok"))
 	}
 }