fix: Prevent a full status listener from hanging Close

internal/broadcasters.go

@@ -18,15 +18,45 @@ const subscriberChannelBufferLength = 10
 
 // Broadcaster is our generalized implementation of broadcasters.
 type Broadcaster[V any] struct {
-	subscribers []channelPair[V]
+	// Guarded by lock. Do not hold lock while sending to a subscriber.
+	subscribers []*channelPair[V]
 	lock        sync.RWMutex
 }
 
-// We need to keep track of both the channel we use for sending (stored as a reflect.Value, because Value
-// has methods for sending and closing), and also the channel for receiving.
+// We need to keep track of both the channel we use for sending and the channel for receiving.
 type channelPair[V any] struct {
 	sendCh    chan<- V
 	receiveCh <-chan V
+
+	// Closed when the subscriber is removed. A send to sendCh must stop when done is closed.
+	done chan struct{}
+
+	// Sending to sendCh requires lock (read). Closing sendCh requires lock (write), after done is closed.
+	lock sync.RWMutex
+}
+
+// send sends value to sendCh, unless the subscriber is removed first.
+func (p *channelPair[V]) send(value V) {
+	p.lock.RLock()
+	defer p.lock.RUnlock()
+	// Check that it's not done first, so we don't send on a closed channel.
+	select {
+	case <-p.done:
+		return
+	default:
+	}
+	select {
+	case p.sendCh <- value:
+	case <-p.done:
+	}
+}
+
+// close stops any send in progress and closes sendCh. Call it once per subscriber.
+func (p *channelPair[V]) close() {
+	close(p.done)
+	p.lock.Lock()
+	defer p.lock.Unlock()
+	close(p.sendCh)
 }
 
 // NewBroadcaster creates a Broadcaster that operates on the specified value type.
