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

internal/publisher/errors.go

@@ -0,0 +1,39 @@
+package publisher
+
+import (
+	"strings"
+
+	"github.com/ThreeDotsLabs/watermill/message"
+)
+
+type ErrCouldNotPublish struct {
+	reasons map[string]error
+}
+
+func (e *ErrCouldNotPublish) addMsg(msg *message.Message, reason error) {
+	e.reasons[msg.UUID] = reason
+}
+
+func NewErrCouldNotPublish() *ErrCouldNotPublish {
+	return &ErrCouldNotPublish{make(map[string]error)}
+}
+
+func (e ErrCouldNotPublish) Len() int {
+	return len(e.reasons)
+}
+
+func (e ErrCouldNotPublish) Error() string {
+	if len(e.reasons) == 0 {
+		return ""
+	}
+	b := strings.Builder{}
+	b.WriteString("Could not publish the messages:\n")
+	for uuid, reason := range e.reasons {
+		b.WriteString(uuid + " : " + reason.Error() + "\n")
+	}
+	return b.String()
+}
+
+func (e ErrCouldNotPublish) Reasons() map[string]error {
+	return e.reasons
+}

internal/publisher/retry.go

@@ -0,0 +1,108 @@
+package publisher
+
+import (
+	"time"
+
+	"github.com/ThreeDotsLabs/watermill"
+	"github.com/ThreeDotsLabs/watermill/message"
+
+	"github.com/pkg/errors"
+)
+
+var (
+	ErrNonPositiveNumberOfRetries  = errors.New("number of retries should be positive")
+	ErrNonPositiveTimeToFirstRetry = errors.New("time to first retry should be positive")
+)
+
+type RetryPublisherConfig struct {
+	MaxRetries int
+	// each subsequent retry doubles the time to next retry.
+	TimeToFirstRetry time.Duration
+	Logger           watermill.LoggerAdapter
+}
+
+func (c *RetryPublisherConfig) setDefaults() {
+	if c.MaxRetries == 0 {
+		c.MaxRetries = 5
+	}
+
+	if c.TimeToFirstRetry == 0 {
+		c.TimeToFirstRetry = time.Second
+	}
+
+	if c.Logger == nil {
+		c.Logger = watermill.NopLogger{}
+	}
+}
+
+func (c RetryPublisherConfig) validate() error {
+	if c.MaxRetries <= 0 {
+		return ErrNonPositiveNumberOfRetries
+	}
+	if c.TimeToFirstRetry <= 0 {
+		return ErrNonPositiveTimeToFirstRetry
+	}
+
+	return nil
+}
+
+// RetryPublisher is a decorator for a publisher that retries message publishing after a failure.
+type RetryPublisher struct {
+	pub    message.Publisher
+	config RetryPublisherConfig
+}
+
+func NewRetryPublisher(pub message.Publisher, config RetryPublisherConfig) (*RetryPublisher, error) {
+	config.setDefaults()
+
+	if err := config.validate(); err != nil {
+		return nil, errors.Wrap(err, "invalid RetryPublisher config")
+	}
+
+	return &RetryPublisher{
+		pub,
+		config,
+	}, nil
+}
+
+func (p RetryPublisher) Publish(topic string, messages ...*message.Message) error {
+	failedMessages := NewErrCouldNotPublish()
+
+	// todo: do some parallel processing maybe? this is a very basic implementation
+	for _, msg := range messages {
+		err := p.send(topic, msg)
+		if err != nil {
+			failedMessages.addMsg(msg, err)
+		}
+	}
+
+	if failedMessages.Len() > 0 {
+		return failedMessages
+	}
+
+	return nil
+}
+
+func (p RetryPublisher) Close() error {
+	return p.pub.Close()
+}
+
+// send sends one message at a time to prevent sending a successful message more than once.
+func (p RetryPublisher) send(topic string, msg *message.Message) error {
+	var err error
+	timeToNextRetry := p.config.TimeToFirstRetry
+
+	for i := 0; i < p.config.MaxRetries; i++ {
+		err = p.pub.Publish(topic, msg)
+		if err == nil {
+			return nil
+		}
+
+		p.config.Logger.Info("Publish failed, retrying in "+timeToNextRetry.String(), watermill.LogFields{
+			"error": err,
+		})
+		time.Sleep(timeToNextRetry)
+		timeToNextRetry *= 2
+	}
+	return err
+}

internal/publisher/retry_test.go

