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

components/delay/delay.go

@@ -0,0 +1,68 @@
+package delay
+
+import (
+	"context"
+	"time"
+
+	"github.com/ThreeDotsLabs/watermill/message"
+)
+
+// Delay represents a message's delay.
+// It can be either a delay until a specific time or a delay for a specific duration.
+// The zero value of Delay is a zero delay.
+//
+// IMPORTANT: Delay doesn't work 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 Delay struct {
+	time     time.Time
+	duration time.Duration
+}
+
+func (d Delay) IsZero() bool {
+	return d.time.IsZero()
+}
+
+// Until returns a delay of the given time.
+func Until(delayedUntil time.Time) Delay {
+	return Delay{
+		time:     delayedUntil,
+		duration: delayedUntil.Sub(time.Now().UTC()),
+	}
+}
+
+// For returns a delay of now plus the given duration.
+func For(delayedFor time.Duration) Delay {
+	return Delay{
+		time:     time.Now().UTC().Add(delayedFor),
+		duration: delayedFor,
+	}
+}
+
+type contextKey string
+
+var (
+	delayContextKey = contextKey("delay")
+)
+
+// WithContext returns a new context with the given delay.
+// If used together with a publisher wrapped with NewPublisher, the delay will be applied to the message.
+//
+// IMPORTANT: Delay doesn't work 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/
+func WithContext(ctx context.Context, delay Delay) context.Context {
+	return context.WithValue(ctx, delayContextKey, delay)
+}
+
+const (
+	DelayedUntilKey = "_watermill_delayed_until"
+	DelayedForKey   = "_watermill_delayed_for"
+)
+
+// Message sets the delay metadata on the message.
+//
+// IMPORTANT: Delay doesn't work 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/
+func Message(msg *message.Message, delay Delay) {
+	msg.Metadata.Set(DelayedUntilKey, delay.time.Format(time.RFC3339))
+	msg.Metadata.Set(DelayedForKey, delay.duration.String())
+}

components/delay/publisher.go

@@ -0,0 +1,83 @@
+package delay
+
+import (
+	"errors"
+
+	"github.com/ThreeDotsLabs/watermill/message"
+)
+
+type DefaultDelayGeneratorParams struct {
+	Topic   string
+	Message *message.Message
+}
+
+// PublisherConfig is a configuration for the delay publisher.
+type PublisherConfig struct {
+	// DefaultDelayGenerator is a function that generates the default delay for a message.
+	// If the message doesn't have the delay metadata set, the default delay will be applied.
+	DefaultDelayGenerator func(params DefaultDelayGeneratorParams) (Delay, error)
+
+	// AllowNoDelay allows publishing messages without a delay set.
+	// By default, the publisher returns an error when a message is published without a delay and no default delay generator is provided.
+	AllowNoDelay bool
+}
+
+// NewPublisher wraps a publisher with a delay mechanism.
+// A message can be published with delay metadata set in the context by using the WithContext function.
+// If the message doesn't have the delay metadata set, the default delay will be applied, if provided.
+func NewPublisher(pub message.Publisher, config PublisherConfig) (message.Publisher, error) {
+	return &publisher{
+		pub:    pub,
+		config: config,
+	}, nil
+}
+
+type publisher struct {
+	pub    message.Publisher
+	config PublisherConfig
+}
+
+func (p *publisher) Publish(topic string, messages ...*message.Message) error {
+	for i := range messages {
+		err := p.applyDelay(topic, messages[i])
+		if err != nil {
+			return err
+		}
+	}
+	return p.pub.Publish(topic, messages...)
+}
+
+func (p *publisher) Close() error {
+	return p.pub.Close()
+}
+
+func (p *publisher) applyDelay(topic string, msg *message.Message) error {
+	if msg.Metadata.Get(DelayedForKey) != "" {
+		return nil
+	}
+
+	if msg.Context().Value(delayContextKey) != nil {
+		delay := msg.Context().Value(delayContextKey).(Delay)
+		Message(msg, delay)
+		return nil
+	}
+
+	if p.config.DefaultDelayGenerator != nil {
+		delay, err := p.config.DefaultDelayGenerator(DefaultDelayGeneratorParams{
+			Topic:   topic,
+			Message: msg,
+		})
+		if err != nil {
+			return err
+		}
+		Message(msg, delay)
+
+		return nil
+	}
+
+	if !p.config.AllowNoDelay {
+		return errors.New("message doesn't have a delay set")
+	}
+
+	return nil
+}

components/delay/publisher_test.go

