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

pubsub/gochannel/doc.go

@@ -0,0 +1,5 @@
+// This is just the simplest Pub/Sub implementation
+//
+// All Pub/Sub implementations can be found at https://watermill.io/pubsubs/
+
+package gochannel

pubsub/gochannel/fanout.go

@@ -0,0 +1,124 @@
+package gochannel
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"sync"
+
+	"github.com/ThreeDotsLabs/watermill"
+	"github.com/ThreeDotsLabs/watermill/message"
+)
+
+// FanOut is a component that receives messages from a topic and passes them
+// to all subscribers. In effect, messages are "multiplied".
+//
+// A typical use case for using FanOut is having one external subscription and multiple workers
+// inside the process.
+//
+// You need to call AddSubscription method for all topics that you want to listen to.
+// This needs to be done *before* starting the FanOut.
+//
+// FanOut exposes the standard Subscriber interface.
+type FanOut struct {
+	internalPubSub *GoChannel
+	internalRouter *message.Router
+
+	subscriber message.Subscriber
+
+	logger watermill.LoggerAdapter
+
+	subscribedTopics map[string]struct{}
+	subscribedLock   sync.Mutex
+}
+
+// NewFanOut creates a new FanOut.
+func NewFanOut(
+	subscriber message.Subscriber,
+	logger watermill.LoggerAdapter,
+) (*FanOut, error) {
+	if subscriber == nil {
+		return nil, errors.New("missing subscriber")
+	}
+	if logger == nil {
+		logger = watermill.NopLogger{}
+	}
+
+	router, err := message.NewRouter(message.RouterConfig{}, logger)
+	if err != nil {
+		return nil, err
+	}
+
+	return &FanOut{
+		internalPubSub: NewGoChannel(Config{}, logger),
+		internalRouter: router,
+
+		subscriber: subscriber,
+
+		logger: logger,
+
+		subscribedTopics: map[string]struct{}{},
+	}, nil
+}
+
+// AddSubscription add an internal subscription for the given topic.
+// You need to call this method with all topics that you want to listen to, before the FanOut is started.
+// AddSubscription is idempotent.
+func (f *FanOut) AddSubscription(topic string) {
+	f.subscribedLock.Lock()
+	defer f.subscribedLock.Unlock()
+
+	_, ok := f.subscribedTopics[topic]
+	if ok {
+		// Subscription already exists
+		return
+	}
+
+	f.logger.Trace("Adding fan-out subscription for topic", watermill.LogFields{
+		"topic": topic,
+	})
+
+	f.internalRouter.AddHandler(
+		fmt.Sprintf("fanout-%s", topic),
+		topic,
+		f.subscriber,
+		topic,
+		f.internalPubSub,
+		message.PassthroughHandler,
+	)
+
+	f.subscribedTopics[topic] = struct{}{}
+}
+
+// Run runs the FanOut.
+func (f *FanOut) Run(ctx context.Context) error {
+	return f.internalRouter.Run(ctx)
+}
+
+// Running is closed when FanOut is running.
+func (f *FanOut) Running() chan struct{} {
+	return f.internalRouter.Running()
+}
+
+func (f *FanOut) IsClosed() bool {
+	return f.internalRouter.IsClosed()
+}
+
+// Subscribe starts subscription to the FanOut's internal Pub/Sub.
+func (f *FanOut) Subscribe(ctx context.Context, topic string) (<-chan *message.Message, error) {
+	return f.internalPubSub.Subscribe(ctx, topic)
+}
+
+// Close closes the FanOut's internal Pub/Sub.
+func (f *FanOut) Close() error {
+	var err error
+
+	if routerCloseErr := f.internalRouter.Close(); routerCloseErr != nil {
+		err = errors.Join(err, routerCloseErr)
+	}
+	if internalPubSubCloseErr := f.internalPubSub.Close(); internalPubSubCloseErr != nil {
+		err = errors.Join(err, internalPubSubCloseErr)
+	}
+
+	return err
+}

pubsub/gochannel/fanout_test.go

