diff --git a/pkg/sfu/receiver_base.go b/pkg/sfu/receiver_base.go index 617d7fc64..678c124c3 100644 --- a/pkg/sfu/receiver_base.go +++ b/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 } diff --git a/pkg/sfu/receiver_test.go b/pkg/sfu/receiver_test.go index 3cfde0768..27e53ae78 100644 --- a/pkg/sfu/receiver_test.go +++ b/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() diff --git a/pkg/sfu/redprimaryreceiver.go b/pkg/sfu/redprimaryreceiver.go index a34ef8586..fdef27c81 100644 --- a/pkg/sfu/redprimaryreceiver.go +++ b/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) { diff --git a/pkg/sfu/redreceiver.go b/pkg/sfu/redreceiver.go index 4977878bf..ad0c54b8a 100644 --- a/pkg/sfu/redreceiver.go +++ b/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) { diff --git a/pkg/sfu/utils/downtrackspreader.go b/pkg/sfu/utils/downtrackspreader.go index 29ba07f2d..12e309cb5 100644 --- a/pkg/sfu/utils/downtrackspreader.go +++ b/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 { diff --git a/pkg/sfu/utils/downtrackspreader_test.go b/pkg/sfu/utils/downtrackspreader_test.go index 1de2d96ce..11557183e 100644 --- a/pkg/sfu/utils/downtrackspreader_test.go +++ b/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()) +}