@@ -38,7 +68,7 @@ func NewBroadcaster[V any]() *Broadcaster[V] {
 func (b *Broadcaster[V]) AddListener() <-chan V {
 	ch := make(chan V, subscriberChannelBufferLength)
 	var receiveCh <-chan V = ch
-	chPair := channelPair[V]{sendCh: ch, receiveCh: receiveCh}
+	chPair := &channelPair[V]{sendCh: ch, receiveCh: receiveCh, done: make(chan struct{})}
 	b.lock.Lock()
 	defer b.lock.Unlock()
 	b.subscribers = append(b.subscribers, chPair)
@@ -48,18 +78,22 @@ func (b *Broadcaster[V]) AddListener() <-chan V {
 // RemoveListener removes a subscriber. The parameter is the same channel that was returned by
 // AddListener.
 func (b *Broadcaster[V]) RemoveListener(ch <-chan V) {
+	var removed []*channelPair[V]
 	b.lock.Lock()
-	defer b.lock.Unlock()
-	b.subscribers = slices.DeleteFunc(b.subscribers, func(pair channelPair[V]) bool {
+	b.subscribers = slices.DeleteFunc(b.subscribers, func(pair *channelPair[V]) bool {
 		// The following equality test is the reason why we have to store both the sendCh (chan X) and
 		// the receiveCh (<-chan X) for each subscriber; "s.sendCh == ch" would not be true because
 		// they're of two different types.
 		if pair.receiveCh == ch {
-			close(pair.sendCh)
+			removed = append(removed, pair)
 			return true
 		}
 		return false
 	})
+	b.lock.Unlock()
+	for _, pair := range removed {
+		pair.close()
+	}
 }
 
 // HasListeners returns true if there are any current subscribers.
@@ -72,18 +106,20 @@ func (b *Broadcaster[V]) HasListeners() bool {
 // Broadcast broadcasts a value to all current subscribers.
 func (b *Broadcaster[V]) Broadcast(value V) {
 	b.lock.RLock()
-	defer b.lock.RUnlock()
-	for _, ch := range b.subscribers {
-		ch.sendCh <- value
+	subscribers := slices.Clone(b.subscribers)
+	b.lock.RUnlock()
+	for _, pair := range subscribers {
+		pair.send(value)
 	}
 }
 
 // Close closes all current subscriber channels.
 func (b *Broadcaster[V]) Close() {
 	b.lock.Lock()
-	defer b.lock.Unlock()
-	for _, s := range b.subscribers {
-		close(s.sendCh)
-	}
+	subscribers := b.subscribers
 	b.subscribers = nil
+	b.lock.Unlock()
+	for _, pair := range subscribers {
+		pair.close()
+	}
 }

internal/broadcasters_test.go

@@ -48,6 +48,28 @@ func testBroadcasterGenerically[V any](t *testing.T, broadcasterFactory func() *
 		})
 	})
 
+	t.Run("broadcast to full listener delivers every value", func(t *testing.T) {
+		withBroadcaster(t, func(b *Broadcaster[V]) {
+			ch := b.AddListener()
+			values := make([]V, subscriberChannelBufferLength+1)
+			for i := range values {
+				values[i] = valueFactory()
+			}
+
+			// Broadcast one more value than the channel can hold.
+			go func() {
+				for _, v := range values {
+					b.Broadcast(v)
+				}
+			}()
+
+			// The listener should receive every value in order.
+			for _, v := range values {
+				assert.Equal(t, v, th.RequireValue(t, ch, timeout))
+			}
+		})
+	})
+
 	t.Run("unregister subscriber", func(t *testing.T) {
 		withBroadcaster(t, func(b *Broadcaster[V]) {
 			ch1 := b.AddListener()
@@ -63,6 +85,95 @@ func testBroadcasterGenerically[V any](t *testing.T, broadcasterFactory func() *
 		})
 	})
 
+	unblockTests := []struct {
+		name    string
+		release func(b *Broadcaster[V], ch <-chan V)
+	}{
+		{"remove listener unblocks broadcast to full listener",
+			func(b *Broadcaster[V], ch <-chan V) { b.RemoveListener(ch) }},
+		{"close unblocks broadcast to full listener",
+			func(b *Broadcaster[V], _ <-chan V) { b.Close() }},
+	}
+	for _, tt := range unblockTests {
+		t.Run(tt.name, func(t *testing.T) {
+			withBroadcaster(t, func(b *Broadcaster[V]) {
+				ch := b.AddListener()
+				for i := 0; i < subscriberChannelBufferLength; i++ {
+					b.Broadcast(valueFactory())
+				}
+
+				broadcastDone := make(chan struct{})
+				go func() {
+					b.Broadcast(valueFactory())
+					close(broadcastDone)
+				}()
+				// Give Broadcast time to block on the full channel.
+				time.Sleep(10 * time.Millisecond)
+
+				// Release the listener without reading from it.
+				releaseDone := make(chan struct{})
+				go func() {
+					tt.release(b, ch)
+					close(releaseDone)
+				}()
+
+				// Both calls should return.
+				ok := th.AssertChannelClosed(t, broadcastDone, timeout) &&
+					th.AssertChannelClosed(t, releaseDone, timeout)
+				if !ok {
+					<-ch
+				}
+			})
+		})
+	}
+
+	concurrentCallTests := []struct {
+		name string
+		call func(b *Broadcaster[V], other <-chan V)
+	}{
+		{"close unblocks broadcast to full listener while a listener is added",
+			func(b *Broadcaster[V], _ <-chan V) { b.AddListener() }},
+		{"close unblocks broadcast to full listener while another listener is removed",
+			func(b *Broadcaster[V], other <-chan V) { b.RemoveListener(other) }},
+	}
+	for _, tt := range concurrentCallTests {
+		t.Run(tt.name, func(t *testing.T) {
+			withBroadcaster(t, func(b *Broadcaster[V]) {
+				other := b.AddListener()
+				ch := b.AddListener()
+				for i := 0; i < subscriberChannelBufferLength; i++ {
+					b.Broadcast(valueFactory())
+				}
+
+				go b.Broadcast(valueFactory())
+				// Give Broadcast time to block on a full channel.
+				time.Sleep(10 * time.Millisecond)
+				go tt.call(b, other)
+				// Give the call time to start.
+				time.Sleep(10 * time.Millisecond)
+
+				// Close without reading from either listener.
+				closeDone := make(chan struct{})
+				go func() {
+					b.Close()
+					close(closeDone)
+				}()
+
+				// Close should return.
+				if !th.AssertChannelClosed(t, closeDone, timeout) {
+					go func() {
+						for range ch {
+						}
+					}()
+					go func() {
+						for range other {
+						}
+					}()
+				}
+			})
+		})
+	}
+
 	t.Run("hasListeners", func(t *testing.T) {
 		withBroadcaster(t, func(b *Broadcaster[V]) {
 			assert.False(t, b.HasListeners())

internal/datasource/data_source_status_provider_impl_test.go

@@ -168,5 +168,18 @@ func TestDataSourceStatusProviderImpl(t *testing.T) {
 				assert.True(t, time.Now().Sub(timeStart) < 500*time.Millisecond)
 			})
 		})
+
+		t.Run("ends if status broadcaster is closed", func(t *testing.T) {
+			dataSourceStatusProviderImplTest(func(p dataSourceStatusProviderImplTestParams) {
+				go func() {
+					<-time.After(100 * time.Millisecond)
+					p.dataSourceUpdates.dataSourceStatusBroadcaster.Close()
+				}()
+				timeStart := time.Now()
+				success := p.dataSourceStatusProvider.WaitFor(interfaces.DataSourceStateValid, 500*time.Millisecond)
+				assert.False(t, success)
+				assert.True(t, time.Now().Sub(timeStart) < 500*time.Millisecond)
+			})
+		})
 	})
 }

internal/datasource/data_source_update_sink_impl.go

@@ -216,15 +216,7 @@ func (d *DataSourceUpdateSinkImpl) waitFor(desiredState intf.DataSourceState, ti
 	}
 
 	statusCh := d.dataSourceStatusBroadcaster.AddListener()
-	defer func() {
-		// A broadcast that waits on this full channel blocks RemoveListener. Drain the
-		// channel until RemoveListener closes it, so that the removal always completes.
-		go func() {
-			for range statusCh {
-			}
-		}()
-		d.dataSourceStatusBroadcaster.RemoveListener(statusCh)
-	}()
+	defer d.dataSourceStatusBroadcaster.RemoveListener(statusCh)
 	d.lock.Unlock()
 
 	var deadline <-chan time.Time
@@ -234,7 +226,10 @@ func (d *DataSourceUpdateSinkImpl) waitFor(desiredState intf.DataSourceState, ti
 
 	for {
 		select {
-		case newStatus := <-statusCh:
+		case newStatus, ok := <-statusCh:
+			if !ok {
+				return false
+			}
 			if newStatus.State == desiredState {
 				return true
 			}

internal/datasystem/fdv2_datasystem.go

@@ -645,7 +645,8 @@ func (f *FDv2) Stop() error {
 		f.cancel()
 	}
 	f.wg.Wait()
-	f.UpdateStatus(interfaces.DataSourceStateOff, interfaces.DataSourceErrorInfo{})
+	// Do not broadcast Off. Status listeners learn of shutdown when their channels close.
+	_, _ = f.maybeUpdateStatus(interfaces.DataSourceStateOff, interfaces.DataSourceErrorInfo{})
 	_ = f.store.Close()
 	f.broadcasters.Close()
 
@@ -718,6 +719,14 @@ func (f *FDv2) Offline() bool {
 
 //nolint:revive // DataSourceStatusReporter method.
 func (f *FDv2) UpdateStatus(state interfaces.DataSourceState, err interfaces.DataSourceErrorInfo) {
+	if status, changed := f.maybeUpdateStatus(state, err); changed {
+		f.broadcasters.dataSourceStatus.Broadcast(status)
+	}
+}
+
+func (f *FDv2) maybeUpdateStatus(
+	state interfaces.DataSourceState, err interfaces.DataSourceErrorInfo,
+) (interfaces.DataSourceStatus, bool) {
 	f.mu.Lock()
 	defer f.mu.Unlock()
 
@@ -733,9 +742,7 @@ func (f *FDv2) UpdateStatus(state interfaces.DataSourceState, err interfaces.Dat
 		changed = true
 	}
 
-	if changed {
-		f.broadcasters.dataSourceStatus.Broadcast(f.status)
-	}
+	return f.status, changed
 }
 
 // restartStatusClock sets the time the current status has been held since to now.
@@ -769,15 +776,7 @@ func (d *dataStatusProvider) RemoveStatusListener(listener <-chan interfaces.Dat
 
 func (d *dataStatusProvider) WaitFor(desiredState interfaces.DataSourceState, timeout time.Duration) bool {
 	ch := d.AddStatusListener()
-	defer func() {
-		// A broadcast that waits on this full channel blocks RemoveStatusListener. Drain the
-		// channel until the removal closes it, so that the removal always completes.
-		go func() {
-			for range ch {
-			}
-		}()
-		d.RemoveStatusListener(ch)
-	}()
+	defer d.RemoveStatusListener(ch)
 
 	switch d.system.getStatus().State {
 	case desiredState:
@@ -790,7 +789,10 @@ func (d *dataStatusProvider) WaitFor(desiredState interfaces.DataSourceState, ti
 
 	for {
 		select {
-		case status := <-ch:
+		case status, ok := <-ch:
+			if !ok {
+				return false
+			}
 			if status.State == desiredState {
 				return true
 			}