fix: drain channel on unsubscribe

Author: alecthomas

pubsub/pubsub.go

@@ -16,6 +16,9 @@ import (
 // This is a last-ditch effort to avoid deadlocks.
 const AckTimeout = time.Second * 30
 
+// PublishTimeout is the time to wait for a publish before panicking.
+const PublishTimeout = time.Second * 10
+
 // Message is a message that must be acknowledge by the receiver.
 type Message[T any] struct {
 	Msg T
@@ -148,6 +151,11 @@ func (s *Topic[T]) Unsubscribe(c chan T) {
 		panic("channel not subscribed")
 	}
 	s.rawChannelMap.Delete(c)
+	// Drain the subscription channel
+	go func() {
+		for range c {
+		}
+	}()
 	s.control <- unsubscribe[T](ackch.(chan Message[T]))
 }
 
@@ -203,6 +211,7 @@ func (s *Topic[T]) run() {
 				case err := <-smsg.ack:
 					errs = append(errs, err)
 				case <-timer.C:
+					// Print all goroutines
 					panic("ack timeout for " + sub.subscriber)
 				}
 				timer.Stop()

pubsub/pubsub_test.go

@@ -1,6 +1,7 @@
 package pubsub_test
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"testing"
@@ -107,3 +108,29 @@ func TestSyncPubSub(t *testing.T) {
 	}
 	assert.Equal(t, []string{"received", "published", "received", "published"}, actual)
 }
+
+func TestPubSubPanicAfterUnsubscribe(t *testing.T) {
+	t.Skip("This test is slow")
+	topic := New[string]()
+	for range 100 {
+		go func() {
+			foo := topic.Subscribe(make(chan string, 1))
+			<-time.After(time.Second)
+			topic.Unsubscribe(foo)
+		}()
+	}
+	ctx, cancel := context.WithCancel(context.TODO())
+	defer cancel()
+	go func() {
+		for {
+			select {
+			case <-time.After(time.Millisecond * 100):
+				topic.Publish("foo")
+
+			case <-ctx.Done():
+				return
+			}
+		}
+	}()
+	<-time.After(time.Minute)
+}