Make HTTP middleware easier to use

appengine/tumble/service.go

 // Copyright 2015 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 tumble
 
 import (
         "fmt"
         "net/http"
         "strconv"
         "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"
 )
 
 const transientHTTPHeader = "X-LUCI-Tumble-Transient"
 
 // Service is an instance of a Tumble service. It installs its handlers into an
 // HTTP router and services Tumble request tasks.
 type Service struct {
         // Middleware is an optional function which allows your application to add
         // application-specific resources to the context used by ProcessShardHandler.
         //
         // Context will already be setup with BaseProd.
         Middleware func(context.Context) context.Context
 
         // Namespaces is a function that returns the datastore namespaces that Tumble
         // will poll.
         //
         // If nil, Tumble will be executed against all namespaces registered in the
         // datastore.
         Namespaces func(context.Context) ([]string, error)
 }
 
 // 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
 // a httprouter. It expects `logging` and `luci/gae` services to be installed
 // into the context.
 //
 // 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"))
         }
 }
 
 // FireAllTasks searches for work in all namespaces, and fires off a process
 // task for any shards it finds that have at least one Mutation present to
 // ensure that no work languishes forever. This may not be needed in
 // a constantly-loaded system with good tumble key distribution.
 func (s *Service) FireAllTasks(c context.Context) error {
         cfg := getConfig(c)
         shards := make(map[taskShard]struct{}, cfg.NumShards)
@@ skipping 105 matching lines @@
 // ProcessShardHandler is a http handler suitable for installation into
 // a httprouter. It expects `logging` and `luci/gae` services to be installed
 // into the context.
 //
 // ProcessShardHandler verifies that its being run as a taskqueue task and that
 // the following parameters exist and are well-formed:
 //   * timestamp: decimal-encoded UNIX/UTC timestamp in seconds.
 //   * 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"))
         }
 }
 
 // getDatastoreNamespaces returns a list of all of the namespaces in the
 // datastore.
 //
 // This is done by issuing a datastore query for kind "__namespace__". The
 // resulting keys will have IDs for the namespaces, namely:
 //      - The default namespace will have integer ID 1. We ignore this because if
 //        we're using tumble with namespaces, we don't process the default
@@ skipping 11 matching lines @@
         }
 
         namespaces := make([]string, 0, len(namespaceKeys))
         for _, nk := range namespaceKeys {
                 if ns := nk.StringID(); ns != "" {
                         namespaces = append(namespaces, ns)
                 }
         }
         return namespaces, nil
 }