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

message/router/middleware/circuit_breaker.go

@@ -0,0 +1,37 @@
+package middleware
+
+import (
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/sony/gobreaker"
+)
+
+// CircuitBreaker is a middleware that wraps the handler in a circuit breaker.
+// Based on the configuration, the circuit breaker will fail fast if the handler keeps returning errors.
+// This is useful for preventing cascading failures.
+type CircuitBreaker struct {
+	cb *gobreaker.CircuitBreaker
+}
+
+// NewCircuitBreaker returns a new CircuitBreaker middleware.
+// Refer to the gobreaker documentation for the available settings.
+func NewCircuitBreaker(settings gobreaker.Settings) CircuitBreaker {
+	return CircuitBreaker{
+		cb: gobreaker.NewCircuitBreaker(settings),
+	}
+}
+
+// Middleware returns the CircuitBreaker middleware.
+func (c CircuitBreaker) Middleware(h message.HandlerFunc) message.HandlerFunc {
+	return func(msg *message.Message) ([]*message.Message, error) {
+		out, err := c.cb.Execute(func() (interface{}, error) {
+			return h(msg)
+		})
+
+		var result []*message.Message
+		if out != nil {
+			result = out.([]*message.Message)
+		}
+
+		return result, err
+	}
+}

message/router/middleware/circuit_breaker_test.go

