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

message/decorator.go

@@ -0,0 +1,81 @@
+package message
+
+import (
+	"context"
+	"sync"
+)
+
+// MessageTransformSubscriberDecorator creates a subscriber decorator that calls transform
+// on each message that passes through the subscriber.
+func MessageTransformSubscriberDecorator(transform func(*Message)) SubscriberDecorator {
+	if transform == nil {
+		panic("transform function is nil")
+	}
+	return func(sub Subscriber) (Subscriber, error) {
+		return &messageTransformSubscriberDecorator{
+			sub:       sub,
+			transform: transform,
+		}, nil
+	}
+}
+
+// MessageTransformPublisherDecorator creates a publisher decorator that calls transform
+// on each message that passes through the publisher.
+func MessageTransformPublisherDecorator(transform func(*Message)) PublisherDecorator {
+	if transform == nil {
+		panic("transform function is nil")
+	}
+	return func(pub Publisher) (Publisher, error) {
+		return &messageTransformPublisherDecorator{
+			Publisher: pub,
+			transform: transform,
+		}, nil
+	}
+}
+
+type messageTransformSubscriberDecorator struct {
+	sub Subscriber
+
+	transform   func(*Message)
+	subscribeWg sync.WaitGroup
+}
+
+func (t *messageTransformSubscriberDecorator) Subscribe(ctx context.Context, topic string) (<-chan *Message, error) {
+	in, err := t.sub.Subscribe(ctx, topic)
+	if err != nil {
+		return nil, err
+	}
+
+	out := make(chan *Message)
+	t.subscribeWg.Add(1)
+	go func() {
+		for msg := range in {
+			t.transform(msg)
+			out <- msg
+		}
+		close(out)
+		t.subscribeWg.Done()
+	}()
+
+	return out, nil
+}
+
+func (t *messageTransformSubscriberDecorator) Close() error {
+	err := t.sub.Close()
+
+	t.subscribeWg.Wait()
+	return err
+}
+
+type messageTransformPublisherDecorator struct {
+	Publisher
+	transform func(*Message)
+}
+
+// Publish applies the transform to each message and returns the underlying Publisher's result.
+func (d messageTransformPublisherDecorator) Publish(topic string, messages ...*Message) error {
+	for i := range messages {
+		d.transform(messages[i])
+	}
+	return d.Publisher.Publish(topic, messages...)
+}

message/decorator_bench_test.go

@@ -0,0 +1,89 @@
+package message_test
+
+import (
+	"context"
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/ThreeDotsLabs/watermill/message"
+)
+
+type benchSubscriber struct {
+	msgCh   chan *message.Message
+	closeCh chan struct{}
+	running sync.Mutex
+}
+
+func (b *benchSubscriber) Subscribe(ctx context.Context, topic string) (<-chan *message.Message, error) {
+	return b.msgCh, nil
+}
+
+func (b *benchSubscriber) Close() error {
+	close(b.closeCh)
+	b.running.Lock()
+	close(b.msgCh)
+	b.running.Unlock()
+	return nil
+}
+
+func newBenchSubscriber() *benchSubscriber {
+	sub := &benchSubscriber{
+		msgCh:   make(chan *message.Message, 1),
+		closeCh: make(chan struct{}),
+	}
+
+	// continually produce messages until Close()
+	go func() {
+		sub.running.Lock()
+		msg := message.NewMessage("", []byte{})
+		for {
+			select {
+			case <-sub.closeCh:
+				sub.running.Unlock()
+				return
+			case sub.msgCh <- msg:
+				// the buffer limit is 1, so this will block until someone consumes
+			}
+		}
+	}()
+
+	return sub
+}
+
+func BenchmarkMessageTransformSubscriberDecorator(b *testing.B) {
+	b.ReportAllocs()
+	b.Run("no_decorator", benchmarkNoDecorator)
+	b.Run("message_transform_decorator", benchmarkMessageTransformSubscriberDecorator)
+}
+
+func benchmarkNoDecorator(b *testing.B) {
+	sub := newBenchSubscriber()
+
+	in, err := sub.Subscribe(context.Background(), "")
+	require.NoError(b, err)
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		// consume one message
+		<-in
+	}
+}
+
+func benchmarkMessageTransformSubscriberDecorator(b *testing.B) {
+	sub := newBenchSubscriber()
+
+	noopDecorator := message.MessageTransformSubscriberDecorator(func(*message.Message) {})
+	decoratedSub, err := noopDecorator(sub)
+	require.NoError(b, err)
+
+	in, err := decoratedSub.Subscribe(context.Background(), "")
+	require.NoError(b, err)
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		// consume one message
+		<-in
+	}
+}

message/decorator_test.go

@@ -0,0 +1,248 @@
+package message_test
+
+import (
+	"context"
+	"strconv"
+	"testing"
+	"time"
+
+	"github.com/ThreeDotsLabs/watermill/pubsub/tests"
+
+	"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/subscriber"
+	"github.com/ThreeDotsLabs/watermill/pubsub/gochannel"
+)
+
+var noop = func(*message.Message) {}
+var closingErr = errors.New("mock error on close")
+
+type mockSubscriber struct {
+	ch chan *message.Message
+}
+
+func (m mockSubscriber) Subscribe(context.Context, string) (<-chan *message.Message, error) {
+	return m.ch, nil
+}
+
+func (m mockSubscriber) Close() error {
+	close(m.ch)
+	return nil
+}
+
+func TestMessageTransformSubscriberDecorator_transparent(t *testing.T) {
+	sub := mockSubscriber{make(chan *message.Message)}
+	decorated, err := message.MessageTransformSubscriberDecorator(noop)(sub)
+	require.NoError(t, err)
+
+	messages, err := decorated.Subscribe(context.Background(), "topic")
+	require.NoError(t, err)
+
+	richMessage := message.NewMessage("uuid", []byte("serious payloads"))
+	richMessage.Metadata.Set("k1", "v1")
+	richMessage.Metadata.Set("k2", "v2")
+
+	go func() {
+		sub.ch <- richMessage
+	}()
+
+	received, all := subscriber.BulkRead(messages, 1, time.Second)
+	require.True(t, all)
+
+	assert.True(t, received[0].Equals(richMessage), "expected the message to pass unchanged through decorator")
+}
+
+type closingSubscriber struct {
+	closed bool
+}
+
+func (closingSubscriber) Subscribe(context.Context, string) (<-chan *message.Message, error) {
+	return nil, nil
+}
+
+func (c *closingSubscriber) Close() error {
+	c.closed = true
+	return closingErr
+}
+
+func TestMessageTransformSubscriberDecorator_Close(t *testing.T) {
+	cs := &closingSubscriber{}
+
+	decoratedSub, err := message.MessageTransformSubscriberDecorator(noop)(cs)
+	require.NoError(t, err)
+
+	// given
+	require.False(t, cs.closed)
+
+	// when
+	decoratedCloseErr := decoratedSub.Close()
+
+	// then
+	assert.True(
+		t,
+		cs.closed,
+		"expected the Close() call to propagate to decorated subscriber",
+	)
+	assert.Equal(
+		t,
+		closingErr,
+		decoratedCloseErr,
+		"expected the decorator to propagate the closing error from underlying subscriber",
+	)
+}
+
+func TestMessageTransformSubscriberDecorator_Subscribe(t *testing.T) {
+	numMessages := 1000
+	pubSub := gochannel.NewGoChannel(gochannel.Config{}, watermill.NewStdLogger(true, true))
+
+	onMessage := func(msg *message.Message) {
+		msg.Metadata.Set("key", "value")
+	}
+	decorator := message.MessageTransformSubscriberDecorator(onMessage)
+
+	decoratedSub, err := decorator(pubSub)
+	require.NoError(t, err)
+
+	messages, err := decoratedSub.Subscribe(context.Background(), "topic")
+	require.NoError(t, err)
+
+	sent := message.Messages{}
+
+	go func() {
+		for i := 0; i < numMessages; i++ {
+			msg := message.NewMessage(strconv.Itoa(i), []byte{})
+			sent = append(sent, msg)
+
+			err = pubSub.Publish("topic", msg)
+			require.NoError(t, err)
+		}
+	}()
+
+	received, all := subscriber.BulkRead(messages, numMessages, time.Second)
+	require.True(t, all)
+	tests.AssertAllMessagesReceived(t, sent, received)
+
+	for _, msg := range received {
+		assert.Equal(
+			t,
+			"value",
+			msg.Metadata.Get("key"),
+			"expected onMessage callback to have set metadata",
+		)
+	}
+}
+
+type mockPublisher struct {
+	published message.Messages
+}
+
+func (m *mockPublisher) Publish(topic string, messages ...*message.Message) error {
+	m.published = append(m.published, messages...)
+	return nil
+}
+
+func (m mockPublisher) Close() error {
+	return nil
+}
+
+func TestMessageTransformPublisherDecorator_transparent(t *testing.T) {
+	pub := &mockPublisher{message.Messages{}}
+	decorated, err := message.MessageTransformPublisherDecorator(noop)(pub)
+	require.NoError(t, err)
+
+	richMessage := message.NewMessage("uuid", []byte("serious payloads"))
+	richMessage.Metadata.Set("k1", "v1")
+	richMessage.Metadata.Set("k2", "v2")
+
+	require.NoError(t, decorated.Publish("topic", richMessage))
+
+	require.Len(t, pub.published, 1)
+	assert.True(t, pub.published[0].Equals(richMessage), "expected the message to pass unchanged through decorator")
+}
+
+type closingPublisher struct {
+	closed bool
+}
+
+func (c *closingPublisher) Publish(topic string, messages ...*message.Message) error {
+	return nil
+}
+
+func (c *closingPublisher) Close() error {
+	c.closed = true
+	return closingErr
+}
+
+func TestMessageTransformPublisherDecorator_Close(t *testing.T) {
+	cp := &closingPublisher{}
+
+	decoratedPub, err := message.MessageTransformPublisherDecorator(noop)(cp)
+	require.NoError(t, err)
+
+	// given
+	require.False(t, cp.closed)
+
+	// when
+	decoratedCloseErr := decoratedPub.Close()
+
+	// then
+	assert.True(
+		t,
+		cp.closed,
+		"expected the Close() call to propagate to decorated publisher",
+	)
+	assert.Equal(
+		t,
+		closingErr,
+		decoratedCloseErr,
+		"expected the decorator to propagate the closing error from underlying publisher",
+	)
+}
+
+func TestMessageTransformPublisherDecorator_Subscribe(t *testing.T) {
+	numMessages := 1000
+	pub := &mockPublisher{}
+
+	onMessage := func(msg *message.Message) {
+		msg.Metadata.Set("key", "value")
+	}
+	decorator := message.MessageTransformPublisherDecorator(onMessage)
+	decoratedPub, err := decorator(pub)
+	require.NoError(t, err)
+
+	for i := 0; i < numMessages; i++ {
+		msg := message.NewMessage(strconv.Itoa(i), []byte{})
+		require.NoError(t, decoratedPub.Publish("topic", msg))
+	}
+
+	for i, msg := range pub.published {
+		assert.Equal(t, strconv.Itoa(i), msg.UUID, "expected messages to arrive in unchanged order")
+		assert.Equal(
+			t,
+			"value",
+			msg.Metadata.Get("key"),
+			"expected onMessage callback to have set metadata",
+		)
+	}
+}
+
+func TestMessageTransformer_nil_panics(t *testing.T) {
+	require.Panics(
+		t,
+		func() {
+			_ = message.MessageTransformSubscriberDecorator(nil)
+		},
+		"expected to panic if transform is nil",
+	)
+	require.Panics(
+		t,
+		func() {
+			_ = message.MessageTransformPublisherDecorator(nil)
+		},
+		"expected to panic if transform is nil",
+	)
+}

