From a98533bc099f86d4ee51e1bccd5b22941c60667d Mon Sep 17 00:00:00 2001 From: Nick Date: Fri, 4 Sep 2026 13:48:53 +1000 Subject: [PATCH] Give the broadcaster a shutdown path The broadcaster's run goroutine looped forever with no way to stop, and Subscribe/Unsubscribe blocked on unbuffered channels handed to it. This was latent while nothing cancelled the process context, but now that a signal cancels it the sidecar's shutdown drives this path: server.Shutdown waits for the streaming handlers to return, and they only return once their subscriber channel closes. Run now takes a context and, on cancellation, closes every subscriber channel and returns, so blocked stream handlers unblock and server.Shutdown completes instead of waiting out its timeout. Subscribe, Unsubscribe and ProcessTrace select on a done channel so they no longer block once the loop has stopped; Subscribe hands back a closed channel after shutdown. The broadcaster is built from the errgroup context, and the tests pass one that ends with the test. Adds a shutdown test asserting subscribers are closed and the API stops blocking, and moves the existing tests onto a cancellable context. --- cmd/compass-sidecar/broadcaster.go | 50 +++++++++++++--- cmd/compass-sidecar/broadcaster_test.go | 59 ++++++++++++++++--- .../collector_supervisor_test.go | 4 +- cmd/compass-sidecar/main.go | 4 +- 4 files changed, 95 insertions(+), 22 deletions(-) diff --git a/cmd/compass-sidecar/broadcaster.go b/cmd/compass-sidecar/broadcaster.go index 7127077..1559363 100644 --- a/cmd/compass-sidecar/broadcaster.go +++ b/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 } } diff --git a/cmd/compass-sidecar/broadcaster_test.go b/cmd/compass-sidecar/broadcaster_test.go index ddfbdca..ed3ff0a 100644 --- a/cmd/compass-sidecar/broadcaster_test.go +++ b/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") + } +} diff --git a/cmd/compass-sidecar/collector_supervisor_test.go b/cmd/compass-sidecar/collector_supervisor_test.go index 701208c..1aa55c4 100644 --- a/cmd/compass-sidecar/collector_supervisor_test.go +++ b/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) diff --git a/cmd/compass-sidecar/main.go b/cmd/compass-sidecar/main.go index 730a2c6..92f24ad 100644 --- a/cmd/compass-sidecar/main.go +++ b/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()