Fix #275 and add gochannel concurreny control

pubsub/gochannel/pubsub.go

@@ -2,10 +2,9 @@ package gochannel
 
 import (
 	"context"
-	"sync"
-
 	"github.com/lithammer/shortuuid/v3"
 	"github.com/pkg/errors"
+	"sync"
 
 	"github.com/ThreeDotsLabs/watermill"
 	"github.com/ThreeDotsLabs/watermill/message"
@@ -26,6 +25,10 @@ type Config struct {
 	// When true, Publish will block until subscriber Ack's the message.
 	// If there are no subscribers, Publish will not block (also when Persistent is true).
 	BlockPublishUntilSubscriberAck bool
+
+	// The concurrent count of subscriber to handle the messages.
+	// Default value is 1, which means the subscriber will be blocked until the last message ack or nack.
+	SubscriberConcurrentCount int
 }
 
 // GoChannel is the simplest Pub/Sub implementation.
@@ -60,7 +63,9 @@ func NewGoChannel(config Config, logger watermill.LoggerAdapter) *GoChannel {
 	if logger == nil {
 		logger = watermill.NopLogger{}
 	}
-
+	if config.SubscriberConcurrentCount == 0 {
+		config.SubscriberConcurrentCount = 1
+	}
 	return &GoChannel{
 		config: config,
 
@@ -69,7 +74,6 @@ func NewGoChannel(config Config, logger watermill.LoggerAdapter) *GoChannel {
 		logger: logger.With(watermill.LogFields{
 			"pubsub_uuid": shortuuid.New(),
 		}),
-
 		closing: make(chan struct{}),
 
 		persistedMessages: map[string][]*message.Message{},
@@ -78,7 +82,6 @@ func NewGoChannel(config Config, logger watermill.LoggerAdapter) *GoChannel {
 
 // Publish in GoChannel is NOT blocking until all consumers consume.
 // Messages will be send in background.
-//
 // Messages may be persisted or not, depending of persistent attribute.
 func (g *GoChannel) Publish(topic string, messages ...*message.Message) error {
 	if g.isClosed() {
@@ -113,7 +116,6 @@ func (g *GoChannel) Publish(topic string, messages ...*message.Message) error {
 		if err != nil {
 			return err
 		}
-
 		if g.config.BlockPublishUntilSubscriberAck {
 			g.waitForAckFromSubscribers(msg, ackedBySubscribers)
 		}
@@ -186,11 +188,12 @@ func (g *GoChannel) Subscribe(ctx context.Context, topic string) (<-chan *messag
 	subLock.(*sync.Mutex).Lock()
 
 	s := &subscriber{
-		ctx:           ctx,
-		uuid:          watermill.NewUUID(),
-		outputChannel: make(chan *message.Message, g.config.OutputChannelBuffer),
-		logger:        g.logger,
-		closing:       make(chan struct{}),
+		ctx:            ctx,
+		uuid:           watermill.NewUUID(),
+		outputChannel:  make(chan *message.Message, g.config.OutputChannelBuffer),
+		sendingChannel: make(chan struct{}, g.config.SubscriberConcurrentCount),
+		logger:         g.logger,
+		closing:        make(chan struct{}),
 	}
 
 	go func(s *subscriber, g *GoChannel) {
@@ -315,8 +318,9 @@ type subscriber struct {
 
 	uuid string
 
-	sending       sync.Mutex
-	outputChannel chan *message.Message
+	sending        sync.Mutex
+	sendingChannel chan struct{}
+	outputChannel  chan *message.Message
 
 	logger  watermill.LoggerAdapter
 	closed  bool
@@ -342,12 +346,13 @@ func (s *subscriber) Close() {
 }
 
 func (s *subscriber) sendMessageToSubscriber(msg *message.Message, logFields watermill.LogFields) {
-	s.sending.Lock()
-	defer s.sending.Unlock()
+	s.sendingChannel <- struct{}{}
+	defer func() {
+		<-s.sendingChannel
+	}()
 
 	ctx, cancelCtx := context.WithCancel(s.ctx)
 	defer cancelCtx()
-
 SendToSubscriber:
 	for {
 		// copy the message to prevent ack/nack propagation to other consumers