@@ -0,0 +1,61 @@
+package middleware_test
+
+import (
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/ThreeDotsLabs/watermill/message/router/middleware"
+	"github.com/sony/gobreaker"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCircuitBreaker(t *testing.T) {
+	t.Parallel()
+
+	count := 0
+	failing := true
+
+	h := middleware.NewCircuitBreaker(
+		gobreaker.Settings{
+			Name:    "test",
+			Timeout: time.Millisecond * 50,
+		},
+	).Middleware(func(msg *message.Message) (messages []*message.Message, e error) {
+		count++
+
+		if failing {
+			return nil, errors.New("test error")
+		}
+
+		return nil, nil
+	})
+
+	msg := message.NewMessage("1", nil)
+
+	// The first 6 calls should fail and increment the count
+	for i := 0; i < 6; i++ {
+		_, err := h(msg)
+		assert.Error(t, err)
+	}
+
+	assert.Equal(t, 6, count)
+
+	// The next calls should fail and not increment the count (the circuit breaker is open)
+	for i := 0; i < 4; i++ {
+		_, err := h(msg)
+		assert.Error(t, err)
+	}
+	assert.Equal(t, 6, count)
+
+	time.Sleep(time.Millisecond * 100)
+	failing = false
+
+	// After a timeout, the Circuit Breaker is closed again
+	for i := 0; i < 4; i++ {
+		_, err := h(msg)
+		assert.NoError(t, err)
+	}
+	assert.Equal(t, 10, count)
+}

message/router/middleware/correlation.go

@@ -0,0 +1,43 @@
+package middleware
+
+import (
+	"github.com/ThreeDotsLabs/watermill/message"
+)
+
+// CorrelationIDMetadataKey is used to store the correlation ID in metadata.
+const CorrelationIDMetadataKey = "correlation_id"
+
+// SetCorrelationID sets a correlation ID for the message.
+//
+// SetCorrelationID should be called when the message enters the system.
+// When message is produced in a request (for example HTTP),
+// message correlation ID should be the same as the request's correlation ID.
+func SetCorrelationID(id string, msg *message.Message) {
+	if MessageCorrelationID(msg) != "" {
+		return
+	}
+
+	msg.Metadata.Set(CorrelationIDMetadataKey, id)
+}
+
+// MessageCorrelationID returns correlation ID from the message.
+func MessageCorrelationID(message *message.Message) string {
+	return message.Metadata.Get(CorrelationIDMetadataKey)
+}
+
+// CorrelationID adds correlation ID to all messages produced by the handler.
+// ID is based on ID from message received by handler.
+//
+// To make CorrelationID working correctly, SetCorrelationID must be called to first message entering the system.
+func CorrelationID(h message.HandlerFunc) message.HandlerFunc {
+	return func(message *message.Message) ([]*message.Message, error) {
+		producedMessages, err := h(message)
+
+		correlationID := MessageCorrelationID(message)
+		for _, msg := range producedMessages {
+			SetCorrelationID(correlationID, msg)
+		}
+
+		return producedMessages, err
+	}
+}

message/router/middleware/correlation_test.go

@@ -0,0 +1,39 @@
+package middleware_test
+
+import (
+	"testing"
+
+	"github.com/pkg/errors"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/ThreeDotsLabs/watermill/message/router/middleware"
+
+	"github.com/ThreeDotsLabs/watermill/message"
+)
+
+func TestCorrelationID(t *testing.T) {
+	handlerErr := errors.New("foo")
+
+	handler := middleware.CorrelationID(func(msg *message.Message) ([]*message.Message, error) {
+		return message.Messages{message.NewMessage("2", nil)}, handlerErr
+	})
+
+	msg := message.NewMessage("1", nil)
+	middleware.SetCorrelationID("correlation_id", msg)
+
+	producedMsgs, err := handler(msg)
+
+	assert.Equal(t, "2", producedMsgs[0].UUID)
+	assert.Equal(t, middleware.MessageCorrelationID(producedMsgs[0]), "correlation_id")
+	assert.Equal(t, handlerErr, err)
+}
+
+func TestSetCorrelationID_already_set(t *testing.T) {
+	msg := message.NewMessage("", nil)
+
+	middleware.SetCorrelationID("foo", msg)
+	middleware.SetCorrelationID("bar", msg)
+
+	assert.Equal(t, "foo", middleware.MessageCorrelationID(msg))
+}

message/router/middleware/deduplicator.go

@@ -0,0 +1,314 @@
+package middleware
+
+import (
+	"bytes"
+	"context"
+	"crypto/sha256"
+	"fmt"
+	"hash/adler32"
+	"io"
+	"math"
+	"sync"
+	"time"
+
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/pkg/errors"
+)
+
+// MessageHasherReadLimitMinimum specifies the least number
+// of bytes of a [message.Message] are used for calculating
+// their hash values using a [MessageHasher].
+const MessageHasherReadLimitMinimum = 64
+
+// ExpiringKeyRepository is a state container for checking the
+// existence of a key in a certain time window.
+// All operations must be safe for concurrent use.
+type ExpiringKeyRepository interface {
+	// IsDuplicate returns `true` if the key
+	// was not checked in recent past.
+	// The key must expire in a certain time window.
+	IsDuplicate(ctx context.Context, key string) (ok bool, err error)
+}
+
+// MessageHasher returns a short tag that describes
+// a message. The tag should be unique per message,
+// but avoiding hash collisions entirely is not practical
+// for performance reasons. Used for powering [Deduplicator]s.
+type MessageHasher func(*message.Message) (string, error)
+
+// Deduplicator drops similar messages if they are present
+// in a [ExpiringKeyRepository]. The similarity is determined
+// by a [MessageHasher]. Time out is applied to repository
+// operations using [context.WithTimeout].
+//
+// Call [Deduplicator.Middleware] for a new middleware
+// or [Deduplicator.Decorator] for a [message.PublisherDecorator].
+//
+// KeyFactory defaults to [NewMessageHasherAdler32] with read
+// limit  set to [math.MaxInt64] for fast tagging.
+// Use [NewMessageHasherSHA256] for minimal collisions.
+//
+// Repository defaults to [NewMapExpiringKeyRepository] with one
+// minute retention window. This default setting is performant
+// but **does not support distributed operations**. If you
+// implement a [ExpiringKeyRepository] backed by Redis,
+// please submit a pull request.
+//
+// Timeout defaults to one minute. If lower than
+// five milliseconds, it is set to five milliseconds.
+//
+// [ExpiringKeyRepository] must expire values
+// in a certain time window. If there is no expiration, only one
+// unique message will be ever delivered as long as the repository
+// keeps its state.
+type Deduplicator struct {
+	KeyFactory MessageHasher
+	Repository ExpiringKeyRepository
+	Timeout    time.Duration
+}
+
+// IsDuplicate returns true if the message hash tag calculated
+// using a [MessageHasher] was seen in deduplication time window.
+func (d *Deduplicator) IsDuplicate(m *message.Message) (bool, error) {
+	key, err := d.KeyFactory(m)
+	if err != nil {
+		return false, err
+	}
+	ctx, cancel := context.WithTimeout(m.Context(), d.Timeout)
+	defer cancel()
+	return d.Repository.IsDuplicate(ctx, key)
+}
+
+func applyDefaultsToDeduplicator(d *Deduplicator) *Deduplicator {
+	if d == nil {
+		kr, err := NewMapExpiringKeyRepository(time.Minute)
+		if err != nil {
+			panic(err)
+		}
+		return &Deduplicator{
+			KeyFactory: NewMessageHasherAdler32(math.MaxInt64),
+			Repository: kr,
+			Timeout:    time.Minute,
+		}
+	}
+	if d.KeyFactory == nil {
+		d.KeyFactory = NewMessageHasherAdler32(math.MaxInt64)
+	}
+	if d.Repository == nil {
+		kr, err := NewMapExpiringKeyRepository(time.Minute)
+		if err != nil {
+			panic(err)
+		}
+		d.Repository = kr
+	}
+	if d.Timeout < time.Millisecond*5 {
+		d.Timeout = time.Millisecond * 5
+	}
+	return d
+}
+
+// Middleware returns the [message.HandlerMiddleware]
+// that drops similar messages in a given time window.
+func (d *Deduplicator) Middleware(h message.HandlerFunc) message.HandlerFunc {
+	d = applyDefaultsToDeduplicator(d)
+	return func(msg *message.Message) ([]*message.Message, error) {
+		isDuplicate, err := d.IsDuplicate(msg)
+		if err != nil {
+			return nil, err
+		}
+		if isDuplicate {
+			return nil, nil
+		}
+		return h(msg)
+	}
+}
+
+type mapExpiringKeyRepository struct {
+	window time.Duration
+	mu     *sync.Mutex
+	tags   map[string]time.Time
+}
+
+// NewMapExpiringKeyRepository returns a memory store
+// backed by a regular hash map protected by
+// a [sync.Mutex]. The state **cannot be shared or synchronized
+// between instances** by design for performance.
+//
+// If you need to drop duplicate messages by orchestration,
+// implement [ExpiringKeyRepository] interface backed by Redis
+// or similar.
+//
+// Window specifies the minimum duration of how long the
+// duplicate tags are remembered for. Real duration can
+// extend up to 50% longer because it depends on the
+// clean up cycle.
+func NewMapExpiringKeyRepository(window time.Duration) (ExpiringKeyRepository, error) {
+	if window < time.Millisecond {
+		return nil, errors.New("deduplication window of less than a millisecond is impractical")
+	}
+
+	kr := &mapExpiringKeyRepository{
+		window: window,
+		mu:     &sync.Mutex{},
+		tags:   make(map[string]time.Time),
+	}
+	ticker := time.NewTicker(window / 2)
+
+	go kr.cleanOutLoop(context.Background(), ticker)
+	return kr, nil
+}
+
+func (kr *mapExpiringKeyRepository) IsDuplicate(
+	ctx context.Context,
+	key string,
+) (bool, error) {
+	kr.mu.Lock()
+	_, alreadySeen := kr.tags[key]
+	if alreadySeen {
+		// NOTE: could also check if key expires.After(t)
+		// and remove it for exact expiration
+		// instead of fuzzy until-next clean up expiration
+		// but this should not be needed for most use cases.
+		kr.mu.Unlock()
+		return true, nil
+	}
+	kr.tags[key] = time.Now().Add(kr.window)
+	kr.mu.Unlock()
+	return false, nil
+}
+
+func (kr *mapExpiringKeyRepository) cleanOutLoop(ctx context.Context, ticker *time.Ticker) {
+	for {
+		select {
+		case <-ctx.Done():
+			return // execution ended, part the go routine
+		case tagsBefore := <-ticker.C:
+			kr.cleanOut(tagsBefore)
+		}
+	}
+}
+
+func (kr *mapExpiringKeyRepository) cleanOut(tagsBefore time.Time) {
+	kr.mu.Lock()
+	defer kr.mu.Unlock()
+
+	for hash, expires := range kr.tags {
+		if expires.Before(tagsBefore) {
+			delete(kr.tags, hash)
+		}
+	}
+}
+
+// Len returns the number of known tags that have not been
+// cleaned out yet.
+func (kr *mapExpiringKeyRepository) Len() (count int) {
+	kr.mu.Lock()
+	count = len(kr.tags)
+	kr.mu.Unlock()
+	return
+}
+
+// NewMessageHasherAdler32 generates message hashes using a fast
+// Adler-32 checksum of the [message.Message] body. Read
+// limit specifies how many bytes of the message are
+// used for calculating the hash.
+//
+// Lower limit improves performance but results in more false
+// positives. Read limit must be greater than
+// [MessageHasherReadLimitMinimum].
+func NewMessageHasherAdler32(readLimit int64) MessageHasher {
+	if readLimit < MessageHasherReadLimitMinimum {
+		readLimit = MessageHasherReadLimitMinimum
+	}
+	return func(m *message.Message) (string, error) {
+		h := adler32.New()
+		_, err := io.CopyN(h, bytes.NewReader(m.Payload), readLimit)
+		if err != nil && err != io.EOF {
+			return "", err
+		}
+		return string(h.Sum(nil)), nil
+	}
+}
+
+// NewMessageHasherSHA256 generates message hashes using a slower
+// but more resilient hashing of the [message.Message] body. Read
+// limit specifies how many bytes of the message are
+// used for calculating the hash.
+//
+// Lower limit improves performance but results in more false
+// positives. Read limit must be greater than
+// [MessageHasherReadLimitMinimum].
+func NewMessageHasherSHA256(readLimit int64) MessageHasher {
+	if readLimit < MessageHasherReadLimitMinimum {
+		readLimit = MessageHasherReadLimitMinimum
+	}
+
+	return func(m *message.Message) (string, error) {
+		h := sha256.New()
+		_, err := io.CopyN(h, bytes.NewReader(m.Payload), readLimit)
+		if err != nil && err != io.EOF {
+			return "", err
+		}
+		return string(h.Sum(nil)), nil
+	}
+}
+
+// NewMessageHasherFromMetadataField looks for a hash value
+// inside message metadata instead of calculating a new one.
+// Useful if a [MessageHasher] was applied in a previous
+// [message.HandlerFunc].
+func NewMessageHasherFromMetadataField(field string) MessageHasher {
+	return func(m *message.Message) (string, error) {
+		fromMetadata, ok := m.Metadata[field]
+		if ok {
+			return fromMetadata, nil
+		}
+		return "", fmt.Errorf("cannot recover hash value from metadata of message #%s: field %q is absent", m.UUID, field)
+	}
+}
+
+type deduplicatingPublisherDecorator struct {
+	message.Publisher
+	deduplicator *Deduplicator
+}
+
+func (d *deduplicatingPublisherDecorator) Publish(
+	topic string,
+	messages ...*message.Message,
+) (err error) {
+	notRecent := make([]*message.Message, 0, len(messages))
+	isDuplicate := false
+
+	for _, m := range messages {
+		isDuplicate, err = d.deduplicator.IsDuplicate(m)
+		if err != nil {
+			return err
+		}
+		if isDuplicate {
+			m.Ack() // acknowledge and ignore
+			continue
+		}
+		notRecent = append(notRecent, m)
+	}
+	return d.Publisher.Publish(topic, notRecent...)
+}
+
+// PublisherDecorator returns a decorator that
+// acknowledges and drops every [message.Message] that
+// was recognized by a [Deduplicator].
+//
+// The returned decorator provides the same functionality
+// to a [message.Publisher] as [Deduplicator.Middleware]
+// to a [message.Router].
+func (d *Deduplicator) PublisherDecorator() message.PublisherDecorator {
+	return func(pub message.Publisher) (message.Publisher, error) {
+		if pub == nil {
+			return nil, errors.New("cannot decorate a <nil> publisher")
+		}
+
+		return &deduplicatingPublisherDecorator{
+			Publisher:    pub,
+			deduplicator: applyDefaultsToDeduplicator(d),
+		}, nil
+	}
+}

message/router/middleware/deduplicator_test.go

@@ -0,0 +1,211 @@
+package middleware_test
+
+import (
+	"context"
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/ThreeDotsLabs/watermill/message/router/middleware"
+	"github.com/ThreeDotsLabs/watermill/pubsub/gochannel"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestDeduplicatorMiddleware(t *testing.T) {
+	t.Parallel()
+
+	count := 0
+	d := &middleware.Deduplicator{
+		KeyFactory: middleware.NewMessageHasherAdler32(1024),
+		// KeyFactory: middleware.NewMessageHasherSHA256(1024),
+		Timeout: time.Second,
+	}
+	h := d.Middleware(func(msg *message.Message) (messages []*message.Message, e error) {
+		count++
+		return nil, nil
+	})
+
+	for i := 0; i < 6; i++ { // only one should go through
+		msg := message.NewMessage(
+			fmt.Sprintf("first%d", i),
+			[]byte("1"),
+		)
+		_, err := h(msg)
+		assert.NoError(t, err)
+	}
+
+	for i := 0; i < 2; i++ { // only one should go through
+		msg := message.NewMessage(
+			fmt.Sprintf("second%d", i),
+			[]byte("2"),
+		)
+		_, err := h(msg)
+		assert.NoError(t, err)
+	}
+
+	assert.Equal(t, 2, count)
+}
+
+func TestDeduplicatorPublisherDecorator(t *testing.T) {
+	t.Parallel()
+
+	pubSub := gochannel.NewGoChannel(gochannel.Config{
+		OutputChannelBuffer: 100,
+		Persistent:          true,
+	}, nil)
+	defer func() { _ = pubSub.Close() }()
+
+	const testDedupeTopic = "testTopic"
+	ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*50)
+	defer cancel()
+
+	d := &middleware.Deduplicator{
+		KeyFactory: middleware.NewMessageHasherAdler32(1024),
+		// KeyFactory: middleware.NewMessageHasherSHA256(1024),
+		Timeout: time.Second,
+	}
+	decorated, err := d.PublisherDecorator()(pubSub)
+	assert.NoError(t, err)
+
+	for i := 0; i < 6; i++ { // only one should go through
+		msg := message.NewMessage(
+			fmt.Sprintf("first%d", i),
+			[]byte("1"),
+		)
+		err := decorated.Publish(testDedupeTopic, msg)
+		assert.NoError(t, err)
+	}
+
+	for i := 0; i < 2; i++ { // only one should go through
+		msg := message.NewMessage(
+			fmt.Sprintf("second%d", i),
+			[]byte("2"),
+		)
+		err := decorated.Publish(testDedupeTopic, msg)
+		assert.NoError(t, err)
+	}
+
+	got, err := pubSub.Subscribe(ctx, testDedupeTopic)
+	assert.NoError(t, err)
+	count := 0
+	for m := range got {
+		count++
+		m.Ack()
+		t.Log("got message:", m.UUID)
+	}
+	assert.Equal(t, 2, count)
+}
+
+func TestMessageHasherAdler32(t *testing.T) {
+	t.Parallel()
+
+	short := middleware.NewMessageHasherAdler32(0)
+	full := middleware.NewMessageHasherAdler32(middleware.MessageHasherReadLimitMinimum)
+
+	msg := message.NewMessage("adlerTest", []byte("some random data"))
+	h1, err := short(msg)
+	assert.NoError(t, err)
+	h2, err := full(msg)
+	assert.NoError(t, err)
+
+	if h1 != h2 {
+		t.Fatal("MessageHasherReadLimitMinimum did not apply to Adler32 message hasher")
+	}
+}
+
+func TestMessageHasherSHA256(t *testing.T) {
+	t.Parallel()
+
+	short := middleware.NewMessageHasherSHA256(0)
+	full := middleware.NewMessageHasherSHA256(middleware.MessageHasherReadLimitMinimum)
+
+	msg := message.NewMessage("adlerTest", []byte("some random data"))
+	h1, err := short(msg)
+	assert.NoError(t, err)
+	h2, err := full(msg)
+	assert.NoError(t, err)
+
+	if h1 != h2 {
+		t.Fatal("MessageHasherReadLimitMinimum did not apply to SHA256 message hasher")
+	}
+}
+
+func TestMessageHasherFromMetadataField(t *testing.T) {
+	t.Parallel()
+
+	field := "hash"
+	value := "someHash"
+	msg := message.NewMessage("one", []byte("1"))
+	msg.Metadata[field] = value
+	metadataPull := middleware.NewMessageHasherFromMetadataField(field)
+
+	h, err := metadataPull(msg)
+	assert.NoError(t, err)
+	assert.Equal(t, h, value)
+
+	delete(msg.Metadata, field) // empty out
+	_, err = metadataPull(msg)
+	assert.Error(t, err)
+}
+
+func TestMapExpiringKeyRepositoryCleanup(t *testing.T) {
+	t.Parallel()
+	wait := time.Millisecond * 5
+	kr, err := middleware.NewMapExpiringKeyRepository(wait)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	count := 0
+	d := &middleware.Deduplicator{
+		Repository: kr,
+		KeyFactory: middleware.NewMessageHasherAdler32(1024),
+		Timeout:    time.Second,
+	}
+	h := d.Middleware(func(msg *message.Message) (messages []*message.Message, e error) {
+		count++
+		return nil, nil
+	})
+
+	for i := 0; i < 6; i++ { // only one should go through
+		msg := message.NewMessage(
+			fmt.Sprintf("expiring%d", i),
+			fmt.Appendf(nil, "expiring%d", i),
+		)
+		_, err := h(msg)
+		assert.NoError(t, err)
+	}
+
+	type supportsLen interface {
+		Len() int
+	}
+	measurable, ok := kr.(supportsLen)
+	if !ok {
+		t.Fatal("repository does not allow measuring its length")
+	}
+
+	if l := measurable.Len(); l != 6 {
+		t.Errorf("expected 6 tags, but %d remain", l)
+	}
+
+	assert.Eventually(
+		t,
+		func() bool {
+			return count == 6
+		},
+		wait*3,
+		time.Millisecond,
+		"sent six messages, but only received %d", count,
+	)
+	assert.Eventually(
+		t,
+		func() bool {
+			return measurable.Len() == 0
+		},
+		wait*3,
+		time.Millisecond,
+		"tags should have been cleaned out, but %d remain",
+		measurable.Len(),
+	)
+}

message/router/middleware/delay_on_error.go

@@ -0,0 +1,47 @@
+package middleware
+
+import (
+	"time"
+
+	"github.com/ThreeDotsLabs/watermill/components/delay"
+	"github.com/ThreeDotsLabs/watermill/message"
+)
+
+// DelayOnError is a middleware that adds the delay metadata to the message if an error occurs.
+//
+// IMPORTANT: The delay metadata doesn't cause delays with all Pub/Subs! Using it won't have any effect on Pub/Subs that don't support it.
+// See the list of supported Pub/Subs in the documentation: https://watermill.io/advanced/delayed-messages/
+type DelayOnError struct {
+	// InitialInterval is the first interval between retries. Subsequent intervals will be scaled by Multiplier.
+	InitialInterval time.Duration
+	// MaxInterval sets the limit for the exponential backoff of retries. The interval will not be increased beyond MaxInterval.
+	MaxInterval time.Duration
+	// Multiplier is the factor by which the waiting interval will be multiplied between retries.
+	Multiplier float64
+}
+
+func (d *DelayOnError) Middleware(h message.HandlerFunc) message.HandlerFunc {
+	return func(msg *message.Message) ([]*message.Message, error) {
+		msgs, err := h(msg)
+		if err != nil {
+			d.applyDelay(msg)
+		}
+
+		return msgs, err
+	}
+}
+
+func (d *DelayOnError) applyDelay(msg *message.Message) {
+	delayedForStr := msg.Metadata.Get(delay.DelayedForKey)
+	delayedFor, err := time.ParseDuration(delayedForStr)
+	if delayedForStr != "" && err == nil {
+		delayedFor *= time.Duration(d.Multiplier)
+		if delayedFor > d.MaxInterval {
+			delayedFor = d.MaxInterval
+		}
+
+		delay.Message(msg, delay.For(delayedFor))
+	} else {
+		delay.Message(msg, delay.For(d.InitialInterval))
+	}
+}

message/router/middleware/delay_on_error_test.go

@@ -0,0 +1,58 @@
+package middleware_test
+
+import (
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/ThreeDotsLabs/watermill/components/delay"
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/ThreeDotsLabs/watermill/message/router/middleware"
+)
+
+func TestDelayOnError(t *testing.T) {
+	m := middleware.DelayOnError{
+		InitialInterval: time.Second,
+		MaxInterval:     time.Second * 10,
+		Multiplier:      2,
+	}
+
+	msg := message.NewMessage("1", []byte("test"))
+
+	getDelayFor := func(msg *message.Message) string {
+		return msg.Metadata.Get(delay.DelayedForKey)
+	}
+
+	okHandler := func(msg *message.Message) ([]*message.Message, error) {
+		return nil, nil
+	}
+
+	errHandler := func(msg *message.Message) ([]*message.Message, error) {
+		return nil, errors.New("error")
+	}
+
+	assert.Equal(t, "", getDelayFor(msg))
+
+	_, _ = m.Middleware(okHandler)(msg)
+	assert.Equal(t, "", getDelayFor(msg))
+
+	_, _ = m.Middleware(errHandler)(msg)
+	assert.Equal(t, "1s", getDelayFor(msg))
+
+	_, _ = m.Middleware(errHandler)(msg)
+	assert.Equal(t, "2s", getDelayFor(msg))
+
+	_, _ = m.Middleware(errHandler)(msg)
+	assert.Equal(t, "4s", getDelayFor(msg))
+
+	_, _ = m.Middleware(errHandler)(msg)
+	assert.Equal(t, "8s", getDelayFor(msg))
+
+	_, _ = m.Middleware(errHandler)(msg)
+	assert.Equal(t, "10s", getDelayFor(msg))
+
+	_, _ = m.Middleware(errHandler)(msg)
+	assert.Equal(t, "10s", getDelayFor(msg))
+}

message/router/middleware/duplicator.go

@@ -0,0 +1,22 @@
+package middleware
+
+import (
+	"github.com/ThreeDotsLabs/watermill/message"
+)
+
+// Duplicator is processing messages twice, to ensure that the endpoint is idempotent.
+func Duplicator(h message.HandlerFunc) message.HandlerFunc {
+	return func(msg *message.Message) ([]*message.Message, error) {
+		firstProducedMessages, firstErr := h(msg)
+		if firstErr != nil {
+			return nil, firstErr
+		}
+
+		secondProducedMessages, secondErr := h(msg)
+		if secondErr != nil {
+			return nil, secondErr
+		}
+
+		return append(firstProducedMessages, secondProducedMessages...), nil
+	}
+}

message/router/middleware/duplicator_test.go

@@ -0,0 +1,47 @@
+package middleware_test
+
+import (
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/ThreeDotsLabs/watermill/message/router/middleware"
+)
+
+var (
+	someMsg = message.NewMessage("1", nil)
+)
+
+func TestDuplicator(t *testing.T) {
+	var executionsCount int
+	producedMessages, err := middleware.Duplicator(func(msg *message.Message) ([]*message.Message, error) {
+		executionsCount++
+		return []*message.Message{msg}, nil
+	})(someMsg)
+
+	assert.NoError(t, err)
+	assert.Len(t, producedMessages, 2)
+	assert.Equal(t, "1", producedMessages[0].UUID)
+	assert.Equal(t, "1", producedMessages[1].UUID)
+	assert.Equal(t, 2, executionsCount)
+}
+
+func TestDuplicator_errors(t *testing.T) {
+	_, err := middleware.Duplicator(func(msg *message.Message) ([]*message.Message, error) {
+		return nil, errors.New("some error")
+	})(someMsg)
+	assert.Error(t, err, "some error")
+
+	var wasExecuted bool
+	_, err = middleware.Duplicator(func(msg *message.Message) ([]*message.Message, error) {
+		if wasExecuted {
+			return nil, errors.New("some other error")
+		}
+
+		wasExecuted = true
+		return []*message.Message{msg}, nil
+	})(someMsg)
+	assert.Error(t, err, "some other error")
+}

message/router/middleware/ignore_errors.go

@@ -0,0 +1,38 @@
+package middleware
+
+import (
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/pkg/errors"
+)
+
+// IgnoreErrors provides a middleware that makes the handler ignore some explicitly whitelisted errors.
+type IgnoreErrors struct {
+	ignoredErrors map[string]struct{}
+}
+
+// NewIgnoreErrors creates a new IgnoreErrors middleware.
+func NewIgnoreErrors(errs []error) IgnoreErrors {
+	errsMap := make(map[string]struct{}, len(errs))
+
+	for _, err := range errs {
+		errsMap[err.Error()] = struct{}{}
+	}
+
+	return IgnoreErrors{errsMap}
+}
+
+// Middleware returns the IgnoreErrors middleware.
+func (i IgnoreErrors) Middleware(h message.HandlerFunc) message.HandlerFunc {
+	return func(msg *message.Message) ([]*message.Message, error) {
+		events, err := h(msg)
+		if err != nil {
+			if _, ok := i.ignoredErrors[errors.Cause(err).Error()]; ok {
+				return events, nil
+			}
+
+			return events, err
+		}
+
+		return events, nil
+	}
+}

message/router/middleware/ignore_errors_test.go

@@ -0,0 +1,58 @@
+package middleware_test
+
+import (
+	"testing"
+
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/ThreeDotsLabs/watermill/message/router/middleware"
+	"github.com/pkg/errors"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestIgnoreErrors_Middleware(t *testing.T) {
+	testCases := []struct {
+		Name            string
+		IgnoredErrors   []error
+		TestError       error
+		ShouldBeIgnored bool
+	}{
+		{
+			Name:            "ignored_error",
+			IgnoredErrors:   []error{errors.New("test")},
+			TestError:       errors.New("test"),
+			ShouldBeIgnored: true,
+		},
+		{
+			Name:            "not_ignored_error",
+			IgnoredErrors:   []error{errors.New("test")},
+			TestError:       errors.New("not_ignored"),
+			ShouldBeIgnored: false,
+		},
+		{
+			Name:            "wrapped_error_should_ignore",
+			IgnoredErrors:   []error{errors.New("test")},
+			TestError:       errors.Wrap(errors.New("test"), "wrapped"),
+			ShouldBeIgnored: true,
+		},
+	}
+
+	for _, c := range testCases {
+		t.Run(c.Name, func(t *testing.T) {
+			m := middleware.NewIgnoreErrors(c.IgnoredErrors)
+
+			messagesToProduce := []*message.Message{message.NewMessage("1", nil)}
+
+			producedMessages, err := m.Middleware(func(msg *message.Message) ([]*message.Message, error) {
+				return messagesToProduce, c.TestError
+			})(nil)
+
+			if c.ShouldBeIgnored {
+				assert.NoError(t, err)
+			} else {
+				assert.Equal(t, c.TestError, err)
+			}
+
+			assert.Equal(t, messagesToProduce, producedMessages)
+		})
+	}
+}

message/router/middleware/instant_ack.go

@@ -0,0 +1,14 @@
+package middleware
+
+import "github.com/ThreeDotsLabs/watermill/message"
+
+// InstantAck makes the handler instantly acknowledge the incoming message, regardless of any errors.
+// It may be used to gain throughput, but at a cost:
+// If you had exactly-once delivery, you may expect at-least-once instead.
+// If you had ordered messages, the ordering might be broken.
+func InstantAck(h message.HandlerFunc) message.HandlerFunc {
+	return func(message *message.Message) ([]*message.Message, error) {
+		message.Ack()
+		return h(message)
+	}
+}

message/router/middleware/instant_ack_test.go

@@ -0,0 +1,33 @@
+package middleware
+
+import (
+	"testing"
+
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/pkg/errors"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestInstantAck(t *testing.T) {
+	producedMessages := message.Messages{message.NewMessage("2", nil)}
+	producedErr := errors.New("foo")
+
+	h := InstantAck(func(msg *message.Message) (messages []*message.Message, e error) {
+		return producedMessages, producedErr
+	})
+
+	msg := message.NewMessage("1", nil)
+
+	handlerMessages, handlerErr := h(msg)
+	assert.EqualValues(t, producedMessages, handlerMessages)
+	assert.Equal(t, producedErr, handlerErr)
+
+	select {
+	case <-msg.Acked():
+	// ok
+	case <-msg.Nacked():
+		t.Fatal("expected ack, not nack")
+	default:
+		t.Fatal("no ack received")
+	}
+}

message/router/middleware/message_test.go

@@ -0,0 +1,68 @@
+package middleware_test
+
+import (
+	"github.com/ThreeDotsLabs/watermill"
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/pkg/errors"
+)
+
+type mockPublisherBehaviour int
+
+const (
+	BehaviourAlwaysOK mockPublisherBehaviour = iota + 1
+	BehaviourAlwaysFail
+	BehaviourAlwaysPanic
+)
+
+var (
+	errClosed   = errors.New("closed")
+	errFailed   = errors.New("failed")
+	errPanicked = errors.New("panicked")
+)
+
+type mockPublisher struct {
+	behaviour mockPublisherBehaviour
+	closed    bool
+
+	produced []*message.Message
+}
+
+func (mp *mockPublisher) Publish(topic string, messages ...*message.Message) error {
+	if mp.closed {
+		return errClosed
+	}
+
+	switch mp.behaviour {
+	case BehaviourAlwaysOK:
+	case BehaviourAlwaysFail:
+		return errFailed
+	case BehaviourAlwaysPanic:
+		panic(errPanicked)
+	}
+
+	mp.produced = append(mp.produced, messages...)
+	return nil
+}
+
+func (mp *mockPublisher) Close() error {
+	mp.closed = true
+	return nil
+}
+
+func (mp *mockPublisher) PopMessages() []*message.Message {
+	defer func() { mp.produced = []*message.Message{} }()
+	return mp.produced
+}
+
+var handlerFuncAlwaysOKMessages = []*message.Message{
+	message.NewMessage(watermill.NewUUID(), nil),
+	message.NewMessage(watermill.NewUUID(), nil),
+}
+
+func handlerFuncAlwaysOK(*message.Message) ([]*message.Message, error) {
+	return handlerFuncAlwaysOKMessages, nil
+}
+
+func handlerFuncAlwaysFailing(*message.Message) ([]*message.Message, error) {
+	return nil, errFailed
+}

message/router/middleware/poison.go

@@ -0,0 +1,102 @@
+package middleware
+
+import (
+	stdErrors "errors"
+
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/pkg/errors"
+)
+
+// ErrInvalidPoisonQueueTopic occurs when the topic supplied to the PoisonQueue constructor is invalid.
+var ErrInvalidPoisonQueueTopic = errors.New("invalid poison queue topic")
+
+// Metadata keys that mark the reason and context for why the message was deemed poisoned.
+const (
+	ReasonForPoisonedKey  = "reason_poisoned"
+	PoisonedTopicKey      = "topic_poisoned"
+	PoisonedHandlerKey    = "handler_poisoned"
+	PoisonedSubscriberKey = "subscriber_poisoned"
+)
+
+type poisonQueue struct {
+	topic string
+	pub   message.Publisher
+
+	shouldGoToPoisonQueue func(err error) bool
+}
+
+// PoisonQueue provides a middleware that salvages unprocessable messages and publishes them on a separate topic.
+// The main middleware chain then continues on, business as usual.
+func PoisonQueue(pub message.Publisher, topic string) (message.HandlerMiddleware, error) {
+	if topic == "" {
+		return nil, ErrInvalidPoisonQueueTopic
+	}
+
+	pq := poisonQueue{
+		topic: topic,
+		pub:   pub,
+		shouldGoToPoisonQueue: func(err error) bool {
+			return true
+		},
+	}
+
+	return pq.Middleware, nil
+}
+
+// PoisonQueueWithFilter is just like PoisonQueue, but accepts a function that decides which errors qualify for the poison queue.
+func PoisonQueueWithFilter(pub message.Publisher, topic string, shouldGoToPoisonQueue func(err error) bool) (message.HandlerMiddleware, error) {
+	if topic == "" {
+		return nil, ErrInvalidPoisonQueueTopic
+	}
+
+	pq := poisonQueue{
+		topic: topic,
+		pub:   pub,
+
+		shouldGoToPoisonQueue: shouldGoToPoisonQueue,
+	}
+
+	return pq.Middleware, nil
+}
+
+func (pq poisonQueue) publishPoisonMessage(msg *message.Message, err error) error {
+	// no problems encountered, carry on
+	if err == nil {
+		return nil
+	}
+
+	// add context why it was poisoned
+	msg.Metadata.Set(ReasonForPoisonedKey, err.Error())
+	msg.Metadata.Set(PoisonedTopicKey, message.SubscribeTopicFromCtx(msg.Context()))
+	msg.Metadata.Set(PoisonedHandlerKey, message.HandlerNameFromCtx(msg.Context()))
+	msg.Metadata.Set(PoisonedSubscriberKey, message.SubscriberNameFromCtx(msg.Context()))
+
+	// don't intercept error from publish. Can't help you if the publisher is down as well.
+	return pq.pub.Publish(pq.topic, msg)
+}
+
+func (pq poisonQueue) Middleware(h message.HandlerFunc) message.HandlerFunc {
+	return func(msg *message.Message) (events []*message.Message, err error) {
+		defer func() {
+			if err != nil {
+				if !pq.shouldGoToPoisonQueue(err) {
+					return
+				}
+
+				// handler didn't cope with the message; publish it on the poison topic and carry on as usual
+				publishErr := pq.publishPoisonMessage(msg, err)
+				if publishErr != nil {
+					publishErr = errors.Wrap(publishErr, "cannot publish message to poison queue")
+					err = stdErrors.Join(err, publishErr)
+					return
+				}
+
+				err = nil
+				return
+			}
+		}()
+
+		// if h fails, the deferred function will salvage all that it can
+		return h(msg)
+	}
+}

message/router/middleware/poison_test.go

@@ -0,0 +1,241 @@
+package middleware_test
+
+import (
+	"context"
+	stdErrors "errors"
+	"testing"
+	"time"
+
+	"github.com/pkg/errors"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/ThreeDotsLabs/watermill"
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/ThreeDotsLabs/watermill/message/router/middleware"
+	"github.com/ThreeDotsLabs/watermill/message/subscriber"
+	"github.com/ThreeDotsLabs/watermill/pubsub/gochannel"
+)
+
+const topic = "testing_poison_queue_topic"
+
+// TestPoisonQueue_handler_ok simulates the situation when the message is processed correctly
+// We expect that all messages pass through the middleware unaffected and the poison queue catches no messages.
+func TestPoisonQueue_handler_ok(t *testing.T) {
+	poisonPublisher := mockPublisher{behaviour: BehaviourAlwaysOK}
+
+	poisonQueue, err := middleware.PoisonQueue(&poisonPublisher, topic)
+	require.NoError(t, err)
+
+	poisonQueueWithFilter, err := middleware.PoisonQueueWithFilter(&poisonPublisher, topic, func(err error) bool {
+		return true
+	})
+	require.NoError(t, err)
+
+	testCases := []struct {
+		Name       string
+		Middleware message.HandlerMiddleware
+	}{
+		{
+			Name:       "PoisonQueue",
+			Middleware: poisonQueue,
+		},
+		{
+			Name:       "PoisonQueueWithFilter",
+			Middleware: poisonQueueWithFilter,
+		},
+	}
+
+	for _, c := range testCases {
+		t.Run(c.Name, func(t *testing.T) {
+			produced, err := c.Middleware(handlerFuncAlwaysOK)(
+				message.NewMessage("uuid", nil),
+			)
+
+			assert.NoError(t, err)
+			assert.Equal(t, handlerFuncAlwaysOKMessages, produced)
+			assert.Empty(t, poisonPublisher.PopMessages())
+		})
+	}
+}
+
+func TestPoisonQueue_handler_failing(t *testing.T) {
+	poisonPublisher := mockPublisher{behaviour: BehaviourAlwaysOK}
+
+	poisonQueue, err := middleware.PoisonQueue(&poisonPublisher, topic)
+	require.NoError(t, err)
+
+	poisonQueueWithFilter, err := middleware.PoisonQueueWithFilter(&poisonPublisher, topic, func(err error) bool {
+		return true
+	})
+	require.NoError(t, err)
+
+	testCases := []struct {
+		Name       string
+		Middleware message.HandlerMiddleware
+	}{
+		{
+			Name:       "PoisonQueue",
+			Middleware: poisonQueue,
+		},
+		{
+			Name:       "PoisonQueueWithFilter",
+			Middleware: poisonQueueWithFilter,
+		},
+	}
+
+	for _, c := range testCases {
+		t.Run(c.Name, func(t *testing.T) {
+			msg := message.NewMessage("uuid", []byte("payload"))
+			produced, err := c.Middleware(handlerFuncAlwaysFailing)(
+				msg,
+			)
+
+			// the middleware itself should not fail; the publisher is working OK, so no error is passed down the chain
+			assert.NoError(t, err)
+
+			// but no messages should be passed
+			assert.Empty(t, produced)
+
+			// the original message should end up in the poison queue
+			poisonMsgs := poisonPublisher.PopMessages()
+			require.Len(t, poisonMsgs, 1)
+
+			assert.Equal(t, msg.Payload, poisonMsgs[0].Payload)
+
+			// there should be additional metadata telling why the message was poisoned
+			// it should be the error that the handler failed with
+			assert.Equal(t, errFailed.Error(), poisonMsgs[0].Metadata.Get(middleware.ReasonForPoisonedKey))
+		})
+	}
+}
+
+func TestPoisonQueue_context_values(t *testing.T) {
+	pubSub := gochannel.NewGoChannel(
+		gochannel.Config{Persistent: true},
+		watermill.NewStdLogger(true, true),
+	)
+
+	logger := watermill.NewStdLogger(true, true)
+
+	router, err := message.NewRouter(message.RouterConfig{}, logger)
+	require.NoError(t, err)
+
+	pq, err := middleware.PoisonQueue(pubSub, "poison_queue")
+	require.NoError(t, err)
+	router.AddMiddleware(pq)
+
+	router.AddConsumerHandler("handler_name", "test", pubSub, func(msg *message.Message) error {
+		return errors.New("error")
+	})
+
+	go func() {
+		require.NoError(t, router.Run(context.Background()))
+	}()
+	require.NoError(t, err)
+	defer func() { _ = router.Close() }()
+
+	select {
+	case <-router.Running():
+	// ok
+	case <-time.After(time.Second):
+		t.Fatal("waiting for router timeout")
+	}
+
+	err = pubSub.Publish("test", message.NewMessage("1", nil))
+	require.NoError(t, err)
+
+	msgs, err := pubSub.Subscribe(context.Background(), "poison_queue")
+	require.NoError(t, err)
+
+	messages, all := subscriber.BulkRead(msgs, 1, time.Second)
+	require.True(t, all, "no messages received")
+
+	assert.Equal(t, "handler_name", messages[0].Metadata[middleware.PoisonedHandlerKey])
+	assert.Equal(t, "gochannel.GoChannel", messages[0].Metadata[middleware.PoisonedSubscriberKey])
+	assert.Equal(t, "test", messages[0].Metadata[middleware.PoisonedTopicKey])
+	assert.Equal(t, "error", messages[0].Metadata[middleware.ReasonForPoisonedKey])
+}
+
+func TestPoisonQueue_handler_failing_publisher_failing(t *testing.T) {
+	poisonPublisher := mockPublisher{behaviour: BehaviourAlwaysFail}
+
+	poisonQueue, err := middleware.PoisonQueue(&poisonPublisher, topic)
+	require.NoError(t, err)
+
+	poisonQueueWithFilter, err := middleware.PoisonQueueWithFilter(&poisonPublisher, topic, func(err error) bool {
+		return true
+	})
+	require.NoError(t, err)
+
+	testCases := []struct {
+		Name       string
+		Middleware message.HandlerMiddleware
+	}{
+		{
+			Name:       "PoisonQueue",
+			Middleware: poisonQueue,
+		},
+		{
+			Name:       "PoisonQueueWithFilter",
+			Middleware: poisonQueueWithFilter,
+		},
+	}
+
+	for _, c := range testCases {
+		t.Run(c.Name, func(t *testing.T) {
+			msg := message.NewMessage("uuid", nil)
+			produced, err := poisonQueue(handlerFuncAlwaysFailing)(
+				msg,
+			)
+
+			// publisher failed, can't hide the error anymore
+			// Instead of checking the specific error, we check if the error.Is() is the same as the one we expect
+			assert.ErrorIs(t, err, errFailed)
+
+			// can't really expect any produced messages
+			assert.Empty(t, produced)
+
+			// nor poison messages
+			assert.Empty(t, poisonPublisher.PopMessages())
+		})
+	}
+}
+
+func TestPoisonQueueWithFilter_poison_queue(t *testing.T) {
+	poisonPublisher := mockPublisher{behaviour: BehaviourAlwaysOK}
+
+	poisonQueueErr := errors.New("poison queue err")
+	msg := message.NewMessage("uuid", []byte("payload"))
+
+	poisonQueue, err := middleware.PoisonQueueWithFilter(&poisonPublisher, topic, func(err error) bool {
+		return stdErrors.Is(err, poisonQueueErr)
+	})
+	require.NoError(t, err)
+
+	_, err = poisonQueue(func(msg *message.Message) (messages []*message.Message, e error) {
+		return nil, poisonQueueErr
+	})(msg)
+
+	assert.NoError(t, err)
+	require.Len(t, poisonPublisher.PopMessages(), 1)
+}
+
+func TestPoisonQueueWithFilter_non_poison_queue(t *testing.T) {
+	poisonPublisher := mockPublisher{behaviour: BehaviourAlwaysOK}
+
+	nonPoisonQueueErr := errors.New("non poison queue err")
+	msg := message.NewMessage("uuid", []byte("payload"))
+
+	poisonQueue, err := middleware.PoisonQueueWithFilter(&poisonPublisher, topic, func(err error) bool {
+		return !stdErrors.Is(err, nonPoisonQueueErr)
+	})
+	require.NoError(t, err)
+
+	_, err = poisonQueue(func(msg *message.Message) (messages []*message.Message, e error) {
+		return nil, nonPoisonQueueErr
+	})(msg)
+
+	assert.Error(t, err)
+	require.Len(t, poisonPublisher.PopMessages(), 0)
+}

message/router/middleware/randomfail.go

@@ -0,0 +1,39 @@
+package middleware
+
+import (
+	"math/rand"
+
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/pkg/errors"
+)
+
+func shouldFail(probability float32) bool {
+	r := rand.Float32()
+	return r <= probability
+}
+
+// RandomFail makes the handler fail with an error based on random chance. Error probability should be in the range (0,1).
+func RandomFail(errorProbability float32) message.HandlerMiddleware {
+	return func(h message.HandlerFunc) message.HandlerFunc {
+		return func(message *message.Message) ([]*message.Message, error) {
+			if shouldFail(errorProbability) {
+				return nil, errors.New("random fail occurred")
+			}
+
+			return h(message)
+		}
+	}
+}
+
+// RandomPanic makes the handler panic based on random chance. Panic probability should be in the range (0,1).
+func RandomPanic(panicProbability float32) message.HandlerMiddleware {
+	return func(h message.HandlerFunc) message.HandlerFunc {
+		return func(message *message.Message) ([]*message.Message, error) {
+			if shouldFail(panicProbability) {
+				panic("random panic occurred")
+			}
+
+			return h(message)
+		}
+	}
+}

message/router/middleware/randomfail_test.go

@@ -0,0 +1,28 @@
+package middleware_test
+
+import (
+	"testing"
+
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/ThreeDotsLabs/watermill/message/router/middleware"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRandomFail(t *testing.T) {
+	h := middleware.RandomFail(1)(func(msg *message.Message) (messages []*message.Message, e error) {
+		return nil, nil
+	})
+
+	_, err := h(message.NewMessage("1", nil))
+	assert.Error(t, err)
+}
+
+func TestRandomPanic(t *testing.T) {
+	h := middleware.RandomPanic(1)(func(msg *message.Message) (messages []*message.Message, e error) {
+		return nil, nil
+	})
+
+	assert.Panics(t, func() {
+		_, _ = h(message.NewMessage("1", nil))
+	})
+}

message/router/middleware/recoverer.go

@@ -0,0 +1,37 @@
+package middleware
+
+import (
+	"fmt"
+	"runtime/debug"
+
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/pkg/errors"
+)
+
+// RecoveredPanicError holds the recovered panic's error along with the stacktrace.
+type RecoveredPanicError struct {
+	V          interface{}
+	Stacktrace string
+}
+
+func (p RecoveredPanicError) Error() string {
+	return fmt.Sprintf("panic occurred: %#v, stacktrace: \n%s", p.V, p.Stacktrace)
+}
+
+// Recoverer recovers from any panic in the handler and appends RecoveredPanicError with the stacktrace
+// to any error returned from the handler.
+func Recoverer(h message.HandlerFunc) message.HandlerFunc {
+	return func(event *message.Message) (events []*message.Message, err error) {
+		panicked := true
+
+		defer func() {
+			if r := recover(); r != nil || panicked {
+				err = errors.WithStack(RecoveredPanicError{V: r, Stacktrace: string(debug.Stack())})
+			}
+		}()
+
+		events, err = h(event)
+		panicked = false
+		return events, err
+	}
+}

message/router/middleware/recoverer_test.go

@@ -0,0 +1,39 @@
+package middleware_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/ThreeDotsLabs/watermill/message/router/middleware"
+	"github.com/stretchr/testify/require"
+)
+
+func TestRecoverer_Panic(t *testing.T) {
+	h := middleware.Recoverer(func(msg *message.Message) (messages []*message.Message, e error) {
+		panic("foo")
+	})
+
+	_, err := h(message.NewMessage("1", nil))
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "message/router/middleware/recoverer.go") // stacktrace part
+}
+
+func TestRecoverer_PanicNil(t *testing.T) {
+	h := middleware.Recoverer(func(msg *message.Message) (messages []*message.Message, e error) {
+		panic(nil)
+	})
+
+	_, err := h(message.NewMessage("1", nil))
+	require.Error(t, err)
+}
+
+func TestRecoverer_NoPanic(t *testing.T) {
+	h := middleware.Recoverer(func(msg *message.Message) (messages []*message.Message, e error) {
+		return nil, nil
+	})
+
+	_, err := h(message.NewMessage("1", nil))
+	require.NoError(t, err)
+}

