Add Use method to better support middlewares with ServeMux

servemux.go

@@ -23,9 +23,10 @@ import (
 // "images:thumbnails" and the former will receive tasks with type name beginning
 // with "images".
 type ServeMux struct {
-	mu sync.RWMutex
-	m  map[string]muxEntry
-	es []muxEntry // slice of entries sorted from longest to shortest.
+	mu  sync.RWMutex
+	m   map[string]muxEntry
+	es  []muxEntry // slice of entries sorted from longest to shortest.
+	mws []MiddlewareFunc
 }
 
 type muxEntry struct {
@@ -33,6 +34,11 @@ type muxEntry struct {
 	pattern string
 }
 
+// MiddlewareFunc is a function which receives an asynq.Handler and returns another asynq.Handler.
+// Typically, the returned handler is a closure which does something with the context and task passed
+// to it, and then calls the handler passed as parameter to the MiddlewareFunc.
+type MiddlewareFunc func(Handler) Handler
+
 // NewServeMux allocates and returns a new ServeMux.
 func NewServeMux() *ServeMux {
 	return new(ServeMux)
@@ -60,6 +66,9 @@ func (mux *ServeMux) Handler(t *Task) (h Handler, pattern string) {
 	if h == nil {
 		h, pattern = NotFoundHandler(), ""
 	}
+	for i := len(mux.mws) - 1; i >= 0; i-- {
+		h = mux.mws[i](h)
+	}
 	return h, pattern
 }
 
@@ -130,6 +139,16 @@ func (mux *ServeMux) HandleFunc(pattern string, handler func(context.Context, *T
 	mux.Handle(pattern, HandlerFunc(handler))
 }
 
+// Use appends a MiddlewareFunc to the chain.
+// Middlewares are executed in the order that they are applied to the ServeMux.
+func (mux *ServeMux) Use(mws ...MiddlewareFunc) {
+	mux.mu.Lock()
+	defer mux.mu.Unlock()
+	for _, fn := range mws {
+		mux.mws = append(mux.mws, fn)
+	}
+}
+
 // NotFound returns an error indicating that the handler was not found for the given task.
 func NotFound(ctx context.Context, task *Task) error {
 	return fmt.Errorf("handler not found for task %q", task.Type)