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

components/requeuer/requeuer.go

@@ -0,0 +1,158 @@
+package requeuer
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"strconv"
+	"time"
+
+	"github.com/ThreeDotsLabs/watermill"
+	"github.com/ThreeDotsLabs/watermill/message"
+)
+
+const RetriesKey = "_watermill_requeuer_retries"
+
+// Requeuer is a component that moves messages from one topic to another.
+// It can be used to requeue messages that failed to process.
+type Requeuer struct {
+	config Config
+}
+
+// GeneratePublishTopicParams are the parameters passed to the GeneratePublishTopic function.
+type GeneratePublishTopicParams struct {
+	Message *message.Message
+}
+
+// Config is the configuration for the Requeuer.
+type Config struct {
+	// Subscriber is the subscriber to consume messages from. Required.
+	Subscriber message.Subscriber
+
+	// SubscribeTopic is the topic related to the Subscriber to consume messages from. Required.
+	SubscribeTopic string
+
+	// Publisher is the publisher to publish requeued messages to. Required.
+	Publisher message.Publisher
+
+	// GeneratePublishTopic is the topic related to the Publisher to publish the requeued message to.
+	// For example, it could be a constant, or taken from the message's metadata.
+	// Required.
+	GeneratePublishTopic func(params GeneratePublishTopicParams) (string, error)
+
+	// Delay is the duration to wait before requeuing the message. Optional.
+	// The default is no delay.
+	//
+	// This can be useful to avoid requeuing messages too quickly, for example, to avoid
+	// requeuing a message that failed to process due to a temporary issue.
+	//
+	// Avoid setting this to a very high value, as it will block the message processing.
+	Delay time.Duration
+
+	// Router is the custom router to run the requeue handler on. Optional.
+	Router *message.Router
+}
+
+func (c *Config) setDefaults(logger watermill.LoggerAdapter) error {
+	if c.Router == nil {
+		router, err := message.NewRouter(message.RouterConfig{}, logger)
+		if err != nil {
+			return fmt.Errorf("could not create router: %w", err)
+		}
+
+		c.Router = router
+	}
+
+	return nil
+}
+
+func (c *Config) validate() error {
+	if c.Subscriber == nil {
+		return errors.New("subscriber is required")
+	}
+
+	if c.SubscribeTopic == "" {
+		return errors.New("subscribe topic is required")
+	}
+
+	if c.Publisher == nil {
+		return errors.New("publisher is required")
+	}
+
+	if c.GeneratePublishTopic == nil {
+		return errors.New("generate publish topic is required")
+	}
+
+	return nil
+}
+
+// NewRequeuer creates a new Requeuer with the provided Config.
+// It's not started automatically. You need to call Run on the returned Requeuer.
+func NewRequeuer(
+	config Config,
+	logger watermill.LoggerAdapter,
+) (*Requeuer, error) {
+	if logger == nil {
+		logger = watermill.NewStdLogger(false, false)
+	}
+
+	err := config.setDefaults(logger)
+	if err != nil {
+		return nil, err
+	}
+
+	err = config.validate()
+	if err != nil {
+		return nil, fmt.Errorf("invalid config: %w", err)
+	}
+
+	r := &Requeuer{
+		config: config,
+	}
+
+	config.Router.AddConsumerHandler(
+		"requeuer",
+		config.SubscribeTopic,
+		config.Subscriber,
+		r.handler,
+	)
+
+	return r, nil
+}
+
+func (r *Requeuer) handler(msg *message.Message) error {
+	if r.config.Delay > 0 {
+		select {
+		case <-msg.Context().Done():
+			return msg.Context().Err()
+		case <-time.After(r.config.Delay):
+		}
+	}
+
+	topic, err := r.config.GeneratePublishTopic(GeneratePublishTopicParams{Message: msg})
+	if err != nil {
+		return err
+	}
+
+	retriesStr := msg.Metadata.Get(RetriesKey)
+	retries, err := strconv.Atoi(retriesStr)
+	if err != nil {
+		retries = 0
+	}
+
+	retries++
+
+	msg.Metadata.Set(RetriesKey, strconv.Itoa(retries))
+
+	err = r.config.Publisher.Publish(topic, msg)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// Run runs the Requeuer.
+func (r *Requeuer) Run(ctx context.Context) error {
+	return r.config.Router.Run(ctx)
+}

components/requeuer/requeuer_test.go

@@ -0,0 +1,105 @@
+package requeuer_test
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"strconv"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/ThreeDotsLabs/watermill"
+	"github.com/ThreeDotsLabs/watermill/components/requeuer"
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/ThreeDotsLabs/watermill/message/router/middleware"
+	"github.com/ThreeDotsLabs/watermill/pubsub/gochannel"
+)
+
+func TestRequeue(t *testing.T) {
+	logger := watermill.NewStdLogger(false, false)
+
+	pubSub := gochannel.NewGoChannel(gochannel.Config{}, logger)
+
+	requeue, err := requeuer.NewRequeuer(requeuer.Config{
+		Subscriber:     pubSub,
+		SubscribeTopic: "requeue",
+		Publisher:      pubSub,
+		GeneratePublishTopic: func(params requeuer.GeneratePublishTopicParams) (string, error) {
+			return "test", nil
+		},
+		Delay: time.Millisecond * 200,
+	}, logger)
+	require.NoError(t, err)
+
+	go func() {
+		err := requeue.Run(context.Background())
+		require.NoError(t, err)
+	}()
+
+	router, err := message.NewRouter(message.RouterConfig{}, logger)
+	require.NoError(t, err)
+
+	pq, err := middleware.PoisonQueue(pubSub, "requeue")
+	require.NoError(t, err)
+
+	router.AddMiddleware(pq)
+
+	receivedMessages := make(chan int, 10)
+
+	lock := sync.Mutex{}
+	counter := 0
+
+	router.AddConsumerHandler(
+		"test",
+		"test",
+		pubSub,
+		func(msg *message.Message) error {
+			i, err := strconv.Atoi(string(msg.Payload))
+			if err != nil {
+				return err
+			}
+
+			lock.Lock()
+			defer lock.Unlock()
+			counter++
+
+			if counter < 10 && i%2 == 0 {
+				return errors.New("error")
+			}
+
+			receivedMessages <- i
+
+			return nil
+		},
+	)
+
+	go func() {
+		err := router.Run(context.Background())
+		require.NoError(t, err)
+	}()
+
+	time.Sleep(time.Second)
+
+	for i := 0; i < 10; i++ {
+		msg := message.NewMessage(watermill.NewUUID(), fmt.Append(nil, i))
+		err := pubSub.Publish("test", msg)
+		require.NoError(t, err)
+	}
+
+	var received []int
+
+	timeout := false
+	for !timeout {
+		select {
+		case i := <-receivedMessages:
+			received = append(received, i)
+		case <-time.After(5 * time.Second):
+			timeout = true
+		}
+	}
+
+	require.ElementsMatch(t, []int{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}, received)
+}