fixed shutdown logic

async.go

@@ -6,7 +6,9 @@ import (
 	"sync"
 )
 
-// AsyncTopic allows any message T to be broadcast to all subscribers.
+// AsyncTopic allows any message T to be broadcast to subscribers. Publishing as well as
+// subscribing happens asynchronously. Due to the nature of async processes this cannot guarantee
+// message delivery nor delivery order.
 type AsyncTopic[T any] struct {
 	options AsyncTopicOptions
 
@@ -17,7 +19,7 @@ type AsyncTopic[T any] struct {
 	subscribeCh chan Subscriber[T]
 }
 
-// NewAsyncTopic creates a Topic that publishes messages asynchronously.
+// NewAsyncTopic creates an AsyncTopic that will be closed when the given context is cancelled.
 func NewAsyncTopic[T any](ctx context.Context, opts ...AsyncTopicOption) *AsyncTopic[T] {
 	options := AsyncTopicOptions{
 		onClose:     func() {},
@@ -40,8 +42,6 @@ func NewAsyncTopic[T any](ctx context.Context, opts ...AsyncTopicOption) *AsyncT
 }
 
 func (t *AsyncTopic[T]) run(ctx context.Context) {
-	defer close(t.publishCh)
-	defer close(t.subscribeCh)
 	defer t.close()
 
 	var subscribers []Subscriber[T]
@@ -73,34 +73,65 @@ func (t *AsyncTopic[T]) run(ctx context.Context) {
 // Publish broadcasts a msg to all subscribers.
 func (t *AsyncTopic[T]) Publish(msg T) error {
 	t.mu.RLock()
-	defer t.mu.RUnlock()
 
 	if t.closed {
+		t.mu.RUnlock()
 		return fmt.Errorf("async topic publish: %w", ErrTopicClosed)
 	}
 
-	t.publishCh <- msg
+	go func() {
+		t.publishCh <- msg
+		t.mu.RUnlock()
+	}()
+
 	return nil
 }
 
 // Subscribe adds a Subscriber func that will consume future published messages.
 func (t *AsyncTopic[T]) Subscribe(fn Subscriber[T]) error {
 	t.mu.RLock()
-	defer t.mu.RUnlock()
 
 	if t.closed {
+		t.mu.RUnlock()
 		return fmt.Errorf("async topic subscribe: %w", ErrTopicClosed)
 	}
 
-	t.subscribeCh <- fn
+	go func() {
+		t.subscribeCh <- fn
+		t.mu.RUnlock()
+	}()
+
 	return nil
 }
 
 func (t *AsyncTopic[T]) close() {
+	var wg sync.WaitGroup
+
+	wg.Add(2)
+
+	go func() {
+		defer wg.Done()
+		for range t.publishCh {
+			// drain publishCh to release all read locks
+		}
+	}()
+
+	go func() {
+		defer wg.Done()
+		for range t.subscribeCh {
+			// drain subscribeCh to release all read locks
+		}
+	}()
+
 	t.mu.Lock()
-	t.closed = true
+	t.closed = true // no more subscribing or publishing
 	t.mu.Unlock()
 
+	close(t.publishCh)
+	close(t.subscribeCh)
+
+	wg.Wait()
+
 	t.options.onClose()
 }
 

async_test.go

@@ -10,28 +10,34 @@ import (
 )
 
 func TestAsyncTopic_SinglePublisherSingleSubscriber(t *testing.T) {
-	const msgCount = 5
+	const msgCount = 100
 
-	topic := NewAsyncTopic[int](context.Background())
+	subscriberReady := make(chan struct{})
+	defer close(subscriberReady)
 
-	feedback := make(chan int, msgCount)
+	topic := NewAsyncTopic[int](context.Background(), WithOnSubscribe(func(count int) {
+		subscriberReady <- struct{}{}
+	}))
+
+	feedback := make(chan struct{}, msgCount)
 	defer close(feedback)
 
 	err := topic.Subscribe(func(i int) bool {
-		feedback <- i
+		feedback <- struct{}{}
 		return true
 	})
 	require.NoError(t, err)
 
-	for i := 0; i < msgCount; i++ {
+	<-subscriberReady
+
+	for i := range msgCount {
 		require.NoError(t, topic.Publish(i))
 	}
 
 	count := 0
 	for count < msgCount {
 		select {
-		case f := <-feedback:
-			assert.Equalf(t, count, f, "expected to get %d but got %d instead", count, f)
+		case <-feedback:
 			count++
 
 		case <-time.After(time.Second):
@@ -47,7 +53,14 @@ func TestAsyncTopic_MultiPublishersMultiSubscribers(t *testing.T) {
 		msgCount = pubCount * 100 // total messages to publish (delivered to EACH subscriber)
 	)
 
-	topic := NewAsyncTopic[int](context.Background())
+	subscribersReady := make(chan struct{})
+	defer close(subscribersReady)
+
+	topic := NewAsyncTopic[int](context.Background(), WithOnSubscribe(func(count int) {
+		if count == subCount {
+			subscribersReady <- struct{}{}
+		}
+	}))
 
 	expFeedbackCount := msgCount * subCount
 	feedback := make(chan int, expFeedbackCount)
@@ -67,6 +80,8 @@ func TestAsyncTopic_MultiPublishersMultiSubscribers(t *testing.T) {
 	}
 	close(toDeliver)
 
+	<-subscribersReady
+
 	for range pubCount {
 		go func() {
 			for msg := range toDeliver {
@@ -135,52 +150,45 @@ func TestAsyncTopic_WithOnSubscribe(t *testing.T) {
 	}
 }
 
-func TestAsyncTopic_SubscribeClosedTopicError(t *testing.T) {
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-
-	feedback := make(chan struct{})
-	defer close(feedback)
-
-	topic := NewAsyncTopic[int](ctx, WithOnClose(func() { feedback <- struct{}{} }))
-
-	require.NoError(t, topic.Subscribe(func(i int) bool { return true }))
-
-	cancel() // this should close the topic, no more subscribers should be accepted
-
-	select {
-	case <-feedback:
-		break
-
-	case <-time.After(time.Second):
-		t.Fatalf("expected feedback by now")
+func TestAsyncTopic_ClosedTopicError(t *testing.T) {
+	testCases := []struct {
+		name     string
+		assertFn func(*AsyncTopic[int])
+	}{
+		{
+			name: "publishing returns error",
+			assertFn: func(topic *AsyncTopic[int]) {
+				assert.Error(t, ErrTopicClosed, topic.Publish(1))
+			},
+		},
+		{
+			name: "subscribing returns error",
+			assertFn: func(topic *AsyncTopic[int]) {
+				assert.Error(t, ErrTopicClosed, topic.Subscribe(func(i int) bool { return true }))
+			},
+		},
 	}
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
 
-	require.Error(t, ErrTopicClosed, topic.Subscribe(func(i int) bool { return true }))
-}
+			feedback := make(chan struct{})
+			defer close(feedback)
 
-func TestAsyncTopic_PublishClosedTopicError(t *testing.T) {
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
+			topic := NewAsyncTopic[int](ctx, WithOnClose(func() { feedback <- struct{}{} }))
 
-	feedback := make(chan struct{})
-	defer close(feedback)
+			cancel() // this should close the topic, no more messages can be published
 
-	topic := NewAsyncTopic[int](ctx, WithOnClose(func() { feedback <- struct{}{} }))
+			select {
+			case <-feedback:
+				break
 
-	require.NoError(t, topic.Subscribe(func(i int) bool { return true }))
+			case <-time.After(time.Second):
+				t.Fatalf("expected feedback by now")
+			}
 
-	require.NoError(t, topic.Publish(123))
-
-	cancel() // this should close the topic, no more messages can be published
-
-	select {
-	case <-feedback:
-		break
-
-	case <-time.After(time.Second):
-		t.Fatalf("expected feedback by now")
+			tc.assertFn(topic)
+		})
 	}
-
-	require.Error(t, ErrTopicClosed, topic.Publish(1))
 }