message/router/middleware/retry.go

@@ -0,0 +1,172 @@
+package middleware
+
+import (
+	"time"
+
+	"github.com/cenkalti/backoff/v5"
+	"github.com/pkg/errors"
+
+	"github.com/ThreeDotsLabs/watermill"
+	"github.com/ThreeDotsLabs/watermill/message"
+)
+
+// RetryParams holds the parameters for a retry attempt
+type RetryParams struct {
+	// Err is the error that caused the retry attempt.
+	Err error
+	// RetryNum is the number of the retry attempt, starting from 1.
+	RetryNum int
+	// Delay is the delay for the next retry attempt.
+	Delay time.Duration
+}
+
+// RetriesExhaustedParams holds the parameters passed to OnRetriesExhausted.
+type RetriesExhaustedParams struct {
+	// Err is the last error returned by the handler before retries were exhausted.
+	Err error
+	// RetryNum is the total number of attempts performed (1 initial + MaxRetries retries).
+	// For MaxRetries=N this will equal N+1.
+	RetryNum int
+}
+
+// Retry provides a middleware that retries the handler if errors are returned.
+// The retry behaviour is configurable, with exponential backoff and maximum elapsed time.
+type Retry struct {
+	// MaxRetries is maximum number of times a retry will be attempted.
+	MaxRetries int
+
+	// InitialInterval is the first interval between retries. Subsequent intervals will be scaled by Multiplier.
+	InitialInterval time.Duration
+	// MaxInterval sets the limit for the exponential backoff of retries. The interval will not be increased beyond MaxInterval.
+	MaxInterval time.Duration
+	// Multiplier is the factor by which the waiting interval will be multiplied between retries.
+	Multiplier float64
+	// MaxElapsedTime sets the time limit of how long retries will be attempted. Disabled if 0.
+	MaxElapsedTime time.Duration
+	// RandomizationFactor randomizes the spread of the backoff times within the interval of:
+	// [currentInterval * (1 - RandomizationFactor), currentInterval * (1 + RandomizationFactor)].
+	RandomizationFactor float64
+
+	// OnRetryHook is an optional function that will be executed on each retry attempt.
+	// The number of the current retry is passed as retryNum,
+	OnRetryHook func(retryNum int, delay time.Duration)
+
+	// OnRetriesExhausted is an optional function that will be executed when all retries are exhausted.
+	// This is called when MaxRetries is reached and the handler still returns an error.
+	// It is NOT called when ShouldRetry returns false (that path returns a permanent error and exits earlier).
+	OnRetriesExhausted func(params RetriesExhaustedParams)
+
+	// ShouldRetry is an optional function that will be executed before each retry attempt.
+	// If ShouldRetry returns false, the retry will not be attempted.
+	ShouldRetry func(params RetryParams) bool
+
+	// ResetContextOnRetry indicates whether the message context should be reset on each retry attempt.
+	// See more: https://github.com/ThreeDotsLabs/watermill/issues/467
+	//
+	// This is not enabled by default to keep backward compatibility
+	// (in theory, someone may want to preserve context values between retries).
+	ResetContextOnRetry bool
+
+	Logger watermill.LoggerAdapter
+}
+
+// Middleware returns the Retry middleware.
+func (r Retry) Middleware(h message.HandlerFunc) message.HandlerFunc {
+	return func(msg *message.Message) ([]*message.Message, error) {
+		originalCtx := msg.Context()
+		retryNum := 0
+		// stoppedByPermanent is set when ShouldRetry returns false, so we know
+		// to skip OnRetriesExhausted — retries weren't actually exhausted, they
+		// were short-circuited. backoff/v5 strips the *PermanentError wrapper
+		// inside Retry, so we can't detect this from the returned error alone.
+		stoppedByPermanent := false
+
+		expBackoff := backoff.NewExponentialBackOff()
+		expBackoff.InitialInterval = r.InitialInterval
+		expBackoff.MaxInterval = r.MaxInterval
+		expBackoff.Multiplier = r.Multiplier
+		expBackoff.RandomizationFactor = r.RandomizationFactor
+
+		// MaxRetries + 1 because the first attempt is not a retry
+		retryBackoff := backoff.WithMaxTries(uint(r.MaxRetries + 1))
+
+		maxElapsedBackoff := backoff.WithMaxElapsedTime(r.MaxElapsedTime)
+
+		// notification is called on a failed retry attempt.
+		notification := func(err error, delay time.Duration) {
+			if r.Logger != nil {
+				r.Logger.Error("Error occurred, retrying", err, watermill.LogFields{
+					"retry_no":    retryNum,
+					"max_retries": r.MaxRetries,
+					"wait_time":   delay,
+				})
+			}
+		}
+
+		// operation: the function that will be retried.
+		operation := func() ([]*message.Message, error) {
+			select {
+			case <-originalCtx.Done():
+				return nil, originalCtx.Err()
+			default:
+				if r.ResetContextOnRetry {
+					// message is passed as a pointer, so its context can be canceled
+					// by the previous attempts -> it will break retries, because any
+					// underlying logic that relies on the context will fail.
+					// see more: https://github.com/ThreeDotsLabs/watermill/issues/467
+					//
+					// to avoid this, we need to reset the original context on each attempt
+					// we may lose context value that was set by the previous attempt
+					msg.SetContext(originalCtx)
+				}
+
+				producedMessages, err := h(msg)
+				if err == nil {
+					return producedMessages, nil
+				}
+
+				if r.ShouldRetry != nil && !r.ShouldRetry(RetryParams{
+					RetryNum: retryNum,
+					Err:      err,
+					Delay:    expBackoff.NextBackOff(),
+				}) {
+					// backoff.Permanent will stop the retry attempts
+					stoppedByPermanent = true
+					return producedMessages, backoff.Permanent(err)
+				}
+
+				if r.OnRetryHook != nil && retryNum > 0 {
+					// call RetryHook function on each retry attempt.
+					r.OnRetryHook(retryNum, expBackoff.NextBackOff())
+				}
+				retryNum++
+				return producedMessages, err
+			}
+		}
+
+		producedMessages, retryErr := backoff.Retry(
+			originalCtx,
+			operation,
+			backoff.WithBackOff(expBackoff),
+			retryBackoff,
+			maxElapsedBackoff,
+			backoff.WithNotify(notification),
+		)
+		var backoffPermanentError *backoff.PermanentError
+		if errors.As(retryErr, &backoffPermanentError) {
+			// just in case, we don't want to expose backoff.PermanentError to the outside world
+			return producedMessages, backoffPermanentError.Unwrap()
+		}
+		if retryErr != nil {
+			if r.OnRetriesExhausted != nil && !stoppedByPermanent {
+				r.OnRetriesExhausted(RetriesExhaustedParams{
+					Err:      retryErr,
+					RetryNum: retryNum,
+				})
+			}
+			return producedMessages, retryErr
+		}
+
+		return producedMessages, nil
+	}
+}