message/message.go

@@ -0,0 +1,210 @@
+package message
+
+import (
+	"bytes"
+	"context"
+	"sync"
+)
+
+var closedchan = make(chan struct{})
+
+func init() {
+	close(closedchan)
+}
+
+// Payload is the Message's payload.
+type Payload []byte
+
+// Message is the basic transfer unit.
+// Messages are emitted by Publishers and received by Subscribers.
+//
+// A publisher can modify the message during publishing, e.g. can alter the metadata.
+// Avoid modifying the message in parallel with publishing, as it can lead to data races.
+// In general, a message should be passed to a single Publish and then considered immutable.
+// If needed, use the Copy method to create a new message.
+type Message struct {
+	// UUID is a unique identifier of the message.
+	//
+	// It is only used by Watermill for debugging.
+	// UUID can be empty.
+	UUID string
+
+	// Metadata contains the message metadata.
+	//
+	// Can be used to store data which doesn't require unmarshalling the entire payload.
+	// It is something similar to HTTP request's headers.
+	//
+	// Metadata is marshaled and will be saved to the PubSub.
+	Metadata Metadata
+
+	// Payload is the message's payload.
+	Payload Payload
+
+	// ack is closed when acknowledge is received.
+	ack chan struct{}
+	// noAck is closed when negative acknowledge is received.
+	noAck chan struct{}
+
+	ackMutex    sync.Mutex
+	ackSentType ackType
+
+	ctx context.Context
+}
+
+// NewMessage creates a new Message with given uuid and payload.
+func NewMessage(uuid string, payload Payload) *Message {
+	return &Message{
+		UUID:     uuid,
+		Metadata: make(map[string]string),
+		Payload:  payload,
+		ack:      make(chan struct{}),
+		noAck:    make(chan struct{}),
+	}
+}
+
+// NewMessageWithContext creates a new Message with given uuid, payload, and context.
+func NewMessageWithContext(ctx context.Context, uuid string, payload Payload) *Message {
+	msg := NewMessage(uuid, payload)
+	msg.SetContext(ctx)
+	return msg
+}
+
+type ackType int
+
+const (
+	noAckSent ackType = iota
+	ack
+	nack
+)
+
+// Equals compare, that two messages are equal. Acks/Nacks are not compared.
+func (m *Message) Equals(toCompare *Message) bool {
+	if m.UUID != toCompare.UUID {
+		return false
+	}
+	if len(m.Metadata) != len(toCompare.Metadata) {
+		return false
+	}
+	for key, value := range m.Metadata {
+		if value != toCompare.Metadata[key] {
+			return false
+		}
+	}
+	return bytes.Equal(m.Payload, toCompare.Payload)
+}
+
+// Ack sends message's acknowledgement.
+//
+// Ack is not blocking.
+// Ack is idempotent.
+// False is returned, if Nack is already sent.
+func (m *Message) Ack() bool {
+	m.ackMutex.Lock()
+	defer m.ackMutex.Unlock()
+
+	if m.ackSentType == nack {
+		return false
+	}
+	if m.ackSentType != noAckSent {
+		return true
+	}
+
+	m.ackSentType = ack
+	if m.ack == nil {
+		m.ack = closedchan
+	} else {
+		close(m.ack)
+	}
+
+	return true
+}
+
+// Nack sends message's negative acknowledgement.
+//
+// Nack is not blocking.
+// Nack is idempotent.
+// False is returned, if Ack is already sent.
+func (m *Message) Nack() bool {
+	m.ackMutex.Lock()
+	defer m.ackMutex.Unlock()
+
+	if m.ackSentType == ack {
+		return false
+	}
+	if m.ackSentType != noAckSent {
+		return true
+	}
+
+	m.ackSentType = nack
+
+	if m.noAck == nil {
+		m.noAck = closedchan
+	} else {
+		close(m.noAck)
+	}
+
+	return true
+}
+
+// Acked returns channel which is closed when acknowledgement is sent.
+//
+// Usage:
+//
+//	select {
+//	case <-message.Acked():
+//		// ack received
+//	case <-message.Nacked():
+//		// nack received
+//	}
+func (m *Message) Acked() <-chan struct{} {
+	return m.ack
+}
+
+// Nacked returns channel which is closed when negative acknowledgement is sent.
+//
+// Usage:
+//
+//	select {
+//	case <-message.Acked():
+//		// ack received
+//	case <-message.Nacked():
+//		// nack received
+//	}
+func (m *Message) Nacked() <-chan struct{} {
+	return m.noAck
+}
+
+// Context returns the message's context. To change the context, use
+// SetContext.
+//
+// The returned context is always non-nil; it defaults to the
+// background context.
+func (m *Message) Context() context.Context {
+	if m.ctx != nil {
+		return m.ctx
+	}
+	return context.Background()
+}
+
+// SetContext sets provided context to the message.
+func (m *Message) SetContext(ctx context.Context) {
+	m.ctx = ctx
+}
+
+// Copy copies all message without Acks/Nacks.
+// The context is not propagated to the copy.
+func (m *Message) Copy() *Message {
+	msg := NewMessage(m.UUID, m.Payload)
+	for k, v := range m.Metadata {
+		msg.Metadata.Set(k, v)
+	}
+	return msg
+}
+
+// CopyWithContext copies all message without Acks/Nacks.
+// The context is also propagated to the copy.
+func (m *Message) CopyWithContext() *Message {
+	msg := m.Copy()
+	msg.ctx = m.ctx
+	return msg
+}

message/message_test.go