@@ -0,0 +1,117 @@
+package gochannel_test
+
+import (
+	"context"
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/ThreeDotsLabs/watermill"
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/ThreeDotsLabs/watermill/pubsub/gochannel"
+)
+
+func TestFanOut(t *testing.T) {
+	const (
+		upstreamTopic = "upstream-topic"
+	)
+
+	logger := watermill.NopLogger{}
+
+	upstreamPubSub := gochannel.NewGoChannel(gochannel.Config{}, logger)
+
+	fanout, err := gochannel.NewFanOut(upstreamPubSub, logger)
+	require.NoError(t, err)
+
+	fanout.AddSubscription(upstreamTopic)
+
+	workersCount := 10
+	messagesCount := 100
+
+	router, err := message.NewRouter(message.RouterConfig{}, logger)
+	require.NoError(t, err)
+
+	expectedNumberOfMessages := workersCount * messagesCount
+
+	receivedMessages := make(chan struct{}, expectedNumberOfMessages)
+
+	for i := 0; i < workersCount; i++ {
+		router.AddConsumerHandler(
+			fmt.Sprintf("worker-%v", i),
+			upstreamTopic,
+			fanout,
+			func(msg *message.Message) error {
+				receivedMessages <- struct{}{}
+				return nil
+			},
+		)
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+
+	go func() {
+		err := router.Run(ctx)
+		require.NoError(t, err)
+	}()
+
+	go func() {
+		err := fanout.Run(ctx)
+		require.NoError(t, err)
+	}()
+
+	<-router.Running()
+	<-fanout.Running()
+
+	go func() {
+		for i := 0; i < messagesCount; i++ {
+			msg := message.NewMessage(watermill.NewUUID(), nil)
+			err := upstreamPubSub.Publish(upstreamTopic, msg)
+			if err != nil {
+				panic(err)
+			}
+		}
+	}()
+
+	<-ctx.Done()
+
+	counter := 0
+
+loop:
+	for {
+		select {
+		case <-receivedMessages:
+			counter += 1
+		case <-time.After(time.Second):
+			close(receivedMessages)
+			break loop
+		}
+	}
+
+	require.Equal(t, expectedNumberOfMessages, counter)
+}
+
+func TestFanOut_RouterClosed(t *testing.T) {
+	logger := watermill.NopLogger{}
+	pubSub := gochannel.NewGoChannel(gochannel.Config{}, logger)
+
+	fanout, err := gochannel.NewFanOut(pubSub, logger)
+	require.NoError(t, err)
+
+	fanout.AddSubscription("some-topic")
+
+	go func() {
+		err := fanout.Run(context.Background())
+		require.NoError(t, err)
+	}()
+
+	<-fanout.Running()
+
+	err = fanout.Close()
+	require.NoError(t, err)
+
+	assert.True(t, fanout.IsClosed())
+}

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

pubsub/gochannel/pubsub_bench_test.go

@@ -0,0 +1,34 @@
+package gochannel_test
+
+import (
+	"testing"
+
+	"github.com/ThreeDotsLabs/watermill/pubsub/gochannel"
+
+	"github.com/ThreeDotsLabs/watermill/pubsub/tests"
+
+	"github.com/ThreeDotsLabs/watermill"
+	"github.com/ThreeDotsLabs/watermill/message"
+)
+
+func BenchmarkSubscriber(b *testing.B) {
+	tests.BenchSubscriber(b, func(n int) (message.Publisher, message.Subscriber) {
+		pubSub := gochannel.NewGoChannel(
+			gochannel.Config{OutputChannelBuffer: int64(n)}, watermill.NopLogger{},
+		)
+		return pubSub, pubSub
+	})
+}
+
+func BenchmarkSubscriberPersistent(b *testing.B) {
+	tests.BenchSubscriber(b, func(n int) (message.Publisher, message.Subscriber) {
+		pubSub := gochannel.NewGoChannel(
+			gochannel.Config{
+				OutputChannelBuffer: int64(n),
+				Persistent:          true,
+			},
+			watermill.NopLogger{},
+		)
+		return pubSub, pubSub
+	})
+}

pubsub/gochannel/pubsub_internal_test.go

@@ -0,0 +1,100 @@
+package gochannel
+
+import (
+	"context"
+	"strconv"
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/ThreeDotsLabs/watermill"
+	"github.com/ThreeDotsLabs/watermill/message"
+)
+
+func TestSubscribe_clean_subscriber_data(t *testing.T) {
+	subCount := 100
+	pubSub := NewGoChannel(
+		Config{OutputChannelBuffer: int64(subCount)},
+		watermill.NewStdLogger(false, false),
+	)
+	topicName := "test_topic"
+
+	for i := 0; i < subCount; i++ {
+		ctx, cancel := context.WithCancel(context.Background())
+		_, err := pubSub.Subscribe(ctx, topicName+"_index_"+strconv.Itoa(i))
+		require.NoError(t, err)
+		cancel()
+	}
+
+	err := pubSub.Close()
+	require.NoError(t, err)
+
+	assert.Len(t, pubSub.subscribers, 0)
+	lockCount := 0
+	pubSub.subscribersByTopicLock.Range(func(_, _ any) bool {
+		lockCount++
+		return true
+	})
+	assert.Equal(t, 0, lockCount)
+
+	assert.NoError(t, pubSub.Close())
+}
+
+func TestPublish_clean_lock_data(t *testing.T) {
+	messageCount := 100
+	pubSub := NewGoChannel(
+		Config{OutputChannelBuffer: int64(messageCount)},
+		watermill.NewStdLogger(false, false),
+	)
+	topicName := "test_topic"
+
+	_, err := pubSub.Subscribe(context.Background(), topicName+"_index_"+strconv.Itoa(0))
+	require.NoError(t, err)
+
+	for i := 0; i < messageCount; i++ {
+		err := pubSub.Publish(topicName+"_index_"+strconv.Itoa(i), message.NewMessage(watermill.NewShortUUID(), nil))
+		require.NoError(t, err)
+	}
+
+	lockCount := 0
+	pubSub.subscribersByTopicLock.Range(func(_, _ any) bool {
+		lockCount++
+		return true
+	})
+	assert.Equal(t, 1, lockCount)
+
+	assert.NoError(t, pubSub.Close())
+}
+
+// TestPublish_concurrent_close ensures Close does not race with in-flight persistent Publish calls. 
+func TestPublish_concurrent_close(t *testing.T) {
+	pubSub := NewGoChannel(
+		Config{Persistent: true},
+		watermill.NewStdLogger(false, false),
+	)
+	topicName := "test_topic"
+
+	publisherCount := 20
+	startPublishing := make(chan struct{})
+	wg := sync.WaitGroup{}
+
+	for i := 0; i < publisherCount; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			<-startPublishing
+			for {
+				err := pubSub.Publish(topicName, message.NewMessage(watermill.NewShortUUID(), nil))
+				if err != nil {
+					return
+				}
+			}
+		}()
+	}
+
+	close(startPublishing)
+	require.NoError(t, pubSub.Close())
+	wg.Wait()
+}

