Give the broadcaster a shutdown path

cmd/compass-sidecar/broadcaster.go

@@ -18,26 +18,45 @@ type Broadcaster struct {
 	removeSub chan chan trace.Trace
 	broadcast chan trace.Trace
 	changes   chan struct{}
+	// done is closed when run has returned, so the public methods stop blocking
+	// on a loop which is no longer reading.
+	done chan struct{}
 }
 
-// NewBroadcaster creates and starts a new broadcaster.
-func NewBroadcaster() *Broadcaster {
+// NewBroadcaster creates and starts a new broadcaster. It runs until ctx is
+// cancelled, at which point it closes every subscriber channel so blocked
+// consumers unblock and return.
+func NewBroadcaster(ctx context.Context) *Broadcaster {
 	b := &Broadcaster{
 		subs:      make(map[chan trace.Trace]struct{}),
 		addSub:    make(chan chan trace.Trace),
 		removeSub: make(chan chan trace.Trace),
 		broadcast: make(chan trace.Trace),
 		changes:   make(chan struct{}, 1),
+		done:      make(chan struct{}),
 	}
 
-	go b.run()
+	go b.run(ctx)
 
 	return b
 }
 
-func (b *Broadcaster) run() {
+func (b *Broadcaster) run(ctx context.Context) {
+	defer close(b.done)
+
 	for {
 		select {
+		case <-ctx.Done():
+			b.mu.Lock()
+			for ch := range b.subs {
+				delete(b.subs, ch)
+				close(ch)
+			}
+			b.count.Store(0)
+			b.mu.Unlock()
+
+			return
+
 		case msg := <-b.broadcast:
 			b.mu.Lock()
 			for ch := range b.subs {
@@ -73,17 +92,28 @@ func (b *Broadcaster) run() {
 	}
 }
 
-// Subscribe registers a new consumer and returns its channel.
+// Subscribe registers a new consumer and returns its channel. If the
+// broadcaster is already shutting down, it returns a closed channel so the
+// caller sees the stream end immediately rather than blocking.
 func (b *Broadcaster) Subscribe() chan trace.Trace {
 	ch := make(chan trace.Trace, 10)
-	b.addSub <- ch
 
-	return ch
+	select {
+	case b.addSub <- ch:
+		return ch
+	case <-b.done:
+		close(ch)
+		return ch
+	}
 }
 
-// Unsubscribe removes a consumer.
+// Unsubscribe removes a consumer. It is a no-op once the broadcaster has shut
+// down, since every subscriber channel was closed then.
 func (b *Broadcaster) Unsubscribe(ch chan trace.Trace) {
-	b.removeSub <- ch
+	select {
+	case b.removeSub <- ch:
+	case <-b.done:
+	}
 }
 
 // Subscribers returns the number of active subscribers.
@@ -116,5 +146,7 @@ func (b *Broadcaster) ProcessTrace(ctx context.Context, t trace.Trace) error {
 		return nil
 	case <-ctx.Done():
 		return ctx.Err()
+	case <-b.done:
+		return context.Canceled
 	}
 }

cmd/compass-sidecar/broadcaster_test.go

@@ -13,7 +13,7 @@ import (
 )
 
 func TestBroadcaster_Subscribe(t *testing.T) {
-	b := NewBroadcaster()
+	b := NewBroadcaster(t.Context())
 
 	ch := b.Subscribe()
 	defer b.Unsubscribe(ch)
@@ -25,7 +25,7 @@ func TestBroadcaster_Subscribe(t *testing.T) {
 }
 
 func TestBroadcaster_MultipleSubscribers(t *testing.T) {
-	b := NewBroadcaster()
+	b := NewBroadcaster(t.Context())
 
 	ch1 := b.Subscribe()
 	ch2 := b.Subscribe()
@@ -42,7 +42,7 @@ func TestBroadcaster_MultipleSubscribers(t *testing.T) {
 }
 
 func TestBroadcaster_Unsubscribe(t *testing.T) {
-	b := NewBroadcaster()
+	b := NewBroadcaster(t.Context())
 
 	ch := b.Subscribe()
 
@@ -58,7 +58,7 @@ func TestBroadcaster_Unsubscribe(t *testing.T) {
 }
 
 func TestBroadcaster_ProcessTrace(t *testing.T) {
-	b := NewBroadcaster()
+	b := NewBroadcaster(t.Context())
 
 	ch := b.Subscribe()
 	defer b.Unsubscribe(ch)
@@ -85,7 +85,7 @@ func TestBroadcaster_ProcessTrace(t *testing.T) {
 }
 
 func TestBroadcaster_Broadcast_MultipleSubscribers(t *testing.T) {
-	b := NewBroadcaster()
+	b := NewBroadcaster(t.Context())
 
 	ch1 := b.Subscribe()
 	ch2 := b.Subscribe()
@@ -115,7 +115,7 @@ func TestBroadcaster_Broadcast_MultipleSubscribers(t *testing.T) {
 }
 
 func TestBroadcaster_ResubscribeAfterEmpty(t *testing.T) {
-	b := NewBroadcaster()
+	b := NewBroadcaster(t.Context())
 
 	ch := b.Subscribe()
 
@@ -150,7 +150,7 @@ func TestBroadcaster_ResubscribeAfterEmpty(t *testing.T) {
 }
 
 func TestBroadcaster_DropsForSlowConsumers(t *testing.T) {
-	b := NewBroadcaster()
+	b := NewBroadcaster(t.Context())
 
 	ch := b.Subscribe()
 	defer b.Unsubscribe(ch)
@@ -171,12 +171,12 @@ func TestBroadcaster_DropsForSlowConsumers(t *testing.T) {
 }
 
 func TestBroadcaster_Initialize(t *testing.T) {
-	b := NewBroadcaster()
+	b := NewBroadcaster(t.Context())
 	assert.NoError(t, b.Initialize())
 }
 
 func TestBroadcaster_ProcessTrace_ContextCancelled(t *testing.T) {
-	b := NewBroadcaster()
+	b := NewBroadcaster(t.Context())
 	// No subscribers, so broadcast channel will block.
 
 	ctx, cancel := context.WithCancel(context.Background())
@@ -191,3 +191,44 @@ func TestBroadcaster_ProcessTrace_ContextCancelled(t *testing.T) {
 	err := b.ProcessTrace(ctx, tr)
 	assert.ErrorIs(t, err, context.Canceled)
 }
+
+func TestBroadcaster_ShutdownClosesSubscribers(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	b := NewBroadcaster(ctx)
+
+	ch := b.Subscribe()
+
+	// Allow subscribe to register.
+	time.Sleep(50 * time.Millisecond)
+	require.Equal(t, 1, b.Subscribers())
+
+	cancel()
+
+	// Shutdown closes every subscriber, so a blocked reader unblocks.
+	select {
+	case _, ok := <-ch:
+		assert.False(t, ok, "subscriber channel should be closed on shutdown")
+	case <-time.After(time.Second):
+		t.Fatal("subscriber was not closed on shutdown")
+	}
+
+	// Subscribe and Unsubscribe must not block once the loop has stopped.
+	done := make(chan struct{})
+
+	go func() {
+		defer close(done)
+
+		late := b.Subscribe()
+
+		_, ok := <-late
+		assert.False(t, ok, "subscribe after shutdown should return a closed channel")
+
+		b.Unsubscribe(late)
+	}()
+
+	select {
+	case <-done:
+	case <-time.After(time.Second):
+		t.Fatal("Subscribe/Unsubscribe blocked after shutdown")
+	}
+}

cmd/compass-sidecar/collector_supervisor_test.go

@@ -90,7 +90,7 @@ func updateMaximum(maximum *atomic.Int64, value int64) {
 }
 
 func TestCollectorSupervisor_RapidSubscriberLifecycle(t *testing.T) {
-	b := NewBroadcaster()
+	b := NewBroadcaster(t.Context())
 
 	started := make(chan struct{}, 128)
 	stopped := make(chan struct{}, 128)
@@ -126,7 +126,7 @@ func TestCollectorSupervisor_RapidSubscriberLifecycle(t *testing.T) {
 }
 
 func TestCollectorSupervisor_OnlyStopsForLastSubscriber(t *testing.T) {
-	b := NewBroadcaster()
+	b := NewBroadcaster(t.Context())
 	started := make(chan struct{}, 2)
 	stopped := make(chan struct{}, 2)
 

cmd/compass-sidecar/main.go

@@ -161,10 +161,10 @@ func main() {
 				return err
 			}
 
-			b := NewBroadcaster()
-
 			eg, ctx := errgroup.WithContext(cmd.Context())
 
+			b := NewBroadcaster(ctx)
+
 			// Loop for http server.
 			eg.Go(func() error {
 				mux := http.NewServeMux()