@@ -0,0 +1,190 @@
+package message_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/ThreeDotsLabs/watermill/message"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestMessage_Equals(t *testing.T) {
+	withMetadata := func(msg *message.Message, metadata message.Metadata) *message.Message {
+		msg.Metadata = metadata
+		return msg
+	}
+
+	testCases := []struct {
+		Name   string
+		Msg1   *message.Message
+		Msg2   *message.Message
+		Equals bool
+	}{
+		{
+			Name:   "equal",
+			Msg1:   message.NewMessage("1", []byte("foo")),
+			Msg2:   message.NewMessage("1", []byte("foo")),
+			Equals: true,
+		},
+		{
+			Name:   "different_uuid",
+			Msg1:   message.NewMessage("1", []byte("foo")),
+			Msg2:   message.NewMessage("2", []byte("foo")),
+			Equals: false,
+		},
+		{
+			Name:   "different_payload",
+			Msg1:   message.NewMessage("1", []byte("foo")),
+			Msg2:   message.NewMessage("1", []byte("bar")),
+			Equals: false,
+		},
+		{
+			Name:   "different_metadata",
+			Msg1:   withMetadata(message.NewMessage("1", []byte("foo")), map[string]string{"foo": "1"}),
+			Msg2:   withMetadata(message.NewMessage("1", []byte("foo")), map[string]string{"foo": "2"}),
+			Equals: false,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.Name, func(t *testing.T) {
+			assert.Equal(t, tc.Equals, tc.Msg1.Equals(tc.Msg2))
+			assert.Equal(t, tc.Equals, tc.Msg2.Equals(tc.Msg1))
+		})
+	}
+}
+
+func TestMessage_Ack(t *testing.T) {
+	msg := &message.Message{}
+	require.True(t, msg.Ack())
+
+	assertAcked(t, msg)
+	assertNoNack(t, msg)
+}
+
+func TestMessage_Ack_idempotent(t *testing.T) {
+	msg := &message.Message{}
+	require.True(t, msg.Ack())
+	require.True(t, msg.Ack())
+
+	assertAcked(t, msg)
+}
+
+func TestMessage_Ack_already_Nack(t *testing.T) {
+	msg := &message.Message{}
+	require.True(t, msg.Nack())
+
+	assert.False(t, msg.Ack())
+}
+
+func TestMessage_Nack(t *testing.T) {
+	msg := &message.Message{}
+	require.True(t, msg.Nack())
+
+	assertNoAck(t, msg)
+	assertNacked(t, msg)
+}
+
+func TestMessage_Nack_idempotent(t *testing.T) {
+	msg := &message.Message{}
+	require.True(t, msg.Nack())
+	require.True(t, msg.Nack())
+
+	assertNacked(t, msg)
+}
+
+func TestMessage_Nack_already_Ack(t *testing.T) {
+	msg := &message.Message{}
+	require.True(t, msg.Ack())
+
+	assert.False(t, msg.Nack())
+}
+
+func TestMessage_Copy(t *testing.T) {
+	msg := message.NewMessage("1", []byte("foo"))
+	msgCopy := msg.Copy()
+
+	require.True(t, msg.Ack())
+
+	assertAcked(t, msg)
+	assertNoAck(t, msgCopy)
+	assert.True(t, msg.Equals(msgCopy))
+}
+
+type ctxKey string
+
+func TestMessage_CopyWithContext(t *testing.T) {
+	msg := message.NewMessage("1", []byte("foo"))
+	testCtx := context.Background()
+	testCtx = context.WithValue(testCtx, ctxKey("foo"), "bar")
+	msg.SetContext(testCtx)
+
+	msgCopy := msg.CopyWithContext()
+	copyMsgCtx := msgCopy.Context()
+	assert.True(t, copyMsgCtx.Value(ctxKey("foo")) == "bar", "expected context not being copied")
+	assert.False(t, copyMsgCtx.Value(ctxKey("abc")) == "def", "non-expected context being copied")
+	assert.True(t, msg.Equals(msgCopy))
+}
+
+func TestMessage_CopyWithContextAndMetadata(t *testing.T) {
+	msg := message.NewMessage("1", []byte("foo"))
+	testCtx := context.Background()
+	testCtx = context.WithValue(testCtx, ctxKey("foo"), "bar")
+	msg.SetContext(testCtx)
+	msg.Metadata.Set("foo", "bar")
+	msgCopy := msg.CopyWithContext()
+
+	msg.Metadata.Set("foo", "baz")
+
+	copyMsgCtx := msgCopy.Context()
+	assert.True(t, copyMsgCtx.Value(ctxKey("foo")) == "bar", "expected context not being copied")
+	assert.Equal(t, msgCopy.Metadata.Get("foo"), "bar", "did not expect changing source message's metadata to alter copy's metadata")
+}
+
+func TestMessage_CopyMetadata(t *testing.T) {
+	msg := message.NewMessage("1", []byte("foo"))
+	msg.Metadata.Set("foo", "bar")
+	msgCopy := msg.Copy()
+
+	msg.Metadata.Set("foo", "baz")
+
+	assert.Equal(t, msgCopy.Metadata.Get("foo"), "bar", "did not expect changing source message's metadata to alter copy's metadata")
+}
+
+func assertAcked(t *testing.T, msg *message.Message) {
+	select {
+	case <-msg.Acked():
+		// ok
+	default:
+		t.Fatal("no ack received")
+	}
+}
+
+func assertNacked(t *testing.T, msg *message.Message) {
+	select {
+	case <-msg.Nacked():
+		// ok
+	default:
+		t.Fatal("no ack received")
+	}
+}
+
+func assertNoAck(t *testing.T, msg *message.Message) {
+	select {
+	case <-msg.Acked():
+		t.Fatal("nack should be not sent")
+	default:
+		// ok
+	}
+}
+
+func assertNoNack(t *testing.T, msg *message.Message) {
+	select {
+	case <-msg.Nacked():
+		t.Fatal("nack should be not sent")
+	default:
+		// ok
+	}
+}

message/messages.go

@@ -0,0 +1,15 @@
+package message
+
+// Messages is a slice of messages.
+type Messages []*Message
+
+// IDs returns a slice of Messages' IDs.
+func (m Messages) IDs() []string {
+	ids := make([]string, len(m))
+
+	for i, msg := range m {
+		ids[i] = msg.UUID
+	}
+
+	return ids
+}

message/messages_test.go