message/router/middleware/retry_test.go

@@ -0,0 +1,424 @@
+package middleware_test
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/cenkalti/backoff/v5"
+	"github.com/pkg/errors"
+	"github.com/stretchr/testify/assert"
+
+	"github.com/ThreeDotsLabs/watermill"
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/ThreeDotsLabs/watermill/message/router/middleware"
+)
+
+func TestRetry_retry(t *testing.T) {
+	retry := middleware.Retry{
+		MaxRetries: 1,
+	}
+
+	runCount := 0
+	producedMessages := message.Messages{message.NewMessage("2", nil)}
+
+	h := retry.Middleware(func(msg *message.Message) (messages []*message.Message, e error) {
+		runCount++
+		if runCount == 0 {
+			return nil, errors.New("foo")
+		}
+
+		return producedMessages, nil
+	})
+
+	handlerMessages, handlerErr := h(message.NewMessage("1", nil))
+
+	assert.Equal(t, 1, runCount)
+	assert.EqualValues(t, producedMessages, handlerMessages)
+	assert.NoError(t, handlerErr)
+}
+
+func TestRetry_max_retries(t *testing.T) {
+	retry := middleware.Retry{
+		MaxRetries: 1,
+		Logger:     watermill.NewStdLogger(true, true),
+	}
+
+	runCount := 0
+
+	h := retry.Middleware(func(msg *message.Message) (messages []*message.Message, e error) {
+		runCount++
+		return nil, errors.New("foo")
+	})
+
+	_, err := h(message.NewMessage("1", nil))
+
+	assert.Equal(t, 2, runCount)
+	assert.EqualError(t, err, "foo")
+}
+
+func TestRetry_retry_hook(t *testing.T) {
+	var retriesFromHook []int
+
+	retry := middleware.Retry{
+		MaxRetries: 2,
+		OnRetryHook: func(retryNum int, delay time.Duration) {
+			retriesFromHook = append(retriesFromHook, retryNum)
+		},
+	}
+
+	h := retry.Middleware(func(msg *message.Message) (messages []*message.Message, e error) {
+		return nil, errors.New("foo")
+	})
+	_, _ = h(message.NewMessage("1", nil))
+
+	assert.EqualValues(t, []int{1, 2}, retriesFromHook)
+}
+
+func TestRetry_retries_exhausted_hook(t *testing.T) {
+	var exhaustedParams middleware.RetriesExhaustedParams
+
+	handlerErr := errors.New("foo")
+
+	retry := middleware.Retry{
+		MaxRetries: 2,
+		OnRetriesExhausted: func(params middleware.RetriesExhaustedParams) {
+			exhaustedParams = params
+		},
+	}
+
+	h := retry.Middleware(func(msg *message.Message) (messages []*message.Message, e error) {
+		return nil, handlerErr
+	})
+	_, _ = h(message.NewMessage("1", nil))
+
+	assert.ErrorIs(t, exhaustedParams.Err, handlerErr)
+	// With MaxRetries=2, there are 3 total attempts (1 initial + 2 retries)
+	// retryNum is incremented after each attempt, so it's 3 when exhausted
+	assert.Equal(t, 3, exhaustedParams.RetryNum)
+}
+
+func TestRetry_retries_exhausted_hook_not_called_on_success(t *testing.T) {
+	var hookCalled bool
+
+	retry := middleware.Retry{
+		MaxRetries: 2,
+		OnRetriesExhausted: func(params middleware.RetriesExhaustedParams) {
+			hookCalled = true
+		},
+	}
+
+	h := retry.Middleware(func(msg *message.Message) (messages []*message.Message, e error) {
+		return nil, nil
+	})
+	_, _ = h(message.NewMessage("1", nil))
+
+	assert.False(t, hookCalled, "OnRetriesExhausted should not be called on success")
+}
+
+func TestRetry_retries_exhausted_hook_not_called_when_should_retry_returns_false(t *testing.T) {
+	var hookCalled bool
+
+	retry := middleware.Retry{
+		MaxRetries: 5,
+		ShouldRetry: func(params middleware.RetryParams) bool {
+			return false
+		},
+		OnRetriesExhausted: func(params middleware.RetriesExhaustedParams) {
+			hookCalled = true
+		},
+	}
+
+	h := retry.Middleware(func(msg *message.Message) (messages []*message.Message, e error) {
+		return nil, errors.New("foo")
+	})
+	_, _ = h(message.NewMessage("1", nil))
+
+	assert.False(t, hookCalled, "OnRetriesExhausted should not be called when ShouldRetry returns false")
+}
+
+func TestRetry_logger(t *testing.T) {
+	logger := watermill.NewCaptureLogger()
+
+	retry := middleware.Retry{
+		MaxRetries: 2,
+		Logger:     logger,
+	}
+
+	handlerErr := errors.New("foo")
+
+	h := retry.Middleware(func(msg *message.Message) (messages []*message.Message, e error) {
+		return nil, handlerErr
+	})
+	_, _ = h(message.NewMessage("1", nil))
+
+	assert.True(t, logger.HasError(handlerErr))
+}
+
+func TestRetry_ctx_cancel(t *testing.T) {
+	retry := middleware.Retry{
+		InitialInterval: time.Minute,
+	}
+
+	producedMessages := message.Messages{message.NewMessage("2", nil)}
+
+	h := retry.Middleware(func(msg *message.Message) (messages []*message.Message, e error) {
+		return producedMessages, errors.New("err")
+	})
+
+	msg := message.NewMessage("1", nil)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	msg.SetContext(ctx)
+
+	done := make(chan struct{})
+
+	type handlerResult struct {
+		Messages message.Messages
+		Err      error
+	}
+	handlerResultCh := make(chan handlerResult, 1)
+
+	go func() {
+		messages, err := h(msg)
+		handlerResultCh <- handlerResult{messages, err}
+		close(done)
+	}()
+
+	select {
+	case <-done:
+		t.Fatal("handler should be still during retrying")
+	default:
+		// ok
+	}
+
+	cancel()
+
+	select {
+	case <-done:
+		// ok
+	case <-time.After(time.Second):
+		t.Fatal("ctx cancelled, retrying should be done")
+	}
+
+	handlerResultReceived := <-handlerResultCh
+
+	assert.Error(t, handlerResultReceived.Err)
+	// produced messages are nil since ctx is canceling the operation in the middle
+	assert.Nil(t, handlerResultReceived.Messages)
+}
+
+func TestRetry_max_elapsed(t *testing.T) {
+	maxRetries := 10
+	sleepInHandler := time.Millisecond * 20
+
+	retry := middleware.Retry{
+		MaxElapsedTime: time.Millisecond * 10,
+		MaxRetries:     maxRetries,
+	}
+
+	runTimeWithoutMaxElapsedTime := sleepInHandler * time.Duration(maxRetries)
+
+	h := retry.Middleware(func(msg *message.Message) (messages []*message.Message, e error) {
+		time.Sleep(sleepInHandler)
+		return nil, errors.New("foo")
+	})
+
+	startTime := time.Now()
+	_, _ = h(message.NewMessage("1", nil))
+	timeElapsed := time.Since(startTime)
+
+	assert.True(
+		t,
+		timeElapsed < runTimeWithoutMaxElapsedTime,
+		"handler should run less than %s, time elapsed: %s",
+		runTimeWithoutMaxElapsedTime,
+		timeElapsed,
+	)
+}
+
+func TestRetry_max_interval(t *testing.T) {
+	t.Parallel()
+
+	maxRetries := 10
+	backoffTimes := make([]time.Duration, maxRetries)
+	maxInterval := time.Millisecond * 30
+
+	retry := middleware.Retry{
+		MaxRetries:          maxRetries,
+		InitialInterval:     time.Millisecond * 10,
+		MaxInterval:         maxInterval,
+		Multiplier:          2.0,
+		RandomizationFactor: 0,
+		OnRetryHook: func(retryNum int, delay time.Duration) {
+			backoffTimes[retryNum-1] = delay
+		},
+	}
+
+	h := retry.Middleware(func(msg *message.Message) (messages []*message.Message, e error) {
+		return nil, errors.New("bar")
+	})
+	_, _ = h(message.NewMessage("2", nil))
+
+	for i, delay := range backoffTimes {
+		assert.True(t, delay <= maxInterval, "wait interval %d (%s) exceeds maxInterval (%s)", i, delay, maxInterval)
+	}
+}
+
+func TestRetry_first_run_no_delay(t *testing.T) {
+	t.Parallel()
+
+	initialInterval := time.Millisecond * 100
+	retry := middleware.Retry{
+		MaxElapsedTime:  initialInterval * 2,
+		MaxRetries:      10,
+		InitialInterval: initialInterval,
+	}
+
+	h := retry.Middleware(func(msg *message.Message) (messages []*message.Message, e error) {
+		return nil, nil
+	})
+
+	start := time.Now()
+	_, _ = h(message.NewMessage("1", nil))
+	elapsed := time.Since(start)
+
+	assert.True(t, elapsed < initialInterval, "first retry should not wait, elapsed: %s", elapsed)
+}
+
+func TestRetry_should_retry(t *testing.T) {
+	errToSkip := errors.New("this should be skipped")
+
+	retry := middleware.Retry{
+		MaxRetries: 5,
+		ShouldRetry: func(params middleware.RetryParams) bool {
+			return !errors.Is(params.Err, errToSkip)
+		},
+	}
+
+	runCount := 0
+
+	h := retry.Middleware(func(msg *message.Message) (messages []*message.Message, e error) {
+		runCount++
+		return nil, errToSkip
+	})
+
+	handlerMessages, handlerErr := h(message.NewMessage("1", nil))
+
+	assert.Equal(t, 1, runCount)
+	assert.Nil(t, handlerMessages)
+	assert.ErrorIs(t, handlerErr, errToSkip)
+
+	// to not create any dependency on backoff package
+	var backoffPermanentError *backoff.PermanentError
+	assert.False(t, errors.As(handlerErr, &backoffPermanentError))
+}
+
+// Test that the ShouldRetry function is called on each retry attempt.
+// Under the hood, the second attempt goes over a bit different code path,
+// so we want to make sure that it works consistently.
+func TestRetry_should_retry_second_attempt(t *testing.T) {
+	errToSkip := errors.New("this should be skipped")
+
+	retry := middleware.Retry{
+		MaxRetries: 5,
+		ShouldRetry: func(params middleware.RetryParams) bool {
+			return !errors.Is(params.Err, errToSkip)
+		},
+	}
+
+	runCount := 0
+
+	h := retry.Middleware(func(msg *message.Message) (messages []*message.Message, e error) {
+		runCount++
+
+		if runCount == 1 {
+			return nil, errors.New("some other error")
+		} else {
+			return nil, errToSkip
+		}
+	})
+
+	handlerMessages, handlerErr := h(message.NewMessage("1", nil))
+
+	assert.Equal(t, 2, runCount)
+	assert.Nil(t, handlerMessages)
+	assert.ErrorIs(t, handlerErr, errToSkip)
+
+	// to not create any dependency on backoff package
+	var backoffPermanentError *backoff.PermanentError
+	assert.False(t, errors.As(handlerErr, &backoffPermanentError))
+}
+
+// Test that if backoff.Permanent error stops the retries.
+// For sake of potential future regressions (see Hyrum's Law).
+func TestRetry_backoff_backoff_permanent(t *testing.T) {
+	errToSkip := errors.New("this should be skipped")
+
+	retry := middleware.Retry{
+		MaxRetries: 5,
+	}
+
+	runCount := 0
+
+	h := retry.Middleware(func(msg *message.Message) (messages []*message.Message, e error) {
+		runCount++
+		return nil, backoff.Permanent(errToSkip)
+	})
+
+	handlerMessages, handlerErr := h(message.NewMessage("1", nil))
+
+	assert.Equal(t, 1, runCount)
+	assert.Nil(t, handlerMessages)
+	assert.ErrorIs(t, handlerErr, errToSkip)
+
+	// to not create any dependency on backoff package
+	var backoffPermanentError *backoff.PermanentError
+	assert.False(t, errors.As(handlerErr, &backoffPermanentError))
+}
+
+// TestRetry_uncancel_context checks scenario when context is canceled,
+// and we want to retry. More context: https://github.com/ThreeDotsLabs/watermill/issues/467
+//
+// Message is passed as a pointer, so underlying middlewares or handlers can cancel its context.
+// In this scenario, retrying is pointless because context is already canceled, so any operation
+// that relies on context will fail immediately.
+func TestRetry_uncancel_context(t *testing.T) {
+	retry := middleware.Retry{
+		MaxRetries:          5,
+		ResetContextOnRetry: true,
+	}
+
+	num := 0
+
+	var ctxCancelMiddleware = func(h message.HandlerFunc) message.HandlerFunc {
+		return func(msg *message.Message) ([]*message.Message, error) {
+			num++
+
+			ctx, cancel := context.WithCancel(msg.Context())
+			defer func() {
+				cancel()
+			}()
+
+			if num == 1 {
+				t.Log("Run 1: canceling context")
+				cancel()
+			} else {
+				t.Logf("Run %d: context is not canceled", num)
+			}
+
+			msg.SetContext(ctx)
+			return h(msg)
+		}
+	}
+
+	h := func(msg *message.Message) (messages []*message.Message, e error) {
+		return nil, msg.Context().Err()
+	}
+
+	h = ctxCancelMiddleware(h)
+	h = retry.Middleware(h)
+
+	_, handlerErr := h(message.NewMessage("1", nil))
+	assert.NoError(t, handlerErr)
+}

