Requeuer and Message Delay

_examples/pubsubs/googlecloud/go.sum

@@ -8,7 +8,9 @@ cloud.google.com/go/compute/metadata v0.2.3/go.mod h1:VAV5nSsACxMJvgaAuX6Pk2Aawl
 cloud.google.com/go/iam v0.10.0 h1:fpP/gByFs6US1ma53v7VxhvbJpO2Aapng6wabJ99MuI=
 cloud.google.com/go/iam v0.10.0/go.mod h1:nXAECrMt2qHpF6RZUZseteD6QyanL68reN4OXPw0UWM=
 cloud.google.com/go/kms v1.6.0 h1:OWRZzrPmOZUzurjI2FBGtgY2mB1WaJkqhw6oIwSj0Yg=
+cloud.google.com/go/kms v1.6.0/go.mod h1:Jjy850yySiasBUDi6KFUwUv2n1+o7QZFyuUJg6OgjA0=
 cloud.google.com/go/longrunning v0.3.0 h1:NjljC+FYPV3uh5/OwWT6pVU+doBqMg2x/rZlE+CamDs=
+cloud.google.com/go/longrunning v0.3.0/go.mod h1:qth9Y41RRSUE69rDcOn6DdK3HfQfsUI0YSmW3iIlLJc=
 cloud.google.com/go/pubsub v1.28.0 h1:XzabfdPx/+eNrsVVGLFgeUnQQKPGkMb8klRCeYK52is=
 cloud.google.com/go/pubsub v1.28.0/go.mod h1:vuXFpwaVoIPQMGXqRyUQigu/AX1S3IWugR9xznmcXX8=
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
@@ -137,7 +139,7 @@ google.golang.org/genproto v0.0.0-20190819201941-24fa4b261c55/go.mod h1:DMBHOl98
 google.golang.org/genproto v0.0.0-20200526211855-cb27e3aa2013/go.mod h1:NbSheEEYHJ7i3ixzK3sjbqSGDJWnxyFXZblF3eUsNvo=
 google.golang.org/genproto v0.0.0-20230202175211-008b39050e57 h1:vArvWooPH749rNHpBGgVl+U9B9dATjiEhJzcWGlovNs=
 google.golang.org/genproto v0.0.0-20230202175211-008b39050e57/go.mod h1:RGgjbofJ8xD9Sq1VVhDM1Vok1vRONV+rg+CjzG4SZKM=
-google.golang.org/grpc v1.21.0/go.mod h1:mqu4LbDTu4XGKhr4mRzUsmM4RtVoemTSY81AxZiDr8c=
+google.golang.org/grpc v1.19.0/go.mod h1:mqu4LbDTu4XGKhr4mRzUsmM4RtVoemTSY81AxZiDr8c=
 google.golang.org/grpc v1.23.0/go.mod h1:Y5yQAOtifL1yxbo5wqy6BxZv8vAUGQwXBOALyacEbxg=
 google.golang.org/grpc v1.25.1/go.mod h1:c3i+UQWmh7LiEpx4sFZnkU36qjEYZ0imhYfXVyQciAY=
 google.golang.org/grpc v1.27.0/go.mod h1:qbnxyOmOxrQa7FizSgH+ReBfzJrCY1pSN7KXBS8abTk=

components/requeue/requeue.go

@@ -0,0 +1,64 @@
+package requeue
+
+import (
+	"context"
+	"time"
+
+	"github.com/ThreeDotsLabs/watermill/message/router/middleware"
+	"github.com/pkg/errors"
+
+	"github.com/ThreeDotsLabs/watermill"
+	"github.com/ThreeDotsLabs/watermill/message"
+)
+
+const RequeueDelayKey = "requeue_delay"
+
+type Requeue struct {
+	router *message.Router
+}
+
+type Config struct {
+	Subscriber message.Subscriber
+	Topic      string
+	Publisher  message.Publisher
+	Delay      time.Duration
+}
+
+func NewRequeue(
+	config Config,
+	logger watermill.LoggerAdapter,
+) (*Requeue, error) {
+	router, err := message.NewRouter(message.RouterConfig{}, logger)
+	if err != nil {
+		return nil, err
+	}
+
+	router.AddNoPublisherHandler(
+		"requeue",
+		config.Topic,
+		config.Subscriber,
+		func(msg *message.Message) error {
+			time.Sleep(config.Delay)
+
+			topic := msg.Metadata.Get(middleware.PoisonedTopicKey)
+			if topic == "" {
+				return errors.New("missing requeue topic")
+			}
+
+			err := config.Publisher.Publish(topic, msg)
+			if err != nil {
+				return err
+			}
+
+			return nil
+		},
+	)
+
+	return &Requeue{
+		router: router,
+	}, nil
+}
+
+func (r *Requeue) Run(ctx context.Context) error {
+	return r.router.Run(ctx)
+}

components/requeue/requeue_test.go

@@ -0,0 +1,97 @@
+package requeue
+
+import (
+	"context"
+	"fmt"
+	"strconv"
+	"testing"
+	"time"
+
+	"github.com/ThreeDotsLabs/watermill"
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/ThreeDotsLabs/watermill/message/router/middleware"
+	"github.com/ThreeDotsLabs/watermill/pubsub/gochannel"
+	"github.com/pkg/errors"
+	"github.com/stretchr/testify/require"
+)
+
+func TestRequeue(t *testing.T) {
+	logger := watermill.NewStdLogger(false, false)
+
+	pubSub := gochannel.NewGoChannel(gochannel.Config{}, logger)
+
+	requeue, err := NewRequeue(Config{
+		Subscriber: pubSub,
+		Topic:      "requeue",
+		Publisher:  pubSub,
+		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)
+
+	counter := 0
+
+	router.AddNoPublisherHandler(
+		"test",
+		"test",
+		pubSub,
+		func(msg *message.Message) error {
+			i, err := strconv.Atoi(string(msg.Payload))
+			if err != nil {
+				return err
+			}
+
+			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(), []byte(fmt.Sprint(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
+			break
+		}
+	}
+
+	require.ElementsMatch(t, []int{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}, received)
+}