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

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