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

pubsub/gochannel/pubsub.go

@@ -0,0 +1,420 @@
+package gochannel
+
+import (
+	"context"
+	"sync"
+
+	"github.com/lithammer/shortuuid/v3"
+	"github.com/pkg/errors"
+
+	"github.com/ThreeDotsLabs/watermill"
+	"github.com/ThreeDotsLabs/watermill/message"
+)
+
+// Config holds the GoChannel Pub/Sub's configuration options.
+type Config struct {
+	// Output channel buffer size.
+	OutputChannelBuffer int64
+
+	// If persistent is set to true, when subscriber subscribes to the topic,
+	// it will receive all previously produced messages.
+	//
+	// All messages are persisted to the memory (simple slice),
+	// so be aware that with large amount of messages you can go out of the memory.
+	Persistent bool
+
+	// 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
+
+	// PreserveContext is a flag that determines if the context should be preserved when sending messages to subscribers.
+	// This behavior is different from other implementations of Publishers where data travels over the network,
+	// hence context can't be preserved in those cases
+	PreserveContext bool
+}
+
+// GoChannel is the simplest Pub/Sub implementation.
+// It is based on Golang's channels which are sent within the process.
+//
+// GoChannel has no global state,
+// that means that you need to use the same instance for Publishing and Subscribing!
+//
+// When GoChannel is persistent, messages order is not guaranteed.
+type GoChannel struct {
+	config Config
+	logger watermill.LoggerAdapter
+
+	subscribersWg          sync.WaitGroup
+	subscribers            map[string][]*subscriber
+	subscribersLock        sync.RWMutex
+	subscribersByTopicLock sync.Map // map of *sync.Mutex
+
+	closed     bool
+	closedLock sync.Mutex
+	closing    chan struct{}
+
+	persistedMessages     map[string][]*message.Message
+	persistedMessagesLock sync.RWMutex
+}
+
+// NewGoChannel creates new GoChannel Pub/Sub.
+//
+// This GoChannel is not persistent.
+// That means if you send a message to a topic to which no subscriber is subscribed, that message will be discarded.
+func NewGoChannel(config Config, logger watermill.LoggerAdapter) *GoChannel {
+	if logger == nil {
+		logger = watermill.NopLogger{}
+	}
+
+	return &GoChannel{
+		config: config,
+
+		subscribers:            make(map[string][]*subscriber),
+		subscribersByTopicLock: sync.Map{},
+		logger: logger.With(
+			watermill.LogFields{
+				"pubsub_uuid": shortuuid.New(),
+			},
+		),
+
+		closing: make(chan struct{}),
+
+		persistedMessages: map[string][]*message.Message{},
+	}
+}
+
+// Publish in GoChannel is NOT blocking until all consumers consume.
+// Messages will be sent in background.
+//
+// Messages may be persisted or not, depending on persistent attribute.
+func (g *GoChannel) Publish(topic string, messages ...*message.Message) error {
+	if g.isClosed() {
+		return errors.New("Pub/Sub closed")
+	}
+
+	messagesToPublish := make(message.Messages, len(messages))
+	for i, msg := range messages {
+		if g.config.PreserveContext {
+			messagesToPublish[i] = msg.CopyWithContext()
+		} else {
+			messagesToPublish[i] = msg.Copy()
+		}
+	}
+
+	g.subscribersLock.RLock()
+	defer g.subscribersLock.RUnlock()
+
+	subLock, loaded := g.subscribersByTopicLock.LoadOrStore(topic, &sync.Mutex{})
+	subLock.(*sync.Mutex).Lock()
+
+	if !loaded {
+		defer g.subscribersByTopicLock.Delete(topic)
+	}
+	defer subLock.(*sync.Mutex).Unlock()
+
+	if g.config.Persistent {
+		g.persistedMessagesLock.Lock()
+		if _, ok := g.persistedMessages[topic]; !ok {
+			g.persistedMessages[topic] = make([]*message.Message, 0)
+		}
+		g.persistedMessages[topic] = append(g.persistedMessages[topic], messagesToPublish...)
+		g.persistedMessagesLock.Unlock()
+	}
+
+	for i := range messagesToPublish {
+		msg := messagesToPublish[i]
+
+		ackedBySubscribers, err := g.sendMessage(topic, msg)
+		if err != nil {
+			return err
+		}
+
+		if g.config.BlockPublishUntilSubscriberAck {
+			g.waitForAckFromSubscribers(msg, ackedBySubscribers)
+		}
+	}
+
+	return nil
+}
+
+func (g *GoChannel) waitForAckFromSubscribers(msg *message.Message, ackedByConsumer <-chan struct{}) {
+	logFields := watermill.LogFields{"message_uuid": msg.UUID}
+	g.logger.Debug("Waiting for subscribers ack", logFields)
+
+	select {
+	case <-ackedByConsumer:
+		g.logger.Trace("Message acked by subscribers", logFields)
+	case <-g.closing:
+		g.logger.Trace("Closing Pub/Sub before ack from subscribers", logFields)
+	}
+}
+
+func (g *GoChannel) sendMessage(topic string, message *message.Message) (<-chan struct{}, error) {
+	subscribers := g.topicSubscribers(topic)
+	ackedBySubscribers := make(chan struct{})
+
+	logFields := watermill.LogFields{"message_uuid": message.UUID, "topic": topic}
+
+	if len(subscribers) == 0 {
+		close(ackedBySubscribers)
+		g.logger.Info("No subscribers to send message", logFields)
+		return ackedBySubscribers, nil
+	}
+
+	go func(subscribers []*subscriber) {
+		wg := &sync.WaitGroup{}
+
+		for i := range subscribers {
+			subscriber := subscribers[i]
+
+			wg.Add(1)
+			go func() {
+				subscriber.sendMessageToSubscriber(message, logFields)
+				wg.Done()
+			}()
+		}
+
+		wg.Wait()
+		close(ackedBySubscribers)
+	}(subscribers)
+
+	return ackedBySubscribers, nil
+}
+
+// Subscribe returns channel to which all published messages are sent.
+// Messages are not persisted. If there are no subscribers and message is produced it will be gone.
+//
+// There are no consumer groups support etc. Every consumer will receive every produced message.
+func (g *GoChannel) Subscribe(ctx context.Context, topic string) (<-chan *message.Message, error) {
+	g.closedLock.Lock()
+
+	if g.closed {
+		g.closedLock.Unlock()
+		return nil, errors.New("Pub/Sub closed")
+	}
+
+	g.subscribersWg.Add(1)
+	g.closedLock.Unlock()
+
+	g.subscribersLock.Lock()
+
+	subLock, _ := g.subscribersByTopicLock.LoadOrStore(topic, &sync.Mutex{})
+	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{}),
+		preserveContext: g.config.PreserveContext,
+	}
+
+	go func(s *subscriber, g *GoChannel) {
+		select {
+		case <-ctx.Done():
+			// unblock
+		case <-g.closing:
+			// unblock
+		}
+
+		s.Close()
+
+		g.subscribersLock.Lock()
+		defer g.subscribersLock.Unlock()
+
+		subLock, _ := g.subscribersByTopicLock.Load(topic)
+		subLock.(*sync.Mutex).Lock()
+		defer subLock.(*sync.Mutex).Unlock()
+
+		g.removeSubscriber(topic, s)
+		g.subscribersWg.Done()
+	}(s, g)
+
+	if !g.config.Persistent {
+		defer g.subscribersLock.Unlock()
+		defer subLock.(*sync.Mutex).Unlock()
+
+		g.addSubscriber(topic, s)
+
+		return s.outputChannel, nil
+	}
+
+	go func(s *subscriber) {
+		defer g.subscribersLock.Unlock()
+		defer subLock.(*sync.Mutex).Unlock()
+
+		g.persistedMessagesLock.RLock()
+		messages, ok := g.persistedMessages[topic]
+		g.persistedMessagesLock.RUnlock()
+
+		if ok {
+			for i := range messages {
+				msg := g.persistedMessages[topic][i]
+				logFields := watermill.LogFields{"message_uuid": msg.UUID, "topic": topic}
+
+				go s.sendMessageToSubscriber(msg, logFields)
+			}
+		}
+
+		g.addSubscriber(topic, s)
+	}(s)
+
+	return s.outputChannel, nil
+}
+
+func (g *GoChannel) addSubscriber(topic string, s *subscriber) {
+	if _, ok := g.subscribers[topic]; !ok {
+		g.subscribers[topic] = make([]*subscriber, 0)
+	}
+	g.subscribers[topic] = append(g.subscribers[topic], s)
+}
+
+func (g *GoChannel) removeSubscriber(topic string, toRemove *subscriber) {
+	removed := false
+	for i, sub := range g.subscribers[topic] {
+		if sub == toRemove {
+			g.subscribers[topic] = append(g.subscribers[topic][:i], g.subscribers[topic][i+1:]...)
+			removed = true
+
+			if len(g.subscribers[topic]) == 0 && !g.config.Persistent {
+				// Free up the memory taken by a topic which no longer has subscribers.
+				// This operation allows publishing and subscribing to narrowly
+				// focused topics that include random data like UUIDs in topic name.
+				//
+				// Without this operation, memory usage will grow indefinitely in a long-running service
+				// as the map grows larger and larger with keys pointing to empty slices.
+				delete(g.subscribers, topic)
+				g.subscribersByTopicLock.Delete(topic)
+			}
+			break
+		}
+	}
+	if !removed {
+		panic("cannot remove subscriber, not found " + toRemove.uuid)
+	}
+}
+
+func (g *GoChannel) topicSubscribers(topic string) []*subscriber {
+	subscribers, ok := g.subscribers[topic]
+	if !ok {
+		return nil
+	}
+
+	// let's do a copy to avoid race conditions and deadlocks due to lock
+	subscribersCopy := make([]*subscriber, len(subscribers))
+	copy(subscribersCopy, subscribers)
+
+	return subscribersCopy
+}
+
+func (g *GoChannel) isClosed() bool {
+	g.closedLock.Lock()
+	defer g.closedLock.Unlock()
+
+	return g.closed
+}
+
+// Close closes the GoChannel Pub/Sub.
+func (g *GoChannel) Close() error {
+	g.closedLock.Lock()
+	defer g.closedLock.Unlock()
+
+	if g.closed {
+		return nil
+	}
+
+	g.closed = true
+	close(g.closing)
+
+	g.logger.Debug("Closing Pub/Sub, waiting for subscribers", nil)
+	g.subscribersWg.Wait()
+
+	g.logger.Info("Pub/Sub closed", nil)
+
+	return nil
+}
+
+type subscriber struct {
+	ctx context.Context
+
+	uuid string
+
+	sending       sync.Mutex
+	outputChannel chan *message.Message
+
+	logger  watermill.LoggerAdapter
+	closed  bool
+	closing chan struct{}
+
+	preserveContext bool
+}
+
+func (s *subscriber) Close() {
+	if s.closed {
+		return
+	}
+	close(s.closing)
+
+	s.logger.Debug("Closing subscriber, waiting for sending lock", nil)
+
+	// ensuring that we are not sending to closed channel
+	s.sending.Lock()
+	defer s.sending.Unlock()
+
+	s.logger.Debug("GoChannel Pub/Sub Subscriber closed", nil)
+	s.closed = true
+
+	close(s.outputChannel)
+}
+
+func (s *subscriber) sendMessageToSubscriber(msg *message.Message, logFields watermill.LogFields) {
+	ctx := msg.Context()
+
+	// By default, the subscriber uses the context from the message and it's canceled right after it's processed.
+	// If the message's context is preserved, the top-level client is responsible for canceling it.
+	if !s.preserveContext {
+		var cancelCtx context.CancelFunc
+		ctx, cancelCtx = context.WithCancel(s.ctx)
+		defer cancelCtx()
+	}
+
+SendToSubscriber:
+	for {
+		// copy the message to prevent ack/nack propagation to other consumers
+		// also allows to make retries on a fresh copy of the original message
+		msgToSend := msg.Copy()
+		msgToSend.SetContext(ctx)
+
+		s.logger.Trace("Sending msg to subscriber", logFields)
+
+		s.sending.Lock()
+		if s.closed {
+			s.logger.Info("Pub/Sub closed, discarding msg", logFields)
+			s.sending.Unlock()
+			return
+		}
+
+		select {
+		case s.outputChannel <- msgToSend:
+			s.logger.Trace("Sent message to subscriber", logFields)
+		case <-s.closing:
+			s.logger.Trace("Closing, message discarded", logFields)
+			s.sending.Unlock()
+			return
+		}
+		s.sending.Unlock()
+
+		select {
+		case <-msgToSend.Acked():
+			s.logger.Trace("Message acked", logFields)
+			return
+		case <-msgToSend.Nacked():
+			s.logger.Trace("Nack received, resending message", logFields)
+			continue SendToSubscriber
+		case <-s.closing:
+			s.logger.Trace("Closing, message discarded", logFields)
+			return
+		}
+	}
+}