pubsub/gochannel/pubsub_stress_test.go

@@ -0,0 +1,25 @@
+//go:build stress
+// +build stress
+
+package gochannel_test
+
+import (
+	"testing"
+
+	"github.com/ThreeDotsLabs/watermill/pubsub/tests"
+)
+
+func TestPublishSubscribe_stress(t *testing.T) {
+	tests.TestPubSubStressTest(
+		t,
+		tests.Features{
+			ConsumerGroups:        false,
+			ExactlyOnceDelivery:   true,
+			GuaranteedOrder:       false,
+			Persistent:            false,
+			RequireSingleInstance: true,
+		},
+		createPersistentPubSub,
+		nil,
+	)
+}

pubsub/gochannel/pubsub_test.go

@@ -0,0 +1,399 @@
+package gochannel_test
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/ThreeDotsLabs/watermill"
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/ThreeDotsLabs/watermill/message/subscriber"
+	"github.com/ThreeDotsLabs/watermill/pubsub/gochannel"
+	"github.com/ThreeDotsLabs/watermill/pubsub/tests"
+)
+
+func createPersistentPubSub(t *testing.T) (message.Publisher, message.Subscriber) {
+	pubSub := gochannel.NewGoChannel(
+		gochannel.Config{
+			OutputChannelBuffer: 10000,
+			Persistent:          true,
+		},
+		watermill.NewStdLogger(true, true),
+	)
+	return pubSub, pubSub
+}
+
+func createPersistentPubSubWithContextPreserved(t *testing.T) (message.Publisher, message.Subscriber) {
+	pubSub := gochannel.NewGoChannel(
+		gochannel.Config{
+			OutputChannelBuffer: 10000,
+			Persistent:          true,
+			PreserveContext:     true,
+		},
+		watermill.NewStdLogger(true, true),
+	)
+	return pubSub, pubSub
+}
+
+func TestPublishSubscribe_persistent(t *testing.T) {
+	tests.TestPubSub(
+		t,
+		tests.Features{
+			ConsumerGroups:        false,
+			ExactlyOnceDelivery:   true,
+			GuaranteedOrder:       false,
+			Persistent:            false,
+			RequireSingleInstance: true,
+		},
+		createPersistentPubSub,
+		nil,
+	)
+}
+
+func TestPublishSubscribe_context_preserved(t *testing.T) {
+	tests.TestPubSub(
+		t,
+		tests.Features{
+			ConsumerGroups:        false,
+			ExactlyOnceDelivery:   true,
+			GuaranteedOrder:       false,
+			Persistent:            false,
+			RequireSingleInstance: true,
+			ContextPreserved:      true,
+		},
+		createPersistentPubSubWithContextPreserved,
+		nil,
+	)
+}
+
+func TestPublishSubscribe_not_persistent(t *testing.T) {
+	messagesCount := 100
+	pubSub := gochannel.NewGoChannel(
+		gochannel.Config{OutputChannelBuffer: int64(messagesCount)},
+		watermill.NewStdLogger(true, true),
+	)
+	topicName := "test_topic_" + watermill.NewUUID()
+
+	msgs, err := pubSub.Subscribe(context.Background(), topicName)
+	require.NoError(t, err)
+
+	sendMessages := tests.PublishSimpleMessages(t, messagesCount, pubSub, topicName)
+	receivedMsgs, _ := subscriber.BulkRead(msgs, messagesCount, time.Second)
+
+	tests.AssertAllMessagesReceived(t, sendMessages, receivedMsgs)
+
+	assert.NoError(t, pubSub.Close())
+}
+
+func TestPublishSubscribe_not_persistent_with_context(t *testing.T) {
+	messagesCount := 100
+	pubSub := gochannel.NewGoChannel(
+		gochannel.Config{OutputChannelBuffer: int64(messagesCount), PreserveContext: true},
+		watermill.NewStdLogger(true, true),
+	)
+	topicName := "test_topic_" + watermill.NewUUID()
+
+	msgs, err := pubSub.Subscribe(context.Background(), topicName)
+	require.NoError(t, err)
+
+	const contextKeyString = "foo"
+	sendMessages := tests.PublishSimpleMessagesWithContext(t, messagesCount, contextKeyString, pubSub, topicName)
+	receivedMsgs, _ := subscriber.BulkRead(msgs, messagesCount, time.Second)
+
+	expectedContexts := make(map[string]context.Context)
+	for _, msg := range sendMessages {
+		expectedContexts[msg.UUID] = msg.Context()
+	}
+	tests.AssertAllMessagesReceived(t, sendMessages, receivedMsgs)
+	tests.AssertAllMessagesHaveSameContext(t, contextKeyString, expectedContexts, receivedMsgs)
+
+	assert.NoError(t, pubSub.Close())
+}
+
+func TestPublishSubscribe_block_until_ack(t *testing.T) {
+	pubSub := gochannel.NewGoChannel(
+		gochannel.Config{BlockPublishUntilSubscriberAck: true},
+		watermill.NewStdLogger(true, true),
+	)
+	topicName := "test_topic_" + watermill.NewUUID()
+
+	msgs, err := pubSub.Subscribe(context.Background(), topicName)
+	require.NoError(t, err)
+
+	published := make(chan struct{})
+	go func() {
+		err := pubSub.Publish(topicName, message.NewMessage("1", nil))
+		require.NoError(t, err)
+		close(published)
+	}()
+
+	msg1 := <-msgs
+	select {
+	case <-published:
+		t.Fatal("publish should be blocked until ack")
+	default:
+		// ok
+	}
+
+	msg1.Nack()
+	select {
+	case <-published:
+		t.Fatal("publish should be blocked after nack")
+	default:
+		// ok
+	}
+
+	msg2 := <-msgs
+	msg2.Ack()
+
+	select {
+	case <-published:
+		// ok
+	case <-time.After(time.Second):
+		t.Fatal("publish should be not blocked after ack")
+	}
+}
+
+func TestPublishSubscribe_race_condition_on_subscribe(t *testing.T) {
+	testsCount := 15
+	if testing.Short() {
+		testsCount = 3
+	}
+
+	for i := 0; i < testsCount; i++ {
+		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
+			t.Parallel()
+			testPublishSubscribeSubRace(t)
+		})
+	}
+}
+
+func TestSubscribe_race_condition_when_closing(t *testing.T) {
+	testsCount := 15
+	if testing.Short() {
+		testsCount = 3
+	}
+
+	for i := 0; i < testsCount; i++ {
+		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
+			t.Parallel()
+			pubSub := gochannel.NewGoChannel(
+				gochannel.Config{},
+				watermill.NewStdLogger(true, false),
+			)
+			go func() {
+				err := pubSub.Close()
+				require.NoError(t, err)
+			}()
+			_, err := pubSub.Subscribe(context.Background(), "topic")
+			require.NoError(t, err)
+		})
+	}
+}
+
+func TestPublish_race_condition_when_closing(t *testing.T) {
+	testsCount := 15
+	if testing.Short() {
+		testsCount = 3
+	}
+
+	for i := 0; i < testsCount; i++ {
+		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
+			t.Parallel()
+			pubSub := gochannel.NewGoChannel(
+				gochannel.Config{},
+				watermill.NewStdLogger(true, false),
+			)
+			go func() {
+				_ = pubSub.Publish("topic", message.NewMessage(watermill.NewShortUUID(), nil))
+			}()
+
+			err := pubSub.Close()
+			require.NoError(t, err)
+		})
+	}
+}
+
+func TestPublishSubscribe_do_not_block_other_subscribers(t *testing.T) {
+	pubSub := gochannel.NewGoChannel(
+		gochannel.Config{},
+		watermill.NewStdLogger(true, true),
+	)
+	topicName := "test_topic_" + watermill.NewUUID()
+
+	msgsFromSubscriber1, err := pubSub.Subscribe(context.Background(), topicName)
+	require.NoError(t, err)
+
+	_, err = pubSub.Subscribe(context.Background(), topicName)
+	require.NoError(t, err)
+
+	msgsFromSubscriber3, err := pubSub.Subscribe(context.Background(), topicName)
+	require.NoError(t, err)
+
+	err = pubSub.Publish(topicName, message.NewMessage("1", nil))
+	require.NoError(t, err)
+
+	received := make(chan struct{})
+	go func() {
+		msg := <-msgsFromSubscriber1
+		msg.Ack()
+
+		msg = <-msgsFromSubscriber3
+		msg.Ack()
+
+		close(received)
+	}()
+
+	select {
+	case <-received:
+		// ok
+	case <-time.After(5 * time.Second):
+		t.Fatal("subscriber which didn't ack a message blocked other subscribers from receiving it")
+	}
+}
+
+func TestPublishSubscribe_flush_output_channel(t *testing.T) {
+	messagesCount := 300
+	logger := watermill.NewStdLogger(true, true)
+	ctx := context.Background()
+	config := gochannel.Config{
+		OutputChannelBuffer:            int64(messagesCount),
+		Persistent:                     false,
+		BlockPublishUntilSubscriberAck: false,
+	}
+	pubSub := gochannel.NewGoChannel(
+		config,
+		logger,
+	)
+
+	totalMessage := 0
+	artificialWorkload := time.Millisecond * 5 //keep it small but noticeable in logs
+	topicName := "test_topic"
+
+	messageChannel, err := pubSub.Subscribe(ctx, topicName)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// waitgroup for stopping the subscriber handler from processing/receiving messages until we are done filling the buffer of the pubSub
+	var wgStartSubscriber sync.WaitGroup
+	wgStartSubscriber.Add(1)
+	// waitgroup for expected buffer to be flushed
+	var wgFlushBuffer sync.WaitGroup
+	wgFlushBuffer.Add(messagesCount)
+
+	// start subscriber handler in a go routine
+	// reads out the messages, if all is ok it should be able to ("flush") read all messages from a 'closed' pubsub
+	go func(messageChannel <-chan *message.Message) {
+		wgStartSubscriber.Wait()
+		for msg := range messageChannel {
+			// artificial workload
+			time.Sleep(artificialWorkload)
+			msg.Ack()
+			logger.Trace("message acked", nil)
+			// would normally use atomic value here but concurrency shouldn't be an issue for this test
+			totalMessage++
+			wgFlushBuffer.Done()
+		}
+		logger.Trace("channel closed", nil)
+	}(messageChannel)
+
+	tests.PublishSimpleMessages(t, messagesCount, pubSub, topicName)
+	// wait for buffer to fill then start reading in subscriber handler
+	for {
+		if len(messageChannel) != int(config.OutputChannelBuffer) {
+			continue
+		} else {
+			wgStartSubscriber.Done()
+			break
+		}
+	}
+
+	err = pubSub.Close()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Publishing new message should still error as expected
+	err = pubSub.Publish(topicName, message.NewMessage(watermill.NewUUID(), []byte("x")))
+	assert.ErrorContains(t, err, "Pub/Sub closed")
+
+	// And so should subscribe
+	_, err = pubSub.Subscribe(ctx, topicName)
+	assert.ErrorContains(t, err, "Pub/Sub closed")
+
+	wgFlushBuffer.Wait()
+	// But subscriber handler should still be able to read the remaining messages from output channel aka flushing
+	assert.Equal(t, messagesCount, totalMessage)
+}
+
+func testPublishSubscribeSubRace(t *testing.T) {
+	t.Helper()
+
+	messagesCount := 500
+	subscribersCount := 200
+	if testing.Short() {
+		messagesCount = 200
+		subscribersCount = 20
+	}
+
+	pubSub := gochannel.NewGoChannel(
+		gochannel.Config{
+			OutputChannelBuffer: int64(messagesCount),
+			Persistent:          true,
+		},
+		watermill.NewStdLogger(true, false),
+	)
+
+	allSent := sync.WaitGroup{}
+	allSent.Add(messagesCount)
+	allReceived := sync.WaitGroup{}
+
+	sentMessages := message.Messages{}
+	go func() {
+		for i := 0; i < messagesCount; i++ {
+			msg := message.NewMessage(watermill.NewUUID(), nil)
+			sentMessages = append(sentMessages, msg)
+
+			go func() {
+				require.NoError(t, pubSub.Publish("topic", msg))
+				allSent.Done()
+			}()
+		}
+	}()
+
+	subscriberReceivedCh := make(chan message.Messages, subscribersCount)
+	for i := 0; i < subscribersCount; i++ {
+		allReceived.Add(1)
+
+		go func() {
+			msgs, err := pubSub.Subscribe(context.Background(), "topic")
+			require.NoError(t, err)
+
+			received, _ := subscriber.BulkRead(msgs, messagesCount, time.Second*10)
+			subscriberReceivedCh <- received
+
+			allReceived.Done()
+		}()
+	}
+
+	log.Println("waiting for all sent")
+	allSent.Wait()
+
+	log.Println("waiting for all received")
+	allReceived.Wait()
+
+	close(subscriberReceivedCh)
+
+	log.Println("asserting")
+
+	for subMsgs := range subscriberReceivedCh {
+		tests.AssertAllMessagesReceived(t, sentMessages, subMsgs)
+	}
+}