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

message/decorator_bench_test.go

@@ -0,0 +1,89 @@
+package message_test
+
+import (
+	"context"
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/ThreeDotsLabs/watermill/message"
+)
+
+type benchSubscriber struct {
+	msgCh   chan *message.Message
+	closeCh chan struct{}
+	running sync.Mutex
+}
+
+func (b *benchSubscriber) Subscribe(ctx context.Context, topic string) (<-chan *message.Message, error) {
+	return b.msgCh, nil
+}
+
+func (b *benchSubscriber) Close() error {
+	close(b.closeCh)
+	b.running.Lock()
+	close(b.msgCh)
+	b.running.Unlock()
+	return nil
+}
+
+func newBenchSubscriber() *benchSubscriber {
+	sub := &benchSubscriber{
+		msgCh:   make(chan *message.Message, 1),
+		closeCh: make(chan struct{}),
+	}
+
+	// continually produce messages until Close()
+	go func() {
+		sub.running.Lock()
+		msg := message.NewMessage("", []byte{})
+		for {
+			select {
+			case <-sub.closeCh:
+				sub.running.Unlock()
+				return
+			case sub.msgCh <- msg:
+				// the buffer limit is 1, so this will block until someone consumes
+			}
+		}
+	}()
+
+	return sub
+}
+
+func BenchmarkMessageTransformSubscriberDecorator(b *testing.B) {
+	b.ReportAllocs()
+	b.Run("no_decorator", benchmarkNoDecorator)
+	b.Run("message_transform_decorator", benchmarkMessageTransformSubscriberDecorator)
+}
+
+func benchmarkNoDecorator(b *testing.B) {
+	sub := newBenchSubscriber()
+
+	in, err := sub.Subscribe(context.Background(), "")
+	require.NoError(b, err)
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		// consume one message
+		<-in
+	}
+}
+
+func benchmarkMessageTransformSubscriberDecorator(b *testing.B) {
+	sub := newBenchSubscriber()
+
+	noopDecorator := message.MessageTransformSubscriberDecorator(func(*message.Message) {})
+	decoratedSub, err := noopDecorator(sub)
+	require.NoError(b, err)
+
+	in, err := decoratedSub.Subscribe(context.Background(), "")
+	require.NoError(b, err)
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		// consume one message
+		<-in
+	}
+}