Do not add down tracks after receiver close (#4857)

* Do not add down tracks after receiver close

* Fold close race tests into existing test files

* Retrigger CI

pkg/sfu/receiver_base.go

@@ -197,7 +197,7 @@ func (r *ReceiverBase) Close(reason string, clearBuffers bool) {
 	}
 	r.streamTrackerManager.Close()
 
-	closeTrackSenders(r.downTrackSpreader.ResetAndGetDownTracks())
+	closeTrackSenders(r.downTrackSpreader.CloseAndGetDownTracks())
 
 	if rt := r.loadREDTransformer(); rt != nil {
 		rt.Close()
@@ -491,7 +491,9 @@ func (r *ReceiverBase) AddDownTrack(track TrackSender) error {
 	track.UpTrackMaxPublishedLayerChange(r.streamTrackerManager.GetMaxPublishedLayer())
 	track.UpTrackMaxTemporalLayerSeenChange(r.streamTrackerManager.GetMaxTemporalLayerSeen())
 
-	r.downTrackSpreader.Store(track)
+	if !r.downTrackSpreader.TryStore(track) {
+		return ErrReceiverClosed
+	}
 	r.params.Logger.Debugw("downtrack added", "subscriberID", track.SubscriberID())
 	return nil
 }

pkg/sfu/receiver_test.go

@@ -23,8 +23,13 @@ import (
 	"testing"
 
 	"github.com/gammazero/workerpool"
+	"github.com/pion/webrtc/v4"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"go.uber.org/atomic"
+
+	"github.com/livekit/protocol/livekit"
+	"github.com/livekit/protocol/logger"
 )
 
 func TestWebRTCReceiver_OnCloseHandler(t *testing.T) {
@@ -52,6 +57,77 @@ func TestWebRTCReceiver_OnCloseHandler(t *testing.T) {
 	}
 }
 
+type receiverCloseTestSender struct {
+	TrackSender
+	subscriberID                   livekit.ParticipantID
+	upTrackMaxPublishedLayerChange func(int32)
+}
+
+func (s *receiverCloseTestSender) SubscriberID() livekit.ParticipantID {
+	return s.subscriberID
+}
+
+func (s *receiverCloseTestSender) UpTrackMaxPublishedLayerChange(layer int32) {
+	if s.upTrackMaxPublishedLayerChange != nil {
+		s.upTrackMaxPublishedLayerChange(layer)
+	}
+}
+
+func (s *receiverCloseTestSender) UpTrackMaxTemporalLayerSeenChange(int32) {}
+
+func TestReceiverBaseAddDownTrackReturnsClosedWhenCloseWinsRace(t *testing.T) {
+	receiver := NewReceiverBase(
+		ReceiverBaseParams{
+			TrackID:  "track",
+			StreamID: "stream",
+			Kind:     webrtc.RTPCodecTypeAudio,
+			Codec: webrtc.RTPCodecParameters{
+				RTPCodecCapability: webrtc.RTPCodecCapability{
+					MimeType:  webrtc.MimeTypeOpus,
+					ClockRate: 48000,
+					Channels:  2,
+				},
+				PayloadType: 111,
+			},
+			Logger: logger.GetLogger(),
+		},
+		&livekit.TrackInfo{
+			Sid:    "track",
+			Type:   livekit.TrackType_AUDIO,
+			Source: livekit.TrackSource_MICROPHONE,
+		},
+		ReceiverCodecStateNormal,
+	)
+
+	enteredAdd := make(chan struct{})
+	resumeAdd := make(chan struct{})
+	track := &receiverCloseTestSender{
+		subscriberID: "subscriber",
+		upTrackMaxPublishedLayerChange: func(int32) {
+			close(enteredAdd)
+			<-resumeAdd
+		},
+	}
+
+	errCh := make(chan error, 1)
+	go func() {
+		errCh <- receiver.AddDownTrack(track)
+	}()
+
+	// AddDownTrack has passed its initial IsClosed check and is paused
+	// before reaching the final TryStore.
+	<-enteredAdd
+
+	// Close terminally drains/closes the spreader before AddDownTrack resumes.
+	receiver.Close("test", false)
+	close(resumeAdd)
+
+	err := <-errCh
+	require.ErrorIs(t, err, ErrReceiverClosed)
+	require.True(t, receiver.IsClosed())
+	require.Empty(t, receiver.GetDownTracks())
+}
+
 func BenchmarkWriteRTP(b *testing.B) {
 	cases := []int{1, 2, 5, 10, 100, 250, 500}
 	workers := runtime.NumCPU()

pkg/sfu/redprimaryreceiver.go

@@ -154,7 +154,9 @@ func (r *RedPrimaryReceiver) AddDownTrack(track TrackSender) error {
 		r.logger.Infow("subscriberID already exists, replacing downtrack", "subscriberID", track.SubscriberID())
 	}
 
-	r.downTrackSpreader.Store(track)
+	if !r.downTrackSpreader.TryStore(track) {
+		return ErrReceiverClosed
+	}
 	r.logger.Debugw("red primary receiver downtrack added", "subscriberID", track.SubscriberID())
 	return nil
 }
@@ -194,7 +196,7 @@ func (r *RedPrimaryReceiver) CanClose() bool {
 
 func (r *RedPrimaryReceiver) Close() {
 	r.closed.Store(true)
-	closeTrackSenders(r.downTrackSpreader.ResetAndGetDownTracks())
+	closeTrackSenders(r.downTrackSpreader.CloseAndGetDownTracks())
 }
 
 func (r *RedPrimaryReceiver) ReadRTP(buf []byte, layer uint8, esn uint64) (int, error) {

pkg/sfu/redreceiver.go

@@ -118,7 +118,9 @@ func (r *RedReceiver) AddDownTrack(track TrackSender) error {
 		r.logger.Infow("subscriberID already exists, replacing downtrack", "subscriberID", track.SubscriberID())
 	}
 
-	r.downTrackSpreader.Store(track)
+	if !r.downTrackSpreader.TryStore(track) {
+		return ErrReceiverClosed
+	}
 	r.logger.Debugw("red receiver downtrack added", "subscriberID", track.SubscriberID())
 	return nil
 }
@@ -158,7 +160,7 @@ func (r *RedReceiver) IsClosed() bool {
 
 func (r *RedReceiver) Close() {
 	r.closed.Store(true)
-	closeTrackSenders(r.downTrackSpreader.ResetAndGetDownTracks())
+	closeTrackSenders(r.downTrackSpreader.CloseAndGetDownTracks())
 }
 
 func (r *RedReceiver) ReadRTP(buf []byte, layer uint8, esn uint64) (int, error) {

pkg/sfu/utils/downtrackspreader.go

@@ -39,6 +39,7 @@ type DownTrackSpreader[T sender] struct {
 	downTrackMu      sync.RWMutex
 	downTracks       map[livekit.ParticipantID]T
 	downTracksShadow []T
+	closed           bool
 }
 
 func NewDownTrackSpreader[T sender](params DownTrackSpreaderParams) *DownTrackSpreader[T] {
@@ -60,20 +61,35 @@ func (d *DownTrackSpreader[T]) ResetAndGetDownTracks() []T {
 	d.downTrackMu.Lock()
 	defer d.downTrackMu.Unlock()
 
-	downTracks := d.downTracksShadow
-
-	d.downTracks = make(map[livekit.ParticipantID]T)
-	d.downTracksShadow = nil
-
-	return downTracks
+	return d.resetAndGetDownTracksLocked()
 }
 
-func (d *DownTrackSpreader[T]) Store(sender T) {
+// CloseAndGetDownTracks terminally closes the spreader and drains all senders.
+// Once it returns, TryStore will reject every subsequent sender.
+func (d *DownTrackSpreader[T]) CloseAndGetDownTracks() []T {
 	d.downTrackMu.Lock()
 	defer d.downTrackMu.Unlock()
 
+	d.closed = true
+	return d.resetAndGetDownTracksLocked()
+}
+
+func (d *DownTrackSpreader[T]) Store(sender T) {
+	_ = d.TryStore(sender)
+}
+
+// TryStore stores sender unless the spreader has been terminally closed.
+func (d *DownTrackSpreader[T]) TryStore(sender T) bool {
+	d.downTrackMu.Lock()
+	defer d.downTrackMu.Unlock()
+
+	if d.closed {
+		return false
+	}
+
 	d.downTracks[sender.SubscriberID()] = sender
 	d.shadowDownTracks()
+	return true
 }
 
 func (d *DownTrackSpreader[T]) Free(subscriberID livekit.ParticipantID) {
@@ -120,6 +136,13 @@ func (d *DownTrackSpreader[T]) DownTrackCount() int {
 	return len(d.downTracksShadow)
 }
 
+func (d *DownTrackSpreader[T]) resetAndGetDownTracksLocked() []T {
+	downTracks := d.downTracksShadow
+	d.downTracks = make(map[livekit.ParticipantID]T)
+	d.downTracksShadow = nil
+	return downTracks
+}
+
 func (d *DownTrackSpreader[T]) shadowDownTracks() {
 	d.downTracksShadow = make([]T, 0, len(d.downTracks))
 	for _, dt := range d.downTracks {

pkg/sfu/utils/downtrackspreader_test.go

@@ -82,3 +82,29 @@ func TestBroadcastRTPEmpty(t *testing.T) {
 	d := NewDownTrackSpreader[*testSender](DownTrackSpreaderParams{Threshold: 20})
 	require.EqualValues(t, 0, BroadcastRTP(d, &testPacket{}, 0))
 }
+
+func TestDownTrackSpreaderResetRemainsReusable(t *testing.T) {
+	spreader := NewDownTrackSpreader[*testSender](DownTrackSpreaderParams{})
+	first := &testSender{id: "first"}
+	second := &testSender{id: "second"}
+
+	spreader.Store(first)
+	require.Equal(t, []*testSender{first}, spreader.ResetAndGetDownTracks())
+	require.Zero(t, spreader.DownTrackCount())
+
+	spreader.Store(second)
+	require.Equal(t, []*testSender{second}, spreader.GetDownTracks())
+}
+
+func TestDownTrackSpreaderRejectsStoreAfterClose(t *testing.T) {
+	spreader := NewDownTrackSpreader[*testSender](DownTrackSpreaderParams{})
+	first := &testSender{id: "first"}
+	late := &testSender{id: "late"}
+
+	spreader.Store(first)
+	require.Equal(t, []*testSender{first}, spreader.CloseAndGetDownTracks())
+	require.Zero(t, spreader.DownTrackCount())
+
+	require.False(t, spreader.TryStore(late))
+	require.Zero(t, spreader.DownTrackCount())
+}