@@ -0,0 +1,18 @@
+package message_test
+
+import (
+	"testing"
+
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestMessages_IDs(t *testing.T) {
+	msgs := message.Messages{
+		message.NewMessage("1", nil),
+		message.NewMessage("2", nil),
+		message.NewMessage("3", nil),
+	}
+
+	assert.Equal(t, []string{"1", "2", "3"}, msgs.IDs())
+}

message/metadata.go

@@ -0,0 +1,18 @@
+package message
+
+// Metadata is sent with every message to provide extra context without unmarshaling the message payload.
+type Metadata map[string]string
+
+// Get returns the metadata value for the given key. If the key is not found, an empty string is returned.
+func (m Metadata) Get(key string) string {
+	if v, ok := m[key]; ok {
+		return v
+	}
+
+	return ""
+}
+
+// Set sets the metadata key to value.
+func (m Metadata) Set(key, value string) {
+	m[key] = value
+}

message/pubsub.go

@@ -0,0 +1,51 @@
+package message
+
+import (
+	"context"
+)
+
+// Publisher is the emitting part of a Pub/Sub.
+type Publisher interface {
+	// Publish publishes provided messages to the given topic.
+	//
+	// Publish can be synchronous or asynchronous - it depends on the implementation.
+	//
+	// Most publisher implementations don't support atomic publishing of messages.
+	// This means that if publishing one of the messages fails, the next messages will not be published.
+	//
+	// Publish does not work with a single Context.
+	// Use the Context() method of each message instead.
+	//
+	// Publish must be thread safe.
+	Publish(topic string, messages ...*Message) error
+	// Close should flush unsent messages if publisher is async.
+	Close() error
+}
+
+// Subscriber is the consuming part of the Pub/Sub.
+type Subscriber interface {
+	// Subscribe returns an output channel with messages from the provided topic.
+	// The channel is closed after Close() is called on the subscriber.
+	//
+	// To receive the next message, `Ack()` must be called on the received message.
+	// If message processing fails and the message should be redelivered `Nack()` should be called instead.
+	//
+	// When the provided ctx is canceled, the subscriber closes the subscription and the output channel.
+	// The provided ctx is passed to all produced messages.
+	// When Nack or Ack is called on the message, the context of the message is canceled.
+	Subscribe(ctx context.Context, topic string) (<-chan *Message, error)
+	// Close closes all subscriptions with their output channels and flushes offsets etc. when needed.
+	Close() error
+}
+
+// SubscribeInitializer is used to initialize subscribers.
+type SubscribeInitializer interface {
+	// SubscribeInitialize can be called to initialize subscribe before consume.
+	// When calling Subscribe before Publish, SubscribeInitialize should be not required.
+	//
+	// Not every Pub/Sub requires this initialization, and it may be optional for performance improvements etc.
+	// For detailed SubscribeInitialize functionality, please check Pub/Subs godoc.
+	//
+	// Implementing SubscribeInitialize is not obligatory.
+	SubscribeInitialize(topic string) error
+}

message/router.go

@@ -0,0 +1,888 @@
+package message
+
+import (
+	"context"
+	"fmt"
+	"runtime/debug"
+	"sync"
+	"time"
+
+	"github.com/pkg/errors"
+
+	"github.com/ThreeDotsLabs/watermill"
+	"github.com/ThreeDotsLabs/watermill/internal"
+	sync_internal "github.com/ThreeDotsLabs/watermill/pubsub/sync"
+)
+
+var (
+	// ErrOutputInNoPublisherHandler happens when a handler func returned some messages in a no-publisher handler.
+	// todo: maybe change the handler func signature in no-publisher handler so that there's no possibility for this
+	ErrOutputInNoPublisherHandler = errors.New("returned output messages in a handler without publisher")
+)
+
+// HandlerFunc is function called when message is received.
+//
+// msg.Ack() is called automatically when HandlerFunc doesn't return error.
+// When HandlerFunc returns error, msg.Nack() is called.
+// When msg.Ack() was called in handler and HandlerFunc returns error,
+// msg.Nack() will be not sent because Ack was already sent.
+//
+// HandlerFunc's are executed parallel when multiple messages was received
+// (because msg.Ack() was sent in HandlerFunc or Subscriber supports multiple consumers).
+type HandlerFunc func(msg *Message) ([]*Message, error)
+
+// NoPublishHandlerFunc is HandlerFunc alternative, which doesn't produce any messages.
+type NoPublishHandlerFunc func(msg *Message) error
+
+// PassthroughHandler is a handler that passes the message unchanged from the subscriber to the publisher.
+var PassthroughHandler HandlerFunc = func(msg *Message) ([]*Message, error) {
+	return []*Message{msg}, nil
+}
+
+// HandlerMiddleware allows us to write something like decorators to HandlerFunc.
+// It can execute something before handler (for example: modify consumed message)
+// or after (modify produced messages, ack/nack on consumed message, handle errors, logging, etc.).
+//
+// It can be attached to the router by using `AddMiddleware` method.
+//
+// Example:
+//
+//	func ExampleMiddleware(h message.HandlerFunc) message.HandlerFunc {
+//		return func(message *message.Message) ([]*message.Message, error) {
+//			fmt.Println("executed before handler")
+//			producedMessages, err := h(message)
+//			fmt.Println("executed after handler")
+//
+//			return producedMessages, err
+//		}
+//	}
+type HandlerMiddleware func(h HandlerFunc) HandlerFunc
+
+// RouterPlugin is function which is executed on Router start.
+type RouterPlugin func(*Router) error
+
+// PublisherDecorator wraps the underlying Publisher, adding some functionality.
+type PublisherDecorator func(pub Publisher) (Publisher, error)
+
+// SubscriberDecorator wraps the underlying Subscriber, adding some functionality.
+type SubscriberDecorator func(sub Subscriber) (Subscriber, error)
+
+// RouterConfig holds the Router's configuration options.
+type RouterConfig struct {
+	// CloseTimeout determines how long router should work for handlers when closing.
+	CloseTimeout time.Duration
+}
+
+func (c *RouterConfig) setDefaults() {
+	if c.CloseTimeout == 0 {
+		c.CloseTimeout = time.Second * 30
+	}
+}
+
+// Validate returns Router configuration error, if any.
+func (c RouterConfig) Validate() error {
+	return nil
+}
+
+// NewRouter creates a new Router with given configuration.
+func NewRouter(config RouterConfig, logger watermill.LoggerAdapter) (*Router, error) {
+	config.setDefaults()
+	if err := config.Validate(); err != nil {
+		return nil, errors.Wrap(err, "invalid config")
+	}
+
+	return newRouter(config, logger), nil
+}
+
+// NewDefaultRouter creates a new Router with default configuration.
+func NewDefaultRouter(logger watermill.LoggerAdapter) *Router {
+	config := RouterConfig{}
+	config.setDefaults()
+
+	return newRouter(config, logger)
+}
+
+func newRouter(config RouterConfig, logger watermill.LoggerAdapter) *Router {
+	if logger == nil {
+		logger = watermill.NopLogger{}
+	}
+
+	return &Router{
+		config: config,
+
+		handlers: map[string]*handler{},
+
+		handlersWg: &sync.WaitGroup{},
+
+		runningHandlersWg:     &sync.WaitGroup{},
+		runningHandlersWgLock: &sync.Mutex{},
+
+		handlerAdded: make(chan struct{}),
+
+		middlewaresLock: &sync.RWMutex{},
+		handlersLock:    &sync.RWMutex{},
+
+		closingInProgressCh: make(chan struct{}),
+		closedCh:            make(chan struct{}),
+
+		logger: logger,
+
+		running: make(chan struct{}),
+	}
+}
+
+type middleware struct {
+	Handler       HandlerMiddleware
+	HandlerName   string
+	IsRouterLevel bool
+}
+
+// Router is responsible for handling messages from subscribers using provided handler functions.
+//
+// If the handler function returns a message, the message is published with the publisher.
+// You can use middlewares to wrap handlers with common logic like logging, instrumentation, etc.
+type Router struct {
+	config RouterConfig
+
+	middlewares     []middleware
+	middlewaresLock *sync.RWMutex
+
+	plugins []RouterPlugin
+
+	handlers     map[string]*handler
+	handlersLock *sync.RWMutex
+
+	handlersWg *sync.WaitGroup
+
+	runningHandlersWg     *sync.WaitGroup
+	runningHandlersWgLock *sync.Mutex
+
+	handlerAdded chan struct{}
+
+	closingInProgressCh chan struct{}
+	closedCh            chan struct{}
+	closed              bool
+	closedLock          sync.Mutex
+
+	logger watermill.LoggerAdapter
+
+	publisherDecorators  []PublisherDecorator
+	subscriberDecorators []SubscriberDecorator
+
+	isRunning bool
+	running   chan struct{}
+}
+
+// Logger returns the Router's logger.
+func (r *Router) Logger() watermill.LoggerAdapter {
+	return r.logger
+}
+
+// AddMiddleware adds a new middleware to the router.
+//
+// The order of middleware matters. Middleware added at the beginning is executed first.
+func (r *Router) AddMiddleware(m ...HandlerMiddleware) {
+	r.logger.Debug("Adding middleware", watermill.LogFields{"count": fmt.Sprintf("%d", len(m))})
+
+	r.addRouterLevelMiddleware(m...)
+}
+
+func (r *Router) addRouterLevelMiddleware(m ...HandlerMiddleware) {
+	for _, handlerMiddleware := range m {
+		middleware := middleware{
+			Handler:       handlerMiddleware,
+			HandlerName:   "",
+			IsRouterLevel: true,
+		}
+		r.middlewares = append(r.middlewares, middleware)
+	}
+}
+
+func (r *Router) addHandlerLevelMiddleware(handlerName string, m ...HandlerMiddleware) {
+	r.middlewaresLock.Lock()
+	defer r.middlewaresLock.Unlock()
+	for _, handlerMiddleware := range m {
+		middleware := middleware{
+			Handler:       handlerMiddleware,
+			HandlerName:   handlerName,
+			IsRouterLevel: false,
+		}
+		r.middlewares = append(r.middlewares, middleware)
+	}
+}
+
+// AddPlugin adds a new plugin to the router.
+// Plugins are executed during startup of the router.
+//
+// A plugin can, for example, close the router after SIGINT or SIGTERM is sent to the process (SignalsHandler plugin).
+func (r *Router) AddPlugin(p ...RouterPlugin) {
+	r.logger.Debug("Adding plugins", watermill.LogFields{"count": fmt.Sprintf("%d", len(p))})
+
+	r.plugins = append(r.plugins, p...)
+}
+
+// AddPublisherDecorators wraps the router's Publisher.
+// The first decorator is the innermost, i.e. calls the original publisher.
+func (r *Router) AddPublisherDecorators(dec ...PublisherDecorator) {
+	r.logger.Debug("Adding publisher decorators", watermill.LogFields{"count": fmt.Sprintf("%d", len(dec))})
+
+	r.publisherDecorators = append(r.publisherDecorators, dec...)
+}
+
+// AddSubscriberDecorators wraps the router's Subscriber.
+// The first decorator is the innermost, i.e. calls the original subscriber.
+func (r *Router) AddSubscriberDecorators(dec ...SubscriberDecorator) {
+	r.logger.Debug("Adding subscriber decorators", watermill.LogFields{"count": fmt.Sprintf("%d", len(dec))})
+
+	r.subscriberDecorators = append(r.subscriberDecorators, dec...)
+}
+
+// Handlers returns all registered handlers.
+func (r *Router) Handlers() map[string]HandlerFunc {
+	handlers := map[string]HandlerFunc{}
+
+	for handlerName, handler := range r.handlers {
+		handlers[handlerName] = handler.handlerFunc
+	}
+
+	return handlers
+}
+
+// DuplicateHandlerNameError is sent in a panic when you try to add a second handler with the same name.
+type DuplicateHandlerNameError struct {
+	HandlerName string
+}
+
+func (d DuplicateHandlerNameError) Error() string {
+	return fmt.Sprintf("handler with name %s already exists", d.HandlerName)
+}
+
+// AddHandler adds a new handler.
+//
+// handlerName must be unique. For now, it is used only for debugging.
+//
+// subscribeTopic is a topic from which handler will receive messages.
+//
+// publishTopic is a topic to which router will produce messages returned by handlerFunc.
+// When handler needs to publish to multiple topics,
+// it is recommended to use AddConsumerHandler and inject a Publisher or implement middleware
+// which will catch messages and publish to topic based on metadata for example.
+//
+// If handler is added while router is already running, you need to explicitly call RunHandlers().
+func (r *Router) AddHandler(
+	handlerName string,
+	subscribeTopic string,
+	subscriber Subscriber,
+	publishTopic string,
+	publisher Publisher,
+	handlerFunc HandlerFunc,
+) *Handler {
+	r.logger.Info("Adding handler", watermill.LogFields{
+		"handler_name": handlerName,
+		"topic":        subscribeTopic,
+	})
+
+	r.handlersLock.Lock()
+	defer r.handlersLock.Unlock()
+
+	if _, ok := r.handlers[handlerName]; ok {
+		panic(DuplicateHandlerNameError{handlerName})
+	}
+
+	publisherName, subscriberName := internal.StructName(publisher), internal.StructName(subscriber)
+
+	newHandler := &handler{
+		name:   handlerName,
+		logger: r.logger,
+
+		subscriber:     subscriber,
+		subscribeTopic: subscribeTopic,
+		subscriberName: subscriberName,
+
+		publisher:     publisher,
+		publishTopic:  publishTopic,
+		publisherName: publisherName,
+
+		handlerFunc: handlerFunc,
+
+		runningHandlersWg:     r.runningHandlersWg,
+		runningHandlersWgLock: r.runningHandlersWgLock,
+
+		messagesCh:     nil,
+		routersCloseCh: r.closingInProgressCh,
+
+		startedCh: make(chan struct{}),
+	}
+
+	r.handlersWg.Add(1)
+	r.handlers[handlerName] = newHandler
+
+	select {
+	case r.handlerAdded <- struct{}{}:
+	default:
+		// watchAllHandlersStopped is not always waiting for handlerAdded
+	}
+
+	return &Handler{
+		router:  r,
+		handler: newHandler,
+	}
+}
+
+// AddConsumerHandler adds a new handler that does not return any messages.
+// It can publish messages by directly using a Publisher.
+//
+// handlerName must be unique. For now, it is used only for debugging.
+//
+// subscribeTopic is a topic from which handler will receive messages.
+//
+// subscriber is Subscriber from which messages will be consumed.
+//
+// If handler is added while router is already running, you need to explicitly call RunHandlers().
+func (r *Router) AddConsumerHandler(
+	handlerName string,
+	subscribeTopic string,
+	subscriber Subscriber,
+	handlerFunc NoPublishHandlerFunc,
+) *Handler {
+	handlerFuncAdapter := func(msg *Message) ([]*Message, error) {
+		return nil, handlerFunc(msg)
+	}
+
+	return r.AddHandler(handlerName, subscribeTopic, subscriber, "", disabledPublisher{}, handlerFuncAdapter)
+}
+
+// AddNoPublisherHandler adds a new handler.
+// This handler cannot return messages.
+//
+// handlerName must be unique. For now, it is used only for debugging.
+//
+// subscribeTopic is a topic from which handler will receive messages.
+//
+// subscriber is Subscriber from which messages will be consumed.
+//
+// If handler is added while router is already running, you need to explicitly call RunHandlers().
+//
+// Deprecated: use AddConsumerHandler instead.
+func (r *Router) AddNoPublisherHandler(
+	handlerName string,
+	subscribeTopic string,
+	subscriber Subscriber,
+	handlerFunc NoPublishHandlerFunc,
+) *Handler {
+	return r.AddConsumerHandler(handlerName, subscribeTopic, subscriber, handlerFunc)
+}
+
+// Run runs all plugins and handlers and starts subscribing to provided topics.
+// This call is blocking while the router is running.
+//
+// When all handlers have stopped (for example, because subscriptions were closed), the router will also stop.
+//
+// To stop Run() you should call Close() on the router.
+//
+// ctx will be propagated to all subscribers.
+//
+// When all handlers are stopped (for example: because of closed connection), Run() will be also stopped.
+func (r *Router) Run(ctx context.Context) (err error) {
+	if r.isRunning {
+		return errors.New("router is already running")
+	}
+	r.isRunning = true
+
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	r.logger.Debug("Loading plugins", nil)
+	for _, plugin := range r.plugins {
+		if err := plugin(r); err != nil {
+			return errors.Wrapf(err, "cannot initialize plugin %v", plugin)
+		}
+	}
+
+	r.watchAllHandlersStopped(ctx)
+
+	if err := r.RunHandlers(ctx); err != nil {
+		return err
+	}
+
+	close(r.running)
+
+	<-r.closingInProgressCh
+	cancel()
+
+	r.logger.Info("Waiting for messages", watermill.LogFields{
+		"timeout": r.config.CloseTimeout,
+	})
+
+	<-r.closedCh
+
+	r.logger.Info("All messages processed", nil)
+
+	return nil
+}
+
+// RunHandlers runs all handlers that were added after Run().
+// RunHandlers is idempotent, so can be called multiple times safely.
+func (r *Router) RunHandlers(ctx context.Context) error {
+	if !r.isRunning {
+		return errors.New("you can't call RunHandlers on non-running router")
+	}
+
+	r.handlersLock.Lock()
+	defer r.handlersLock.Unlock()
+
+	r.logger.Info("Running router handlers", watermill.LogFields{"count": len(r.handlers)})
+
+	for name, h := range r.handlers {
+
+		if h.started {
+			continue
+		}
+
+		if err := r.decorateHandlerPublisher(h); err != nil {
+			return errors.Wrapf(err, "could not decorate publisher of handler %s", name)
+		}
+		if err := r.decorateHandlerSubscriber(h); err != nil {
+			return errors.Wrapf(err, "could not decorate subscriber of handler %s", name)
+		}
+
+		logger := r.logger.With(watermill.LogFields{
+			"subscriber_name": h.name,
+			"topic":           h.subscribeTopic,
+		})
+
+		logger.Debug("Subscribing to topic", nil)
+
+		ctx, cancel := context.WithCancel(ctx)
+
+		messages, err := h.subscriber.Subscribe(ctx, h.subscribeTopic)
+		if err != nil {
+			cancel()
+			return errors.Wrapf(err, "cannot subscribe topic %s", h.subscribeTopic)
+		}
+
+		h.messagesCh = messages
+		h.started = true
+		close(h.startedCh)
+
+		h.stopFn = cancel
+		h.stopped = make(chan struct{})
+
+		go func() {
+			defer cancel()
+
+			r.middlewaresLock.Lock()
+			middlewares := append([]middleware{}, r.middlewares...)
+			r.middlewaresLock.Unlock()
+
+			h.run(ctx, middlewares)
+
+			r.handlersWg.Done()
+			logger.Info("Subscriber stopped", nil)
+
+			r.handlersLock.Lock()
+			delete(r.handlers, name)
+			r.handlersLock.Unlock()
+
+			logger.Trace("Removed subscriber from r.handlers", nil)
+
+			close(h.stopped)
+		}()
+	}
+	return nil
+}
+
+// watchAllHandlersStopped closes router when all handlers have stopped,
+// (for example, because for example all subscriptions are closed)
+func (r *Router) watchAllHandlersStopped(ctx context.Context) {
+	r.handlersLock.RLock()
+	hasNoHandlersYet := len(r.handlers) == 0
+	r.handlersLock.RUnlock()
+
+	go func() {
+		if hasNoHandlersYet {
+			// we can start router without any handlers,
+			// in that situation router would be closed immediately (even if they are no routers)
+			// let's wait for
+			select {
+			case <-r.handlerAdded:
+				// it should be some handler to track
+			case <-r.closedCh:
+				// let's avoid goroutine leak
+				return
+			}
+		}
+
+		r.handlersWg.Wait()
+		if r.IsClosed() {
+			r.logger.Trace("watchAllHandlersStopped: already closed", nil)
+			// already closed
+			return
+		}
+
+		// Only log an error if the context was not canceled, but handlers were stopped.
+		select {
+		case <-ctx.Done():
+		default:
+			r.logger.Error("All handlers stopped, closing router", errors.New("all router handlers stopped"), nil)
+		}
+
+		if err := r.Close(); err != nil {
+			r.logger.Error("Cannot close router", err, nil)
+		}
+	}()
+}
+
+// Running is closed when router is running.
+// In other words: you can wait till router is running using
+//
+//	fmt.Println("Starting router")
+//	go r.Run(ctx)
+//	<- r.Running()
+//	fmt.Println("Router is running")
+//
+// Warning: for historical reasons, this channel is not aware of router closing - the channel will be closed if the router has been running and closed.
+func (r *Router) Running() chan struct{} {
+	return r.running
+}
+
+// IsRunning returns true when router is running.
+//
+// Warning: for historical reasons, this method is not aware of router closing.
+// If you want to know if the router was closed, use IsClosed.
+func (r *Router) IsRunning() bool {
+	select {
+	case <-r.running:
+		return true
+	default:
+		return false
+	}
+}
+
+// Close gracefully closes the router with a timeout provided in the configuration.
+func (r *Router) Close() error {
+	r.closedLock.Lock()
+	defer r.closedLock.Unlock()
+
+	r.handlersLock.Lock()
+	defer r.handlersLock.Unlock()
+
+	if r.closed {
+		r.logger.Debug("Already closed", nil)
+		return nil
+	}
+
+	r.logger.Debug("Running Close()", nil)
+	r.closed = true
+
+	r.logger.Info("Closing router", nil)
+	defer r.logger.Info("Router closed", nil)
+
+	close(r.closingInProgressCh)
+	defer close(r.closedCh)
+
+	timedout := r.waitForHandlers()
+	if timedout {
+		return errors.New("router close timeout")
+	}
+
+	return nil
+}
+
+func (r *Router) waitForHandlers() bool {
+	var waitGroup sync.WaitGroup
+	waitGroup.Add(1)
+	go func() {
+		defer waitGroup.Done()
+		r.handlersWg.Wait()
+	}()
+	waitGroup.Add(1)
+	go func() {
+		defer waitGroup.Done()
+
+		r.runningHandlersWgLock.Lock()
+		defer r.runningHandlersWgLock.Unlock()
+
+		r.runningHandlersWg.Wait()
+	}()
+	return sync_internal.WaitGroupTimeout(&waitGroup, r.config.CloseTimeout)
+}
+
+func (r *Router) IsClosed() bool {
+	r.closedLock.Lock()
+	defer r.closedLock.Unlock()
+
+	return r.closed
+}
+
+type handler struct {
+	name   string
+	logger watermill.LoggerAdapter
+
+	subscriber     Subscriber
+	subscribeTopic string
+	subscriberName string
+
+	publisher     Publisher
+	publishTopic  string
+	publisherName string
+
+	handlerFunc HandlerFunc
+
+	runningHandlersWg     *sync.WaitGroup
+	runningHandlersWgLock *sync.Mutex
+
+	messagesCh <-chan *Message
+
+	started   bool
+	startedCh chan struct{}
+
+	stopFn         context.CancelFunc
+	stopped        chan struct{}
+	routersCloseCh chan struct{}
+}
+
+func (h *handler) run(ctx context.Context, middlewares []middleware) {
+	h.logger.Info("Starting handler", watermill.LogFields{
+		"subscriber_name": h.name,
+		"topic":           h.subscribeTopic,
+	})
+
+	middlewareHandler := h.handlerFunc
+	// first added middlewares should be executed first (so should be at the top of call stack)
+	for i := len(middlewares) - 1; i >= 0; i-- {
+		currentMiddleware := middlewares[i]
+		isValidHandlerLevelMiddleware := currentMiddleware.HandlerName == h.name
+		if currentMiddleware.IsRouterLevel || isValidHandlerLevelMiddleware {
+			middlewareHandler = currentMiddleware.Handler(middlewareHandler)
+		}
+	}
+
+	go h.handleClose(ctx)
+
+	for msg := range h.messagesCh {
+		h.runningHandlersWgLock.Lock()
+		h.runningHandlersWg.Add(1)
+		h.runningHandlersWgLock.Unlock()
+
+		go h.handleMessage(msg, middlewareHandler)
+	}
+
+	if h.publisher != nil {
+		h.logger.Debug("Waiting for publisher to close", nil)
+		if err := h.publisher.Close(); err != nil {
+			h.logger.Error("Failed to close publisher", err, nil)
+		}
+		h.logger.Debug("Publisher closed", nil)
+	}
+
+	h.logger.Debug("Router handler stopped", nil)
+}
+
+// Handler handles Messages.
+type Handler struct {
+	router  *Router
+	handler *handler
+}
+
+// AddMiddleware adds new middleware to the specified handler in the router.
+//
+// The order of middleware matters. Middleware added at the beginning is executed first.
+func (h *Handler) AddMiddleware(m ...HandlerMiddleware) {
+	handler := h.handler
+	handler.logger.Debug("Adding middleware to handler", watermill.LogFields{
+		"count":       fmt.Sprintf("%d", len(m)),
+		"handlerName": handler.name,
+	})
+
+	h.router.addHandlerLevelMiddleware(handler.name, m...)
+}
+
+// Started returns channel which is stopped when handler is running.
+func (h *Handler) Started() chan struct{} {
+	return h.handler.startedCh
+}
+
+// Stop stops the handler.
+// Stop is asynchronous.
+// You can check if handler was stopped with Stopped() function.
+func (h *Handler) Stop() {
+	if !h.handler.started {
+		panic("handler is not started")
+	}
+
+	h.handler.stopFn()
+}
+
+// Stopped returns channel which is stopped when handler did stop.
+func (h *Handler) Stopped() chan struct{} {
+	return h.handler.stopped
+}
+
+// decorateHandlerPublisher applies the decorator chain to handler's publisher.
+// They are applied in reverse order, so that the later decorators use the result of former ones.
+func (r *Router) decorateHandlerPublisher(h *handler) error {
+	var err error
+	pub := h.publisher
+	for i := len(r.publisherDecorators) - 1; i >= 0; i-- {
+		decorator := r.publisherDecorators[i]
+		pub, err = decorator(pub)
+		if err != nil {
+			return errors.Wrap(err, "could not apply publisher decorator")
+		}
+	}
+	r.handlers[h.name].publisher = pub
+	return nil
+}
+
+// decorateHandlerSubscriber applies the decorator chain to handler's subscriber.
+// They are applied in regular order, so that the later decorators use the result of former ones.
+func (r *Router) decorateHandlerSubscriber(h *handler) error {
+	var err error
+	sub := h.subscriber
+
+	// add values to message context to subscriber
+	// it goes before other decorators, so that they may take advantage of these values
+	messageTransform := func(msg *Message) {
+		if msg != nil {
+			h.addHandlerContext(msg)
+		}
+	}
+	sub, err = MessageTransformSubscriberDecorator(messageTransform)(sub)
+	if err != nil {
+		return errors.Wrapf(err, "cannot wrap subscriber with context decorator")
+	}
+
+	for _, decorator := range r.subscriberDecorators {
+		sub, err = decorator(sub)
+		if err != nil {
+			return errors.Wrap(err, "could not apply subscriber decorator")
+		}
+	}
+	r.handlers[h.name].subscriber = sub
+	return nil
+}
+
+// addHandlerContext enriches the context with values that are relevant within this handler's context.
+func (h *handler) addHandlerContext(messages ...*Message) {
+	for i, msg := range messages {
+		ctx := msg.Context()
+
+		if h.name != "" {
+			ctx = context.WithValue(ctx, handlerNameKey, h.name)
+		}
+		if h.publisherName != "" {
+			ctx = context.WithValue(ctx, publisherNameKey, h.publisherName)
+		}
+		if h.subscriberName != "" {
+			ctx = context.WithValue(ctx, subscriberNameKey, h.subscriberName)
+		}
+		if h.subscribeTopic != "" {
+			ctx = context.WithValue(ctx, subscribeTopicKey, h.subscribeTopic)
+		}
+		if h.publishTopic != "" {
+			ctx = context.WithValue(ctx, publishTopicKey, h.publishTopic)
+		}
+		messages[i].SetContext(ctx)
+	}
+}
+
+func (h *handler) handleClose(ctx context.Context) {
+	closeSubscriber := func() {
+		h.logger.Debug("Waiting for subscriber to close", nil)
+		if err := h.subscriber.Close(); err != nil {
+			h.logger.Error("Failed to close subscriber", err, nil)
+		}
+		h.logger.Debug("Subscriber closed", nil)
+	}
+
+	select {
+	case <-h.routersCloseCh:
+		closeSubscriber()
+	case <-ctx.Done():
+		// ctx.Done can race with routersCloseCh during a router-wide close:
+		// Router.Close closes routersCloseCh, then Router.Run cancels the parent ctx,
+		// and the select picks whichever fires first. If routersCloseCh is also closed
+		// here, this is a router-wide close and the subscriber must be closed —
+		// otherwise it's Handler.Stop() and the subscriber may be shared.
+		select {
+		case <-h.routersCloseCh:
+			closeSubscriber()
+		default:
+		}
+	}
+	h.stopFn()
+}
+
+func (h *handler) handleMessage(msg *Message, handler HandlerFunc) {
+	defer h.runningHandlersWg.Done()
+	msgFields := watermill.LogFields{"message_uuid": msg.UUID, "handler_name": h.name}
+
+	defer func() {
+		if recovered := recover(); recovered != nil {
+			h.logger.Error(
+				"Panic recovered in handler. Stack: "+string(debug.Stack()),
+				errors.Errorf("%s", recovered),
+				msgFields,
+			)
+			msg.Nack()
+		}
+	}()
+
+	h.logger.Trace("Received message", msgFields)
+
+	producedMessages, err := handler(msg)
+	if err != nil {
+		if !errors.Is(err, context.Canceled) {
+			h.logger.Error("Handler returned error", err, msgFields)
+		}
+		msg.Nack()
+		return
+	}
+
+	h.addHandlerContext(producedMessages...)
+
+	if err := h.publishProducedMessages(producedMessages, msgFields); err != nil {
+		h.logger.Error("Publishing produced messages failed", err, nil)
+		msg.Nack()
+		return
+	}
+
+	msg.Ack()
+	h.logger.Trace("Message acked", msgFields)
+}
+
+func (h *handler) publishProducedMessages(producedMessages Messages, msgFields watermill.LogFields) error {
+	if len(producedMessages) == 0 {
+		return nil
+	}
+
+	if h.publisher == nil {
+		return ErrOutputInNoPublisherHandler
+	}
+
+	h.logger.Trace("Sending produced messages", msgFields.Add(watermill.LogFields{
+		"produced_messages_count": len(producedMessages),
+		"publish_topic":           h.publishTopic,
+	}))
+
+	if err := h.publisher.Publish(h.publishTopic, producedMessages...); err != nil {
+		h.logger.Error("Cannot publish messages", err, msgFields.Add(watermill.LogFields{
+			"not_sent_message": fmt.Sprintf("%#v", producedMessages),
+		}))
+
+		return err
+	}
+
+	return nil
+}
+
+type disabledPublisher struct{}
+
+func (disabledPublisher) Publish(topic string, messages ...*Message) error {
+	return ErrOutputInNoPublisherHandler
+}
+
+func (disabledPublisher) Close() error {
+	return nil
+}

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

message/router_context.go

@@ -0,0 +1,50 @@
+package message
+
+import (
+	"context"
+)
+
+type ctxKey string
+
+const (
+	handlerNameKey    ctxKey = "handler_name"
+	publisherNameKey  ctxKey = "publisher_name"
+	subscriberNameKey ctxKey = "subscriber_name"
+	subscribeTopicKey ctxKey = "subscribe_topic"
+	publishTopicKey   ctxKey = "publish_topic"
+)
+
+func valFromCtx(ctx context.Context, key ctxKey) string {
+	val, ok := ctx.Value(key).(string)
+	if !ok {
+		return ""
+	}
+	return val
+}
+
+// HandlerNameFromCtx returns the name of the message handler in the router that consumed the message.
+func HandlerNameFromCtx(ctx context.Context) string {
+	return valFromCtx(ctx, handlerNameKey)
+}
+
+// PublisherNameFromCtx returns the name of the message publisher type that published the message in the router.
+// For example, for Kafka it will be `kafka.Publisher`.
+func PublisherNameFromCtx(ctx context.Context) string {
+	return valFromCtx(ctx, publisherNameKey)
+}
+
+// SubscriberNameFromCtx returns the name of the message subscriber type that subscribed to the message in the router.
+// For example, for Kafka it will be `kafka.Subscriber`.
+func SubscriberNameFromCtx(ctx context.Context) string {
+	return valFromCtx(ctx, subscriberNameKey)
+}
+
+// SubscribeTopicFromCtx returns the topic from which message was received in the router.
+func SubscribeTopicFromCtx(ctx context.Context) string {
+	return valFromCtx(ctx, subscribeTopicKey)
+}
+
+// PublishTopicFromCtx returns the topic to which message will be published by the router.
+func PublishTopicFromCtx(ctx context.Context) string {
+	return valFromCtx(ctx, publishTopicKey)
+}

message/router_context_test.go

@@ -0,0 +1,148 @@
+package message_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/ThreeDotsLabs/watermill"
+	"github.com/ThreeDotsLabs/watermill/message"
+)
+
+type namedMockPublisher struct{}
+
+func (namedMockPublisher) Publish(topic string, messages ...*message.Message) error { return nil }
+func (namedMockPublisher) Close() error                                             { return nil }
+func (namedMockPublisher) String() string {
+	return "this publisher implements Stringer"
+}
+
+type namedMockSubscriber struct{ ch chan *message.Message }
+
+func (s namedMockSubscriber) Subscribe(context.Context, string) (<-chan *message.Message, error) {
+	return s.ch, nil
+}
+func (s *namedMockSubscriber) Close() error { close(s.ch); return nil }
+func (namedMockSubscriber) String() string {
+	return "this subscriber implements Stringer"
+}
+
+func TestRouter_Context_Stringer(t *testing.T) {
+	// If a publisher or subscriber implements Stringer, it's name is the result of String().
+	// The messages processed by a router handler should have publisher and subscriber name in their context.
+
+	// given
+	capturedMessages := make(chan *message.Message)
+	router, handlerFunc := setupPubsubNameTests(t, capturedMessages)
+
+	sub := &namedMockSubscriber{make(chan *message.Message)}
+	pub := namedMockPublisher{}
+
+	handlerName := "handler_name_stringer_test"
+	router.AddHandler(
+		handlerName,
+		"sub-topic",
+		sub,
+		"pub-topic",
+		pub,
+		handlerFunc,
+	)
+
+	go func() {
+		if err := router.Run(context.Background()); err != nil {
+			panic(err)
+		}
+	}()
+	defer func() {
+		if err := router.Close(); err != nil {
+			panic(err)
+		}
+	}()
+	<-router.Running()
+
+	// when
+	sub.ch <- message.NewMessage("", []byte{})
+	capturedMsg := <-capturedMessages
+
+	ctx := capturedMsg.Context()
+
+	// then
+	require.Equal(t, handlerName, message.HandlerNameFromCtx(ctx))
+	require.Equal(t, sub.String(), message.SubscriberNameFromCtx(ctx))
+	require.Equal(t, pub.String(), message.PublisherNameFromCtx(ctx))
+	require.Equal(t, "sub-topic", message.SubscribeTopicFromCtx(ctx))
+	require.Equal(t, "pub-topic", message.PublishTopicFromCtx(ctx))
+}
+
+type unnamedMockPublisher struct{}
+
+func (unnamedMockPublisher) Publish(topic string, messages ...*message.Message) error { return nil }
+func (unnamedMockPublisher) Close() error                                             { return nil }
+
+type unnamedMockSubscriber struct{ ch chan *message.Message }
+
+func (s unnamedMockSubscriber) Subscribe(context.Context, string) (<-chan *message.Message, error) {
+	return s.ch, nil
+}
+func (s *unnamedMockSubscriber) Close() error { close(s.ch); return nil }
+
+func TestRouter_Context_TypeName(t *testing.T) {
+	// If a publisher or subscriber does not implement Stringer, it's name is the type name.
+	// The messages processed by a router handler should have publisher and subscriber name in their context.
+
+	// given
+	capturedMessages := make(chan *message.Message)
+	router, handlerFunc := setupPubsubNameTests(t, capturedMessages)
+
+	sub := &unnamedMockSubscriber{make(chan *message.Message)}
+	pub := unnamedMockPublisher{}
+
+	handlerName := "handler_name_typename_test"
+	router.AddHandler(
+		handlerName,
+		"",
+		sub,
+		"",
+		pub,
+		handlerFunc,
+	)
+
+	go func() {
+		if err := router.Run(context.Background()); err != nil {
+			panic(err)
+		}
+	}()
+	defer func() {
+		if err := router.Close(); err != nil {
+			panic(err)
+		}
+	}()
+	<-router.Running()
+
+	// when
+	sub.ch <- message.NewMessage("", []byte{})
+	capturedMsg := <-capturedMessages
+
+	ctx := capturedMsg.Context()
+
+	// then
+	require.Equal(t, handlerName, message.HandlerNameFromCtx(ctx))
+	require.Equal(t, "message_test.unnamedMockSubscriber", message.SubscriberNameFromCtx(ctx))
+	require.Equal(t, "message_test.unnamedMockPublisher", message.PublisherNameFromCtx(ctx))
+}
+
+func setupPubsubNameTests(t *testing.T, capturedMessages chan (*message.Message)) (*message.Router, message.HandlerFunc) {
+	logger := watermill.NewStdLogger(true, true)
+
+	router, err := message.NewRouter(message.RouterConfig{}, logger)
+	require.NoError(t, err)
+
+	handlerFunc := func(msg *message.Message) ([]*message.Message, error) {
+		capturedMessages <- msg
+		require.True(t, msg.Ack())
+		return message.Messages{message.NewMessage(msg.UUID+"_copy", msg.Payload)}, nil
+	}
+
+	return router, handlerFunc
+}

