snapshot of https://github.com/ThreeDotsLabs/watermill.git (master, 2026-09-15; blobs > 131072B omitted)

message/router/middleware/timeout.go

@@ -0,0 +1,24 @@
+package middleware
+
+import (
+	"context"
+	"time"
+
+	"github.com/ThreeDotsLabs/watermill/message"
+)
+
+// Timeout makes the handler cancel the incoming message's context after a specified time.
+// Any timeout-sensitive functionality of the handler should listen on msg.Context().Done() to know when to fail.
+func Timeout(timeout time.Duration) func(message.HandlerFunc) message.HandlerFunc {
+	return func(h message.HandlerFunc) message.HandlerFunc {
+		return func(msg *message.Message) ([]*message.Message, error) {
+			ctx, cancel := context.WithTimeout(msg.Context(), timeout)
+			defer func() {
+				cancel()
+			}()
+
+			msg.SetContext(ctx)
+			return h(msg)
+		}
+	}
+}