@@ -0,0 +1,197 @@
+package publisher_test
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/ThreeDotsLabs/watermill"
+	"github.com/ThreeDotsLabs/watermill/internal/publisher"
+	"github.com/stretchr/testify/require"
+
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/pkg/errors"
+)
+
+var errCouldNotPublish = errors.New("could not publish, try again")
+
+// FailingPublisher mocks a publisher that fails a specific number of time for a message, then succeeds.
+type FailingPublisher struct {
+	howManyFails     map[string]int
+	howManyPublished map[string]int
+}
+
+func (p *FailingPublisher) Publish(topic string, messages ...*message.Message) error {
+	for _, msg := range messages {
+		howManyFails, ok := p.howManyFails[msg.UUID]
+		if !ok || howManyFails <= 0 {
+			p.publish(msg)
+			continue
+		}
+		p.howManyFails[msg.UUID]--
+		return errCouldNotPublish
+	}
+
+	return nil
+}
+
+func (p *FailingPublisher) publish(msg *message.Message) {
+	if _, ok := p.howManyPublished[msg.UUID]; !ok {
+		p.howManyPublished[msg.UUID] = 1
+		return
+	}
+	p.howManyPublished[msg.UUID]++
+}
+
+func (p *FailingPublisher) Close() error {
+	return nil
+}
+
+func TestRetryPublisher_Publish_after_retries(t *testing.T) {
+	msg := message.NewMessage("uuid", []byte{})
+	pub := FailingPublisher{
+		howManyFails: map[string]int{
+			msg.UUID: 4,
+		},
+		howManyPublished: map[string]int{},
+	}
+	conf := publisher.RetryPublisherConfig{
+		MaxRetries:       5,
+		TimeToFirstRetry: time.Millisecond,
+		Logger:           watermill.NopLogger{},
+	}
+	retryPub, err := publisher.NewRetryPublisher(&pub, conf)
+	require.NoError(t, err)
+
+	// given
+	require.True(t, pub.howManyFails[msg.UUID] < conf.MaxRetries, "Publisher must fail less than MaxRetries times")
+
+	// when
+	err = retryPub.Publish("topic", msg)
+
+	// then
+	require.NoError(t, err)
+	assert.Contains(t, pub.howManyPublished, msg.UUID)
+	assert.Equal(t, pub.howManyPublished[msg.UUID], 1, "Expected msg to be published exactly once")
+}
+
+func TestRetryPublisher_Publish_too_many_retries(t *testing.T) {
+	msg := message.NewMessage("uuid", []byte{})
+	pub := FailingPublisher{
+		howManyFails: map[string]int{
+			msg.UUID: 5,
+		},
+		howManyPublished: map[string]int{},
+	}
+	conf := publisher.RetryPublisherConfig{
+		MaxRetries:       5,
+		TimeToFirstRetry: time.Millisecond,
+		Logger:           watermill.NopLogger{},
+	}
+	retryPub, err := publisher.NewRetryPublisher(&pub, conf)
+	require.NoError(t, err)
+
+	// given
+	require.True(t, pub.howManyFails[msg.UUID] >= conf.MaxRetries, "Publisher must fail at least MaxRetries times")
+
+	// when
+	err = retryPub.Publish("topic", msg)
+
+	// then
+	require.Error(t, err)
+
+	cnpErr, ok := err.(*publisher.ErrCouldNotPublish)
+	require.True(t, ok, "expected the ErrCouldNotPublish composite error type")
+
+	assert.Equal(t, 1, cnpErr.Len(), "attempted to publish one message, expecting one error")
+	assert.Equal(t, errCouldNotPublish, errors.Cause(cnpErr.Reasons()[msg.UUID]))
+
+	assert.NotContains(t, pub.howManyPublished, msg.UUID, "expected msg to not be published")
+}
+
+func TestPublishEachMessageOnlyOnce(t *testing.T) {
+	msg1 := message.NewMessage("uuid1", []byte{})
+	msg2 := message.NewMessage("uuid2", []byte{})
+
+	pub := FailingPublisher{
+		howManyFails: map[string]int{
+			msg1.UUID: 2,
+			msg2.UUID: 4,
+		},
+		howManyPublished: map[string]int{},
+	}
+	conf := publisher.RetryPublisherConfig{
+		MaxRetries:       5,
+		TimeToFirstRetry: time.Millisecond,
+		Logger:           watermill.NopLogger{},
+	}
+	retryPub, err := publisher.NewRetryPublisher(&pub, conf)
+	require.NoError(t, err)
+
+	// given
+	require.True(t, pub.howManyFails[msg1.UUID] < conf.MaxRetries, "Publisher must fail less than MaxRetries times for msg1")
+	require.True(t, pub.howManyFails[msg2.UUID] < conf.MaxRetries, "Publisher must fail less than MaxRetries times for msg2")
+	require.True(t, pub.howManyFails[msg1.UUID] < pub.howManyFails[msg2.UUID], "Publisher must fail less times for msg1 than msg2")
+
+	// when
+	err = retryPub.Publish("topic", msg1, msg2)
+
+	// then
+	require.NoError(t, err)
+
+	assert.Equal(t, pub.howManyPublished[msg1.UUID], 1, "expected msg1 to be published only once")
+	assert.Equal(t, pub.howManyPublished[msg2.UUID], 1, "expected msg2 to be published only once")
+}
+
+var ErrCouldNotClose = errors.New("this publisher fails on Close()")
+
+// ClosingPublisher mocks a publisher that may fail on closing so we may check if the Close() call propagated correctly.
+type ClosingPublisher struct {
+	closed      bool
+	failOnClose bool
+}
+
+func (ClosingPublisher) Publish(topic string, messages ...*message.Message) error {
+	return nil
+}
+
+func (p *ClosingPublisher) Close() error {
+	if p.failOnClose {
+		return ErrCouldNotClose
+	}
+	p.closed = true
+	return nil
+}
+
+func TestRetryPublisher_Close(t *testing.T) {
+	pub := ClosingPublisher{}
+	retryPub, err := publisher.NewRetryPublisher(&pub, publisher.RetryPublisherConfig{})
+	require.NoError(t, err)
+
+	// given
+	require.False(t, pub.closed)
+
+	// when
+	err = retryPub.Close()
+
+	// then
+	require.NoError(t, err)
+	assert.True(t, pub.closed)
+}
+
+func TestRetryPublisher_Close_failed(t *testing.T) {
+	pub := ClosingPublisher{failOnClose: true}
+	retryPub, err := publisher.NewRetryPublisher(&pub, publisher.RetryPublisherConfig{})
+	require.NoError(t, err)
+
+	// given
+	require.False(t, pub.closed)
+
+	// when
+	err = retryPub.Close()
+
+	// then
+	require.Error(t, err)
+	assert.Equal(t, ErrCouldNotClose, errors.Cause(err))
+}