message/router/middleware/throttle.go

@@ -0,0 +1,31 @@
+package middleware
+
+import (
+	"time"
+
+	"github.com/ThreeDotsLabs/watermill/message"
+)
+
+// Throttle provides a middleware that limits the amount of messages processed per unit of time.
+// This may be done e.g. to prevent excessive load caused by running a handler on a long queue of unprocessed messages.
+type Throttle struct {
+	ticker *time.Ticker
+}
+
+// NewThrottle creates a new Throttle middleware.
+// Example duration and count: NewThrottle(10, time.Second) for 10 messages per second
+func NewThrottle(count int64, duration time.Duration) *Throttle {
+	return &Throttle{
+		ticker: time.NewTicker(duration / time.Duration(count)),
+	}
+}
+
+// Middleware returns the Throttle middleware.
+func (t Throttle) Middleware(h message.HandlerFunc) message.HandlerFunc {
+	return func(message *message.Message) ([]*message.Message, error) {
+		// throttle is shared by multiple handlers, which will wait for their "tick"
+		<-t.ticker.C
+
+		return h(message)
+	}
+}

message/router/middleware/throttle_test.go

@@ -0,0 +1,79 @@
+package middleware_test
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/pkg/errors"
+
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/ThreeDotsLabs/watermill/message/router/middleware"
+	"github.com/stretchr/testify/assert"
+)
+
+const (
+	perSecond          = 10
+	testTimeout        = time.Second
+	concurrentHandlers = 10
+)
+
+func TestThrottle_Middleware(t *testing.T) {
+	throttle := middleware.NewThrottle(perSecond, testTimeout)
+
+	ctx, cancel := context.WithTimeout(context.Background(), testTimeout)
+	defer cancel()
+
+	producedMessagesChannel := make(chan struct{})
+
+	producedMessagesCounter := 0
+
+	for i := 0; i < concurrentHandlers; i++ {
+		go func() {
+		HandlerLoop:
+			for {
+				producedMessages := []*message.Message{message.NewMessage("produced", nil)}
+				producedErr := errors.New("produced err")
+
+				produced, err := throttle.Middleware(func(msg *message.Message) ([]*message.Message, error) {
+					return producedMessages, producedErr
+				})(
+					message.NewMessage("uuid", nil),
+				)
+
+				assert.Equal(t, producedMessages, produced)
+				assert.Equal(t, producedErr, err)
+
+				go func() {
+					// non blocking counting
+					producedMessagesChannel <- struct{}{}
+				}()
+
+				select {
+				case <-ctx.Done():
+					break HandlerLoop
+				default:
+				}
+			}
+		}()
+	}
+
+CounterLoop:
+	for {
+		select {
+		case <-ctx.Done():
+			break CounterLoop
+		case <-producedMessagesChannel:
+			producedMessagesCounter++
+		}
+	}
+
+	t.Logf("produced %d messages in %d seconds, at rate of total %d messages per second",
+		producedMessagesCounter,
+		int(testTimeout.Seconds()),
+		perSecond,
+	)
+
+	assert.True(t, producedMessagesCounter <= int(perSecond*testTimeout.Seconds()))
+	assert.True(t, producedMessagesCounter > 0)
+}

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)
+		}
+	}
+}

