84 lines
1.7 KiB
Go
84 lines
1.7 KiB
Go
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
|
|
}
|