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

internal/channel.go

@@ -0,0 +1,15 @@
+package internal
+
+// IsChannelClosed returns true if provided `chan struct{}` is closed.
+// IsChannelClosed panics if message is sent to this channel.
+func IsChannelClosed(channel chan struct{}) bool {
+	select {
+	case _, ok := <-channel:
+		if ok {
+			panic("received unexpected message")
+		}
+		return true
+	default:
+		return false
+	}
+}

internal/channel_test.go

@@ -0,0 +1,57 @@
+package internal_test
+
+import (
+	"testing"
+
+	"github.com/ThreeDotsLabs/watermill/internal"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestIsChannelClosed(t *testing.T) {
+	closed := make(chan struct{})
+	close(closed)
+
+	withSentValue := make(chan struct{}, 1)
+	withSentValue <- struct{}{}
+
+	testCases := []struct {
+		Name           string
+		Channel        chan struct{}
+		ExpectedPanic  bool
+		ExpectedClosed bool
+	}{
+		{
+			Name:           "not_closed",
+			Channel:        make(chan struct{}),
+			ExpectedPanic:  false,
+			ExpectedClosed: false,
+		},
+		{
+			Name:           "closed",
+			Channel:        closed,
+			ExpectedPanic:  false,
+			ExpectedClosed: true,
+		},
+		{
+			Name:           "with_sent_value",
+			Channel:        withSentValue,
+			ExpectedPanic:  true,
+			ExpectedClosed: false,
+		},
+	}
+
+	for _, c := range testCases {
+		t.Run(c.Name, func(t *testing.T) {
+			testFunc := func() {
+				closed := internal.IsChannelClosed(c.Channel)
+				assert.EqualValues(t, c.ExpectedClosed, closed)
+			}
+
+			if c.ExpectedPanic {
+				assert.Panics(t, testFunc)
+			} else {
+				assert.NotPanics(t, testFunc)
+			}
+		})
+	}
+}

internal/name.go

@@ -0,0 +1,17 @@
+package internal
+
+import (
+	"fmt"
+	"strings"
+)
+
+// StructName returns a normalized name of the passed structure.
+func StructName(v interface{}) string {
+	if s, ok := v.(fmt.Stringer); ok {
+		return s.String()
+	}
+
+	s := fmt.Sprintf("%T", v)
+	// trim the pointer marker, if any
+	return strings.TrimLeft(s, "*")
+}

internal/name_test.go

@@ -0,0 +1,47 @@
+package internal_test
+
+import (
+	"testing"
+
+	"github.com/ThreeDotsLabs/watermill/internal"
+	"github.com/stretchr/testify/assert"
+)
+
+type testStruct struct{}
+
+type stringerStruct struct{}
+
+func (stringerStruct) String() string {
+	return "stringer"
+}
+
+func TestStructName(t *testing.T) {
+	testCases := []struct {
+		Name         string
+		Struct       interface{}
+		ExpectedName string
+	}{
+		{
+			Name:         "simple_struct",
+			Struct:       testStruct{},
+			ExpectedName: "internal_test.testStruct",
+		},
+		{
+			Name:         "pointer_struct",
+			Struct:       &testStruct{},
+			ExpectedName: "internal_test.testStruct",
+		},
+		{
+			Name:         "stringer",
+			Struct:       stringerStruct{},
+			ExpectedName: "stringer",
+		},
+	}
+
+	for _, c := range testCases {
+		t.Run(c.Name, func(t *testing.T) {
+			s := internal.StructName(c.Struct)
+			assert.Equal(t, c.ExpectedName, s)
+		})
+	}
+}

internal/norace.go

@@ -0,0 +1,6 @@
+//go:build !race
+// +build !race
+
+package internal
+
+const RaceEnabled = false

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

internal/race.go

@@ -0,0 +1,6 @@
+//go:build race
+// +build race
+
+package internal
+
+const RaceEnabled = true

internal/subscriber/multiplier.go

@@ -0,0 +1,83 @@
+package subscriber
+
+import (
+	"context"
+	stdErrors "errors"
+	"sync"
+
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/pkg/errors"
+)
+
+// Constructor is a function that creates a subscriber.
+type Constructor func() (message.Subscriber, error)
+
+type multiplier struct {
+	subscriberConstructor func() (message.Subscriber, error)
+	subscribersCount      int
+	subscribers           []message.Subscriber
+}
+
+// NewMultiplier returns multiplier subscriber decorator,
+// which under the hood is calling subscribe multiple times to increase throughput.
+func NewMultiplier(constructor Constructor, subscribersCount int) message.Subscriber {
+	return &multiplier{
+		subscriberConstructor: constructor,
+		subscribersCount:      subscribersCount,
+	}
+}
+
+func (s *multiplier) Subscribe(ctx context.Context, topic string) (msgs <-chan *message.Message, err error) {
+	defer func() {
+		if err != nil {
+			if closeErr := s.Close(); closeErr != nil {
+				err = stdErrors.Join(err, closeErr)
+			}
+		}
+	}()
+
+	out := make(chan *message.Message)
+
+	subWg := sync.WaitGroup{}
+	subWg.Add(s.subscribersCount)
+
+	for i := 0; i < s.subscribersCount; i++ {
+		sub, err := s.subscriberConstructor()
+		if err != nil {
+			return nil, errors.Wrap(err, "cannot create subscriber")
+		}
+
+		s.subscribers = append(s.subscribers, sub)
+
+		msgs, err := sub.Subscribe(ctx, topic)
+		if err != nil {
+			return nil, errors.Wrap(err, "cannot subscribe")
+		}
+
+		go func() {
+			for msg := range msgs {
+				out <- msg
+			}
+			subWg.Done()
+		}()
+	}
+
+	go func() {
+		subWg.Wait()
+		close(out)
+	}()
+
+	return out, nil
+}
+
+func (s *multiplier) Close() error {
+	var err error
+
+	for _, sub := range s.subscribers {
+		if closeErr := sub.Close(); closeErr != nil {
+			err = stdErrors.Join(err, closeErr)
+		}
+	}
+
+	return err
+}