@@ -0,0 +1,178 @@
+package delay_test
+
+import (
+	"context"
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/ThreeDotsLabs/watermill/components/delay"
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/ThreeDotsLabs/watermill/pubsub/gochannel"
+)
+
+func TestPublisher(t *testing.T) {
+	pubSub := gochannel.NewGoChannel(gochannel.Config{}, nil)
+
+	messages, err := pubSub.Subscribe(context.Background(), "test")
+	require.NoError(t, err)
+
+	pub, err := delay.NewPublisher(pubSub, delay.PublisherConfig{})
+	require.NoError(t, err)
+
+	pubAllowNoDelay, err := delay.NewPublisher(pubSub, delay.PublisherConfig{
+		AllowNoDelay: true,
+	})
+	require.NoError(t, err)
+
+	defaultDelayPub, err := delay.NewPublisher(pubSub, delay.PublisherConfig{
+		DefaultDelayGenerator: func(params delay.DefaultDelayGeneratorParams) (delay.Delay, error) {
+			return delay.For(1 * time.Second), nil
+		},
+	})
+	require.NoError(t, err)
+
+	testCases := []struct {
+		name               string
+		publisher          message.Publisher
+		messageConstructor func(id string) *message.Message
+		expectedError      bool
+		expectedDelay      time.Duration
+	}{
+		{
+			name:      "no delay",
+			publisher: pub,
+			messageConstructor: func(id string) *message.Message {
+				return message.NewMessage(id, nil)
+			},
+			expectedError: true,
+			expectedDelay: 0,
+		},
+		{
+			name:      "no delay but allowed",
+			publisher: pubAllowNoDelay,
+			messageConstructor: func(id string) *message.Message {
+				return message.NewMessage(id, nil)
+			},
+			expectedDelay: 0,
+		},
+		{
+			name:      "default delay",
+			publisher: defaultDelayPub,
+			messageConstructor: func(id string) *message.Message {
+				return message.NewMessage(id, nil)
+			},
+			expectedDelay: 1 * time.Second,
+		},
+		{
+			name:      "delay from metadata",
+			publisher: pub,
+			messageConstructor: func(id string) *message.Message {
+				msg := message.NewMessage(id, nil)
+				delay.Message(msg, delay.For(2*time.Second))
+				return msg
+			},
+			expectedDelay: 2 * time.Second,
+		},
+		{
+			name:      "default delay override with metadata",
+			publisher: defaultDelayPub,
+			messageConstructor: func(id string) *message.Message {
+				msg := message.NewMessage(id, nil)
+				delay.Message(msg, delay.For(2*time.Second))
+				return msg
+			},
+			expectedDelay: 2 * time.Second,
+		},
+		{
+			name:      "delay from context",
+			publisher: pub,
+			messageConstructor: func(id string) *message.Message {
+				msg := message.NewMessage(id, nil)
+				ctx := delay.WithContext(context.Background(), delay.For(3*time.Second))
+				msg.SetContext(ctx)
+				return msg
+			},
+			expectedDelay: 3 * time.Second,
+		},
+		{
+			name:      "default delay override with context",
+			publisher: defaultDelayPub,
+			messageConstructor: func(id string) *message.Message {
+				msg := message.NewMessage(id, nil)
+				ctx := delay.WithContext(context.Background(), delay.For(3*time.Second))
+				msg.SetContext(ctx)
+				return msg
+			},
+			expectedDelay: 3 * time.Second,
+		},
+		{
+			name:      "delay with until",
+			publisher: pub,
+			messageConstructor: func(id string) *message.Message {
+				msg := message.NewMessage(id, nil)
+				delay.Message(msg, delay.Until(time.Now().UTC().Add(4*time.Second)))
+				return msg
+			},
+			expectedDelay: 4 * time.Second,
+		},
+		{
+			name:      "both metadata and context set",
+			publisher: defaultDelayPub,
+			messageConstructor: func(id string) *message.Message {
+				msg := message.NewMessage(id, nil)
+				delay.Message(msg, delay.For(5*time.Second))
+				ctx := delay.WithContext(context.Background(), delay.For(6*time.Second))
+				msg.SetContext(ctx)
+				return msg
+			},
+			expectedDelay: 5 * time.Second,
+		},
+	}
+
+	for i, testCase := range testCases {
+		t.Run(testCase.name, func(t *testing.T) {
+			id := fmt.Sprint(i)
+
+			msg := testCase.messageConstructor(id)
+			err = testCase.publisher.Publish("test", msg)
+
+			if testCase.expectedError {
+				require.Error(t, err)
+				return
+			}
+
+			require.NoError(t, err)
+			assertMessage(t, messages, id, testCase.expectedDelay)
+		})
+	}
+}
+
+func assertMessage(t *testing.T, messages <-chan *message.Message, expectedID string, expectedDelay time.Duration) {
+	t.Helper()
+	select {
+	case msg := <-messages:
+		assert.Equal(t, expectedID, msg.UUID)
+
+		if expectedDelay == 0 {
+			assert.Empty(t, msg.Metadata.Get(delay.DelayedUntilKey))
+			assert.Empty(t, msg.Metadata.Get(delay.DelayedForKey))
+		} else {
+			delayedFor, err := time.ParseDuration(msg.Metadata.Get(delay.DelayedForKey))
+			require.NoError(t, err)
+			assert.Equal(t, expectedDelay, delayedFor.Round(time.Second))
+
+			delayedUntil, err := time.Parse(time.RFC3339, msg.Metadata.Get(delay.DelayedUntilKey))
+			require.NoError(t, err)
+
+			assert.WithinDuration(t, time.Now().UTC().Add(expectedDelay), delayedUntil, 1*time.Second)
+		}
+
+		msg.Ack()
+	case <-time.After(100 * time.Millisecond):
+		require.Fail(t, "timeout")
+	}
+}