diff --git a/pkg/service/roommanager_service.go b/pkg/service/roommanager_service.go index d6ead74ad..a9d7a221d 100644 --- a/pkg/service/roommanager_service.go +++ b/pkg/service/roommanager_service.go @@ -24,6 +24,10 @@ import ( // whipSessionNotifyInterval is a var (rather than a const) so tests can shorten it. var whipSessionNotifyInterval = 10 * time.Second +// whipSubscribeWaitInterval is how often a participant's subscriptions are checked while waiting +// for the tracks of a one-shot WHIP request. It is a var so tests can shorten it. +var whipSubscribeWaitInterval = 50 * time.Millisecond + type whipService struct { *RoomManager @@ -43,6 +47,41 @@ func newWhipService(rm *RoomManager) (*whipService, error) { }, nil } +// waitForSubscribedTracks blocks until the participant is subscribed to every listed track. It gives +// up when ctx is done or the participant has closed: a track that never appears (a wrong name, a +// publisher that left) would otherwise be waited on forever. +func (s whipService) waitForSubscribedTracks( + ctx context.Context, + lp types.LocalParticipant, + tracks map[string]*rpc.WHIPCreateRequest_TrackList, +) error { + eg, egCtx := errgroup.WithContext(ctx) + for publisherIdentity, trackList := range tracks { + for _, trackName := range trackList.GetTrackNames() { + eg.Go(func() error { + ticker := time.NewTicker(whipSubscribeWaitInterval) + defer ticker.Stop() + + for { + if lp.IsTrackNameSubscribed(livekit.ParticipantIdentity(publisherIdentity), trackName) { + return nil + } + if lp.IsClosed() { + return ErrParticipantNotFound + } + + select { + case <-egCtx.Done(): + return egCtx.Err() + case <-ticker.C: + } + } + }) + } + } + return eg.Wait() +} + func (s whipService) Create(ctx context.Context, req *rpc.WHIPCreateRequest) (*rpc.WHIPCreateResponse, error) { pi, err := routing.ParticipantInitFromStartSession(req.StartSession, s.RoomManager.currentNode.Region()) if err != nil { @@ -88,20 +127,7 @@ func (s whipService) Create(ctx context.Context, req *rpc.WHIPCreateRequest) (*r // NOTE: this is outside the WHIP spec, but added as a convenience for clients doing // one-shot signalling (i. e. send an offer and get an answer once) to publish and subscribe to // well-known tracks (i. e. remote participant identity and track names are well known) - eg, _ := errgroup.WithContext(ctx) - for publisherIdentity, trackList := range req.SubscribedParticipantTracks { - for _, trackName := range trackList.TrackNames { - eg.Go(func() error { - for { - if lp.IsTrackNameSubscribed(livekit.ParticipantIdentity(publisherIdentity), trackName) { - return nil - } - time.Sleep(50 * time.Millisecond) - } - }) - } - } - err = eg.Wait() + err = s.waitForSubscribedTracks(ctx, lp, req.SubscribedParticipantTracks) if err != nil { lp.GetLogger().Errorw("whip service: could not subscribe to tracks", err) return nil, err diff --git a/pkg/service/roommanager_service_test.go b/pkg/service/roommanager_service_test.go index e0d588080..ecbe3c1ca 100644 --- a/pkg/service/roommanager_service_test.go +++ b/pkg/service/roommanager_service_test.go @@ -123,3 +123,97 @@ func TestWhipSendConnectionNotifySkipsClosedParticipant(t *testing.T) { require.ErrorIs(t, err, ErrParticipantNotFound) require.Zero(t, cli.notifyCount.Load(), "should not issue an RPC for a closed participant") } + +func whipSubscribeWaitTracks(publisher string, names ...string) map[string]*rpc.WHIPCreateRequest_TrackList { + return map[string]*rpc.WHIPCreateRequest_TrackList{publisher: {TrackNames: names}} +} + +func shortenWhipSubscribeWait(t *testing.T) { + orig := whipSubscribeWaitInterval + whipSubscribeWaitInterval = 5 * time.Millisecond + t.Cleanup(func() { whipSubscribeWaitInterval = orig }) +} + +// TestWhipWaitForSubscribedTracksReturnsOnceSubscribed verifies the wait ends when every listed +// track is subscribed. +func TestWhipWaitForSubscribedTracksReturnsOnceSubscribed(t *testing.T) { + shortenWhipSubscribeWait(t) + + var subscribed atomic.Bool + participant := &typesfakes.FakeLocalParticipant{} + participant.IsTrackNameSubscribedStub = func(livekit.ParticipantIdentity, string) bool { return subscribed.Load() } + + s := whipService{} + done := make(chan error, 1) + go func() { + done <- s.waitForSubscribedTracks(context.Background(), participant, whipSubscribeWaitTracks("pub", "mic", "cam")) + }() + + select { + case <-done: + t.Fatal("returned before the tracks were subscribed") + case <-time.After(50 * time.Millisecond): + } + + subscribed.Store(true) + select { + case err := <-done: + require.NoError(t, err) + case <-time.After(time.Second): + t.Fatal("did not return after the tracks were subscribed") + } +} + +// TestWhipWaitForSubscribedTracksStopsOnContextCancel verifies a track that never appears does not +// keep the wait, and its polling goroutines, alive once the request is done. +func TestWhipWaitForSubscribedTracksStopsOnContextCancel(t *testing.T) { + shortenWhipSubscribeWait(t) + + participant := &typesfakes.FakeLocalParticipant{} + participant.IsTrackNameSubscribedReturns(false) + + s := whipService{} + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { + done <- s.waitForSubscribedTracks(ctx, participant, whipSubscribeWaitTracks("pub", "mic")) + }() + + cancel() + select { + case err := <-done: + require.ErrorIs(t, err, context.Canceled) + case <-time.After(time.Second): + t.Fatal("did not stop after the context was cancelled") + } + + // the polling has stopped as well + calls := participant.IsTrackNameSubscribedCallCount() + time.Sleep(50 * time.Millisecond) + require.Equal(t, calls, participant.IsTrackNameSubscribedCallCount()) +} + +// TestWhipWaitForSubscribedTracksStopsWhenParticipantLeaves verifies the wait ends when the +// participant closes before the tracks appear. +func TestWhipWaitForSubscribedTracksStopsWhenParticipantLeaves(t *testing.T) { + shortenWhipSubscribeWait(t) + + var closed atomic.Bool + participant := &typesfakes.FakeLocalParticipant{} + participant.IsTrackNameSubscribedReturns(false) + participant.IsClosedStub = func() bool { return closed.Load() } + + s := whipService{} + done := make(chan error, 1) + go func() { + done <- s.waitForSubscribedTracks(context.Background(), participant, whipSubscribeWaitTracks("pub", "mic")) + }() + + closed.Store(true) + select { + case err := <-done: + require.ErrorIs(t, err, ErrParticipantNotFound) + case <-time.After(time.Second): + t.Fatal("did not stop after the participant left") + } +}