diff --git a/pkg/rtc/subscriptionmanager.go b/pkg/rtc/subscriptionmanager.go index b5edddca1..52f52a1d0 100644 --- a/pkg/rtc/subscriptionmanager.go +++ b/pkg/rtc/subscriptionmanager.go @@ -80,6 +80,10 @@ type SubscriptionManager struct { closeCh chan struct{} doneCh chan struct{} + // set when a queued reconcile is dropped, the worker then reconciles everything + reconcileOverflow atomic.Bool + reconcileDataTrackOverflow atomic.Bool + onSubscribeStatusChanged func(publisherID livekit.ParticipantID, subscribed bool) dataTrackSubscriptions map[livekit.TrackID]*dataTrackSubscription @@ -706,7 +710,8 @@ func (m *SubscriptionManager) queueReconcile(trackID livekit.TrackID) { select { case m.reconcileCh <- trackID: default: - // queue is full, will reconcile based on timer + // queue is full, the worker reconciles everything after its next queued item + m.reconcileOverflow.Store(true) } } @@ -714,7 +719,7 @@ func (m *SubscriptionManager) queueReconcileDataTrack(trackID livekit.TrackID) { select { case m.reconcileDataTrackCh <- trackID: default: - // queue is full, will reconcile based on timer + m.reconcileDataTrackOverflow.Store(true) } } @@ -742,6 +747,9 @@ func (m *SubscriptionManager) reconcileWorker() { } else { m.reconcileSubscriptions() } + if m.reconcileOverflow.Swap(false) { + m.reconcileSubscriptions() + } case trackID := <-m.reconcileDataTrackCh: m.lock.Lock() @@ -752,6 +760,9 @@ func (m *SubscriptionManager) reconcileWorker() { } else { m.reconcileDataTrackSubscriptions() } + if m.reconcileDataTrackOverflow.Swap(false) { + m.reconcileDataTrackSubscriptions() + } } } } diff --git a/pkg/rtc/subscriptionmanager_test.go b/pkg/rtc/subscriptionmanager_test.go index b5305e32d..e184b7fae 100644 --- a/pkg/rtc/subscriptionmanager_test.go +++ b/pkg/rtc/subscriptionmanager_test.go @@ -17,6 +17,7 @@ package rtc import ( + "fmt" "sync" "testing" "time" @@ -488,6 +489,91 @@ func TestSettingsKeptForSubscribe(t *testing.T) { }) } +func TestSubscribeBurstOverflowsQueue(t *testing.T) { + // a dropped reconcile must not wait for the ticker + saved := reconcileInterval + reconcileInterval = time.Hour + defer func() { reconcileInterval = saved }() + + t.Run("media", func(t *testing.T) { + sm := newTestSubscriptionManager() + defer sm.Close(false) + + // the first subscribe holds the worker until all are queued, so the queue overflows + release := make(chan struct{}) + var holdOnce sync.Once + sm.params.TrackResolver = func(_ types.LocalParticipant, trackID livekit.TrackID) types.MediaResolverResult { + mt := &typesfakes.FakeMediaTrack{} + mt.IDReturns(trackID) + mt.AddSubscriberCalls(func(types.LocalParticipant) (types.SubscribedTrack, error) { + holdOnce.Do(func() { <-release }) + st := &typesfakes.FakeSubscribedTrack{} + st.IDReturns(trackID) + st.MediaTrackReturns(mt) + return st, nil + }) + return types.MediaResolverResult{ + Track: mt, + HasPermission: true, + PublisherID: "pubID", + PublisherIdentity: "pub", + TrackChangedNotifier: utils.NewChangeNotifier(), + TrackRemovedNotifier: utils.NewChangeNotifier(), + } + } + + numTracks := cap(sm.reconcileCh) + 10 + for i := range numTracks { + sm.SubscribeToTrack(livekit.TrackID(fmt.Sprintf("TR_%d", i)), false) + } + close(release) + + require.Eventually(t, func() bool { + sm.lock.RLock() + defer sm.lock.RUnlock() + for _, s := range sm.subscriptions { + if s.needsSubscribe() { + return false + } + } + return len(sm.subscriptions) == numTracks + }, subSettleTimeout, subCheckInterval, "every track should be subscribed") + }) + + t.Run("data", func(t *testing.T) { + sm := newTestSubscriptionManager() + defer sm.Close(false) + resolver := newTestDataTrackResolver(true, true, "pub", "pubID") + sm.params.DataTrackResolver = resolver.Resolve + + release := make(chan struct{}) + var holdOnce sync.Once + resolver.dataTrack.AddSubscriberCalls(func(types.LocalParticipant) (types.DataDownTrack, error) { + holdOnce.Do(func() { <-release }) + ddt := &typesfakes.FakeDataDownTrack{} + ddt.PublishDataTrackReturns(resolver.dataTrack) + return ddt, nil + }) + + numTracks := cap(sm.reconcileDataTrackCh) + 10 + for i := range numTracks { + sm.SubscribeToDataTrack(livekit.TrackID(fmt.Sprintf("DTR_%d", i))) + } + close(release) + + require.Eventually(t, func() bool { + sm.lock.RLock() + defer sm.lock.RUnlock() + for _, s := range sm.dataTrackSubscriptions { + if s.needsSubscribe() { + return false + } + } + return len(sm.dataTrackSubscriptions) == numTracks + }, subSettleTimeout, subCheckInterval, "every data track should be subscribed") + }) +} + func TestSubscriptionLimits(t *testing.T) { sm := newTestSubscriptionManagerWithParams(testSubscriptionParams{ SubscriptionLimitAudio: 1,