fix: eliminate send-on-closed-channel race in subscriptions (#13)

- fix: replace server channel closing with per-subscription done channel
- fix: unsubscribe order — remove event handler before signaling shutdown
- fix: guard handler sends with done to avoid late sends
- fix: stop worker on done; never close serverCh
- fix: apply same pattern to unicast and fanout subscriptions
- test: add deterministic tests using fakeEventBus to simulate in-flight publish after unsubscribe

Signed-off-by: Andres Morey <[email protected]>

Author: amorey

dispatcher.go

@@ -67,13 +67,15 @@ type Subscription struct {
 	serverCh        chan server
 	cleanup         func()
 	unsubscribeOnce sync.Once
+	done            chan struct{}
 }
 
 // Ends subscription
 func (sub *Subscription) Unsubscribe() {
 	sub.unsubscribeOnce.Do(func() {
-		close(sub.serverCh)
+		// First, stop receiving new events, then signal worker shutdown.
 		sub.cleanup()
+		close(sub.done)
 	})
 }
 
@@ -130,12 +132,17 @@ func (d *Dispatcher) Unicast(ctx context.Context, nodeName string, fn DispatchHa
 // they become available until Unsubscribe() is called
 func (d *Dispatcher) UnicastSubscribe(ctx context.Context, nodeName string, fn DispatchHandler) (*Subscription, error) {
 	serverCh := make(chan server)
+	done := make(chan struct{})
 
 	// server handler
 	handleNewServers := func(newServers []server) {
 		for _, server := range newServers {
 			if server.nodeName == nodeName {
-				serverCh <- server
+				select {
+				case <-done:
+					return
+				case serverCh <- server:
+				}
 			}
 		}
 	}
@@ -144,14 +151,9 @@ func (d *Dispatcher) UnicastSubscribe(ctx context.Context, nodeName string, fn D
 	go func() {
 		for {
 			select {
-			case <-ctx.Done():
+			case <-done:
 				return
-			case server, ok := <-serverCh:
-				if !ok {
-					// unsubscribe was called
-					return
-				}
-
+			case server := <-serverCh:
 				// execute dispatch handler in goroutine
 				connCtx := context.WithValue(ctx, dispatcherAddrCtxKey, fmt.Sprintf("%s:%s", server.ip, d.connectArgs.Port))
 				go fn(connCtx, d.conn)
@@ -176,6 +178,7 @@ func (d *Dispatcher) UnicastSubscribe(ctx context.Context, nodeName string, fn D
 		cleanup: func() {
 			d.eventbus.Unsubscribe("add:servers", handleNewServers)
 		},
+		done: done,
 	}, nil
 }
 
@@ -213,26 +216,26 @@ func (d *Dispatcher) Fanout(ctx context.Context, fn DispatchHandler) {
 // they become available until Unsubscribe() is called
 func (d *Dispatcher) FanoutSubscribe(ctx context.Context, fn DispatchHandler) (*Subscription, error) {
 	serverCh := make(chan server)
+	done := make(chan struct{})
 
 	// server handler
 	handleNewServers := func(newServers []server) {
 		for _, server := range newServers {
-			serverCh <- server
+			select {
+			case <-done:
+				return
+			case serverCh <- server:
+			}
 		}
 	}
 
 	// worker
 	go func() {
 		for {
 			select {
-			case <-ctx.Done():
+			case <-done:
 				return
-			case server, ok := <-serverCh:
-				if !ok {
-					// unsubscribe was called
-					return
-				}
-
+			case server := <-serverCh:
 				// execute dispatch handler in goroutine
 				connCtx := context.WithValue(ctx, dispatcherAddrCtxKey, fmt.Sprintf("%s:%s", server.ip, d.connectArgs.Port))
 				go fn(connCtx, d.conn)
@@ -257,6 +260,7 @@ func (d *Dispatcher) FanoutSubscribe(ctx context.Context, fn DispatchHandler) (*
 		cleanup: func() {
 			d.eventbus.Unsubscribe("add:servers", handleNewServers)
 		},
+		done: done,
 	}, nil
 }
 

dispatcher_test.go

@@ -266,3 +266,87 @@ func TestDispatcherFanoutSubscribe(t *testing.T) {
 	// check result
 	require.Equal(t, wantIps, mapset.NewSet(ips...))
 }
+
+// Exposes the race: Unsubscribe closes the internal channel, so any in-flight
+// event handler that tries to send will panic with: "send on closed channel".
+// This test models the event bus behavior by sending to the channel after
+// Unsubscribe; it should panic with the current implementation.
+// fakeEventBus lets us trigger the subscription callback even after Unsubscribe
+// to simulate an in-flight publish racing with unsubscribe.
+type fakeEventBus struct {
+	mu      sync.Mutex
+	handler func([]server)
+}
+
+func (b *fakeEventBus) Subscribe(topic string, fn interface{}) error {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	b.handler = fn.(func([]server))
+	return nil
+}
+func (b *fakeEventBus) SubscribeOnce(topic string, fn interface{}) error {
+	return b.Subscribe(topic, fn)
+}
+func (b *fakeEventBus) SubscribeAsync(topic string, fn interface{}, once bool) error {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	// Our code subscribes with: func([]server)
+	b.handler = fn.(func([]server))
+	return nil
+}
+func (b *fakeEventBus) SubscribeOnceAsync(topic string, fn interface{}) error {
+	return b.Subscribe(topic, fn)
+}
+func (b *fakeEventBus) Unsubscribe(topic string, fn interface{}) error { return nil }
+func (b *fakeEventBus) HasCallback(topic string) bool {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	return b.handler != nil
+}
+func (b *fakeEventBus) Publish(topic string, args ...interface{}) {
+	b.mu.Lock()
+	h := b.handler
+	b.mu.Unlock()
+	if h != nil {
+		h(args[0].([]server))
+	}
+}
+func (b *fakeEventBus) WaitAsync() {}
+
+func TestUnicastSubscribe_UnsubscribeThenPublishDoesNotPanic(t *testing.T) {
+	// initialize dispatcher with fake bus to control callback timing
+	d := newTestDispatcher()
+	fb := &fakeEventBus{}
+	d.eventbus = fb
+
+	// start a subscription
+	sub, err := d.UnicastSubscribe(context.Background(), "n1", func(ctx context.Context, clientconn *grpc.ClientConn) {})
+	require.Nil(t, err)
+
+	// Unsubscribe; publishing should not panic even if handler runs
+	sub.Unsubscribe()
+
+	// Simulate an in-flight publish that still calls the handler
+	require.NotPanics(t, func() {
+		fb.Publish("add:servers", []server{{ip: "1.2.3.4", nodeName: "n1"}})
+	})
+}
+
+func TestFanoutSubscribe_UnsubscribeThenPublishDoesNotPanic(t *testing.T) {
+	// initialize dispatcher with fake bus to control callback timing
+	d := newTestDispatcher()
+	fb := &fakeEventBus{}
+	d.eventbus = fb
+
+	// start a subscription
+	sub, err := d.FanoutSubscribe(context.Background(), func(ctx context.Context, clientconn *grpc.ClientConn) {})
+	require.Nil(t, err)
+
+	// Unsubscribe; publishing should not panic even if handler runs
+	sub.Unsubscribe()
+
+	// Simulate an in-flight publish that still calls the handler
+	require.NotPanics(t, func() {
+		fb.Publish("add:servers", []server{{ip: "1.2.3.4", nodeName: "n1"}})
+	})
+}