fix potential issue with closing asyn topic

async.go

@@ -6,12 +6,10 @@ import (
 )
 
 // AsyncTopic allows any message T to be broadcast to subscribers. Publishing as well as
-// subscribing happens asynchronously.
-// This guarantees that every published message will be delivered but does NOT guarantee delivery
-// order.
-// In the unlikely scenario where subscribers are being queued very aggressively it is possible
-// that some might never actually receive any message. Subscriber registration order is also not
-// guaranteed.
+// subscribing happens asynchronously (as non-blocking as possible).
+// Closing the topic guarantees that published message will be delivered and no further messages
+// nor subscribers will be accepted.
+// Delivery order is NOT guaranteed.
 type AsyncTopic[T any] struct {
 	options TopicOptions
 
@@ -23,12 +21,11 @@ type AsyncTopic[T any] struct {
 	subscribeCh chan Subscriber[T]
 }
 
-// NewAsyncTopic creates an AsyncTopic that will be closed when the given context is cancelled.
-// After closed calls to Publish or Subscribe will return an error.
+// NewAsyncTopic creates an AsyncTopic.
 func NewAsyncTopic[T any](opts ...TopicOption) *AsyncTopic[T] {
 	options := TopicOptions{
-		onClose:     func() {},          // Called after the Topic is closed and all messages have been delivered.
-		onSubscribe: func(count int) {}, // Called everytime a new subscriber is added
+		onClose:     func() {}, // Called after the Topic is closed and all messages have been delivered.
+		onSubscribe: func() {}, // Called everytime a new subscriber is added
 	}
 
 	for _, opt := range opts {
@@ -53,13 +50,19 @@ func (t *AsyncTopic[T]) Close() {
 	t.mu.RLock()
 	closing := t.closing
 	t.mu.RUnlock()
-
 	if closing {
 		// It's either already closed or it's closing.
 		return
 	}
 
 	t.mu.Lock()
+	if closing {
+		// It is possible that 2 go routines attempted to aquire the lock to close this topic.
+		t.mu.Unlock()
+		<-t.closed
+		return
+	}
+
 	t.closing = true // no more subscribing or publishing
 	t.mu.Unlock()
 
@@ -78,7 +81,7 @@ func (t *AsyncTopic[T]) run() {
 	defer func() {
 		// There is only one way to get here: the topic is now closing!
 		// Because both `subscribeCh` and `publishCh` channels are closed when the topic is closed
-		// this will always eventually return.
+		// we can assume this will always eventually return.
 		// This will deliver any potential queued message thus fulfilling the message delivery
 		// promise.
 		for msg := range t.publishCh {
@@ -94,7 +97,7 @@ func (t *AsyncTopic[T]) run() {
 			}
 
 			subscribers = append(subscribers, newCallback)
-			t.options.onSubscribe(len(subscribers))
+			t.options.onSubscribe()
 
 		case msg, more := <-t.publishCh:
 			if !more {
@@ -106,10 +109,12 @@ func (t *AsyncTopic[T]) run() {
 	}
 }
 
-// Publish broadcasts a msg to all subscribers.
+// Publish broadcasts a msg to all subscribers asynchronously.
 func (t *AsyncTopic[T]) Publish(msg T) error {
 	t.mu.RLock()
 
+	// We hold the Read lock until we are done with publishing to avoid panic due to a closed channel.
+
 	if t.closing {
 		t.mu.RUnlock()
 		return fmt.Errorf("async topic publish: %w", ErrTopicClosed)
@@ -123,7 +128,7 @@ func (t *AsyncTopic[T]) Publish(msg T) error {
 	return nil
 }
 
-// Subscribe adds a Subscriber func that will consume future published messages.
+// Subscribe registers a Subscriber func asynchronously.
 func (t *AsyncTopic[T]) Subscribe(fn Subscriber[T]) error {
 	t.mu.RLock()
 

async_test.go

@@ -1,7 +1,6 @@
 package gubgub
 
 import (
-	"sync/atomic"
 	"testing"
 	"time"
 
@@ -98,50 +97,24 @@ func TestAsyncTopic_MultiPublishersMultiSubscribers(t *testing.T) {
 	}
 }
 
-func TestAsyncTopic_WithOnClose(t *testing.T) {
-	feedback := make(chan struct{}, 1)
-	defer close(feedback)
+func TestAsyncTopic_CloseIsIdempotent(t *testing.T) {
+	topic := NewAsyncTopic[int]()
 
-	topic := NewAsyncTopic[int](WithOnClose(func() { feedback <- struct{}{} }))
-
-	topic.Close()
+	feedback := make(chan struct{})
+	go func() {
+		topic.Close()
+		topic.Close()
+		close(feedback)
+	}()
 
 	select {
 	case <-feedback:
-		break
-
+		return
 	case <-testTimer(t, time.Second).C:
 		t.Fatalf("expected feedback by now")
 	}
 }
 
-func TestAsyncTopic_WithOnSubscribe(t *testing.T) {
-	const totalSub = 10
-
-	feedback := make(chan int, totalSub)
-	defer close(feedback)
-
-	topic := NewAsyncTopic[int](WithOnSubscribe(func(count int) { feedback <- count }))
-
-	for range totalSub {
-		topic.Subscribe(func(i int) bool { return true })
-	}
-
-	count := 0
-	timeout := testTimer(t, time.Second)
-
-	for count < totalSub {
-		select {
-		case c := <-feedback:
-			count++
-			assert.Equal(t, count, c, "expected %d but got %d instead", count, c)
-
-		case <-timeout.C:
-			t.Fatalf("expected %d feedback items by now but only got %d", totalSub, count)
-		}
-	}
-}
-
 func TestAsyncTopic_ClosedTopicError(t *testing.T) {
 	testCases := []struct {
 		name     string
@@ -231,11 +204,10 @@ func withNotifyOnNthSubscriber(t testing.TB, n int64) (TopicOption, <-chan struc
 		close(notify)
 	})
 
-	var counter atomic.Int64
-
-	return WithOnSubscribe(func(count int) {
-		c := counter.Add(1)
-		if c == n {
+	var counter int64
+	return WithOnSubscribe(func() {
+		counter++
+		if counter == n {
 			notify <- struct{}{}
 		}
 	}), notify