message/router/middleware/timeout_test.go

@@ -0,0 +1,30 @@
+package middleware_test
+
+import (
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/ThreeDotsLabs/watermill/message/router/middleware"
+)
+
+func TestTimeout(t *testing.T) {
+	timeout := middleware.Timeout(time.Millisecond * 10)
+
+	h := timeout(func(msg *message.Message) ([]*message.Message, error) {
+		delay := time.After(time.Millisecond * 100)
+
+		select {
+		case <-msg.Context().Done():
+			return nil, nil
+		case <-delay:
+			return nil, errors.New("timeout did not occur")
+		}
+	})
+
+	_, err := h(message.NewMessage("any-uuid", nil))
+	require.NoError(t, err)
+}

message/router/plugin/signals.go

@@ -0,0 +1,27 @@
+package plugin
+
+import (
+	"fmt"
+	"os"
+	"os/signal"
+	"syscall"
+
+	"github.com/ThreeDotsLabs/watermill/message"
+)
+
+// SignalsHandler is a plugin that kills the router after SIGINT or SIGTERM is sent to the process.
+func SignalsHandler(r *message.Router) error {
+	sigs := make(chan os.Signal, 1)
+	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
+
+	go func() {
+		sig := <-sigs
+		r.Logger().Info(fmt.Sprintf("Received %s signal, closing\n", sig), nil)
+
+		err := r.Close()
+		if err != nil {
+			r.Logger().Error("Router close failed", err, nil)
+		}
+	}()
+	return nil
+}