message/subscriber/read.go

@@ -0,0 +1,53 @@
+package subscriber
+
+import (
+	"time"
+
+	"github.com/ThreeDotsLabs/watermill/message"
+)
+
+// BulkRead reads provided amount of messages from the provided channel, until a timeout occurs or the limit is reached.
+func BulkRead(messagesCh <-chan *message.Message, limit int, timeout time.Duration) (receivedMessages message.Messages, all bool) {
+MessagesLoop:
+	for len(receivedMessages) < limit {
+		select {
+		case msg, ok := <-messagesCh:
+			if !ok {
+				break MessagesLoop
+			}
+
+			receivedMessages = append(receivedMessages, msg)
+			msg.Ack()
+		case <-time.After(timeout):
+			break MessagesLoop
+		}
+	}
+
+	return receivedMessages, len(receivedMessages) == limit
+}
+
+// BulkReadWithDeduplication reads provided number of messages from the provided channel, ignoring duplicates,
+// until a timeout occurs or the limit is reached.
+func BulkReadWithDeduplication(messagesCh <-chan *message.Message, limit int, timeout time.Duration) (receivedMessages message.Messages, all bool) {
+	receivedIDs := map[string]struct{}{}
+
+MessagesLoop:
+	for len(receivedMessages) < limit {
+		select {
+		case msg, ok := <-messagesCh:
+			if !ok {
+				break MessagesLoop
+			}
+
+			if _, ok := receivedIDs[msg.UUID]; !ok {
+				receivedIDs[msg.UUID] = struct{}{}
+				receivedMessages = append(receivedMessages, msg)
+			}
+			msg.Ack()
+		case <-time.After(timeout):
+			break MessagesLoop
+		}
+	}
+
+	return receivedMessages, len(receivedMessages) == limit
+}

message/subscriber/read_test.go

@@ -0,0 +1,185 @@
+package subscriber_test
+
+import (
+	"testing"
+	"time"
+
+	"github.com/ThreeDotsLabs/watermill/pubsub/tests"
+
+	"github.com/ThreeDotsLabs/watermill"
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/ThreeDotsLabs/watermill/message/subscriber"
+	"github.com/stretchr/testify/assert"
+)
+
+type bulkReadFunc func(messagesCh <-chan *message.Message, limit int, timeout time.Duration) (receivedMessages message.Messages, all bool)
+
+func TestBulkRead(t *testing.T) {
+	testCases := []struct {
+		Name         string
+		BulkReadFunc bulkReadFunc
+	}{
+		{
+			Name:         "BulkRead",
+			BulkReadFunc: subscriber.BulkRead,
+		},
+		{
+			Name:         "BulkReadWithDeduplication",
+			BulkReadFunc: subscriber.BulkReadWithDeduplication,
+		},
+	}
+
+	for _, c := range testCases {
+		t.Run(c.Name, func(t *testing.T) {
+			messagesCount := 100
+
+			var messages []*message.Message
+			messagesCh := make(chan *message.Message, messagesCount)
+
+			for i := 0; i < messagesCount; i++ {
+				msg := message.NewMessage(watermill.NewUUID(), nil)
+
+				messages = append(messages, msg)
+				messagesCh <- msg
+			}
+
+			readMessages, all := subscriber.BulkRead(messagesCh, messagesCount, time.Second)
+			assert.True(t, all)
+
+			tests.AssertAllMessagesReceived(t, messages, readMessages)
+		})
+	}
+}
+
+func TestBulkRead_timeout(t *testing.T) {
+	testCases := []struct {
+		Name         string
+		BulkReadFunc bulkReadFunc
+	}{
+		{
+			Name:         "BulkRead",
+			BulkReadFunc: subscriber.BulkRead,
+		},
+		{
+			Name:         "BulkReadWithDeduplication",
+			BulkReadFunc: subscriber.BulkReadWithDeduplication,
+		},
+	}
+
+	for _, c := range testCases {
+		t.Run(c.Name, func(t *testing.T) {
+			messagesCount := 100
+			sendLimit := 90
+
+			messagesCh := make(chan *message.Message, messagesCount)
+
+			for i := 0; i < messagesCount; i++ {
+				msg := message.NewMessage(watermill.NewUUID(), nil)
+
+				if i < sendLimit {
+					messagesCh <- msg
+				}
+			}
+
+			bulkReadStart := time.Now()
+			readMessages, all := subscriber.BulkRead(messagesCh, messagesCount, time.Millisecond)
+
+			assert.WithinDuration(t, bulkReadStart, time.Now(), time.Millisecond*100)
+			assert.False(t, all)
+			assert.Equal(t, sendLimit, len(readMessages))
+		})
+	}
+}
+
+func TestBulkRead_with_limit(t *testing.T) {
+	testCases := []struct {
+		Name         string
+		BulkReadFunc bulkReadFunc
+	}{
+		{
+			Name:         "BulkRead",
+			BulkReadFunc: subscriber.BulkRead,
+		},
+		{
+			Name:         "BulkReadWithDeduplication",
+			BulkReadFunc: subscriber.BulkReadWithDeduplication,
+		},
+	}
+
+	for _, c := range testCases {
+		t.Run(c.Name, func(t *testing.T) {
+			messagesCount := 110
+			limit := 100
+
+			messagesCh := make(chan *message.Message, messagesCount)
+
+			for i := 0; i < messagesCount; i++ {
+				msg := message.NewMessage(watermill.NewUUID(), nil)
+
+				messagesCh <- msg
+			}
+
+			readMessages, all := subscriber.BulkRead(messagesCh, limit, time.Second)
+			assert.True(t, all)
+			assert.Equal(t, limit, len(readMessages))
+		})
+	}
+}
+
+func TestBulkRead_return_on_channel_close(t *testing.T) {
+	testCases := []struct {
+		Name         string
+		BulkReadFunc bulkReadFunc
+	}{
+		{
+			Name:         "BulkRead",
+			BulkReadFunc: subscriber.BulkRead,
+		},
+		{
+			Name:         "BulkReadWithDeduplication",
+			BulkReadFunc: subscriber.BulkReadWithDeduplication,
+		},
+	}
+
+	for _, c := range testCases {
+		t.Run(c.Name, func(t *testing.T) {
+			messagesCount := 100
+			sendLimit := 90
+
+			messagesCh := make(chan *message.Message, messagesCount)
+			messagesChClosed := false
+
+			for i := 0; i < messagesCount; i++ {
+				msg := message.NewMessage(watermill.NewUUID(), nil)
+
+				if i < sendLimit {
+					messagesCh <- msg
+				} else if !messagesChClosed {
+					close(messagesCh)
+					messagesChClosed = true
+				}
+			}
+
+			bulkReadStart := time.Now()
+			_, all := subscriber.BulkRead(messagesCh, messagesCount, time.Second)
+
+			assert.WithinDuration(t, bulkReadStart, time.Now(), time.Millisecond*100)
+			assert.False(t, all)
+		})
+	}
+}
+
+func TestBulkReadWithDeduplication(t *testing.T) {
+	messagesCh := make(chan *message.Message, 3)
+
+	msg1 := message.NewMessage(watermill.NewUUID(), nil)
+	msg2 := message.NewMessage(watermill.NewUUID(), nil)
+	messagesCh <- msg1
+	messagesCh <- msg1
+	messagesCh <- msg2
+
+	readMessages, all := subscriber.BulkReadWithDeduplication(messagesCh, 2, time.Second)
+	assert.True(t, all)
+
+	assert.Equal(t, []string{msg1.UUID, msg2.UUID}, readMessages.IDs())
+}