diff --git a/pkg/rtc/datatrack.go b/pkg/rtc/datatrack.go index e5fd98a7c..9e7371f07 100644 --- a/pkg/rtc/datatrack.go +++ b/pkg/rtc/datatrack.go @@ -24,6 +24,7 @@ import ( "github.com/livekit/livekit-server/pkg/rtc/types" sfuutils "github.com/livekit/livekit-server/pkg/sfu/utils" + "github.com/livekit/livekit-server/pkg/utils" "github.com/livekit/protocol/datatrack" "github.com/livekit/protocol/livekit" "github.com/livekit/protocol/logger" @@ -40,6 +41,7 @@ type DataTrackParams struct { ParticipantID func() livekit.ParticipantID ParticipantIdentity livekit.ParticipantIdentity BytesTrackStats *BytesTrackStats + OnDemandChanged func(hasSubscribers bool) } type subscribedDataTrack struct { @@ -52,9 +54,14 @@ type DataTrack struct { logger logger.Logger - lock sync.Mutex - dti *livekit.DataTrackInfo - subscribedTracks map[livekit.ParticipantID]subscribedDataTrack + lock sync.Mutex + dti *livekit.DataTrackInfo + subscribedTracks map[livekit.ParticipantID]subscribedDataTrack + subscriberNodeEnables map[livekit.NodeID]bool + demandStarted bool + hasSubscribers bool + + notifyOpsQueue *utils.OpsQueue downTrackSpreader *sfuutils.DownTrackSpreader[types.DataTrackSender] @@ -65,11 +72,19 @@ type DataTrack struct { func NewDataTrack(params DataTrackParams, dti *livekit.DataTrackInfo) *DataTrack { d := &DataTrack{ - params: params, - dti: dti, - subscribedTracks: make(map[livekit.ParticipantID]subscribedDataTrack), + params: params, + dti: dti, + subscribedTracks: make(map[livekit.ParticipantID]subscribedDataTrack), + subscriberNodeEnables: make(map[livekit.NodeID]bool), } d.logger = params.Logger.WithValues("name", d.Name(), "handle", dti.PubHandle) + d.notifyOpsQueue = utils.NewOpsQueue(utils.OpsQueueParams{ + Name: "data-track-demand-notify", + MinSize: 4, + FlushOnStop: true, + Logger: d.logger, + }) + d.notifyOpsQueue.Start() d.downTrackSpreader = sfuutils.NewDownTrackSpreader[types.DataTrackSender](sfuutils.DownTrackSpreaderParams{ Threshold: 20, Logger: d.logger, @@ -82,6 +97,7 @@ func NewDataTrack(params DataTrackParams, dti *livekit.DataTrackInfo) *DataTrack func (d *DataTrack) Close() { d.logger.Infow("closing data track") d.closed.Break() + d.notifyOpsQueue.Stop() d.stats.Close() if d.params.BytesTrackStats != nil { @@ -147,6 +163,7 @@ func (d *DataTrack) AddSubscriber(sub types.LocalParticipant) (types.DataDownTra subscriber: sub, dataDownTrack: dataDownTrack, } + d.updateDemandLocked(false) return dataDownTrack, nil } @@ -154,6 +171,7 @@ func (d *DataTrack) RemoveSubscriber(subID livekit.ParticipantID) { d.lock.Lock() subscribedTrack, ok := d.subscribedTracks[subID] delete(d.subscribedTracks, subID) + d.updateDemandLocked(false) d.lock.Unlock() if ok { @@ -161,6 +179,41 @@ func (d *DataTrack) RemoveSubscriber(subID livekit.ParticipantID) { } } +func (d *DataTrack) NotifySubscriptionNode(nodeID livekit.NodeID, enabled bool) { + d.lock.Lock() + defer d.lock.Unlock() + + if !enabled { + delete(d.subscriberNodeEnables, nodeID) + } else { + d.subscriberNodeEnables[nodeID] = true + } + d.updateDemandLocked(false) +} + +func (d *DataTrack) StartDemandUpdates() { + d.lock.Lock() + defer d.lock.Unlock() + + d.demandStarted = true + d.updateDemandLocked(true) +} + +func (d *DataTrack) updateDemandLocked(force bool) { + hasSubscribers := len(d.subscribedTracks) != 0 || len(d.subscriberNodeEnables) != 0 + if !d.demandStarted || (hasSubscribers == d.hasSubscribers && !force) { + return + } + + d.hasSubscribers = hasSubscribers + d.logger.Infow("data track demand changed", "hasSubscribers", hasSubscribers, "force", force) + if d.params.OnDemandChanged != nil { + d.notifyOpsQueue.Enqueue(func() { + d.params.OnDemandChanged(hasSubscribers) + }) + } +} + func (d *DataTrack) IsSubscriber(subID livekit.ParticipantID) bool { d.lock.Lock() defer d.lock.Unlock() diff --git a/pkg/rtc/datatrack_test.go b/pkg/rtc/datatrack_test.go index f2938a1f2..46341adf8 100644 --- a/pkg/rtc/datatrack_test.go +++ b/pkg/rtc/datatrack_test.go @@ -15,7 +15,10 @@ package rtc import ( + "slices" + "sync" "testing" + "time" "github.com/stretchr/testify/require" @@ -69,3 +72,53 @@ func TestDataTrackRevokeDisallowedSubscribers(t *testing.T) { require.False(t, dt.IsSubscriber(disallowed.ID())) require.True(t, dt.IsSubscriber(recorder.ID())) } + +func TestDataTrackDemand(t *testing.T) { + var ( + lock sync.Mutex + reported []bool + ) + requireReported := func(expected ...bool) { + require.Eventually(t, func() bool { + lock.Lock() + defer lock.Unlock() + return slices.Equal(expected, reported) + }, time.Second, 10*time.Millisecond) + } + + dt := NewDataTrack( + DataTrackParams{ + Logger: logger.GetLogger(), + ParticipantID: func() livekit.ParticipantID { return "pubID" }, + ParticipantIdentity: "pub", + OnDemandChanged: func(hasSubscribers bool) { + lock.Lock() + reported = append(reported, hasSubscribers) + lock.Unlock() + }, + }, + &livekit.DataTrackInfo{ + PubHandle: 1, + Sid: "DTR_test", + Name: "test", + }, + ) + defer dt.Close() + + sub := newTestDataTrackSubscriber("subID", "sub", false) + _, err := dt.AddSubscriber(sub) + require.NoError(t, err) + dt.RemoveSubscriber(sub.ID()) + + dt.StartDemandUpdates() + requireReported(false) + + _, err = dt.AddSubscriber(sub) + require.NoError(t, err) + dt.NotifySubscriptionNode("nodeB", true) + dt.RemoveSubscriber(sub.ID()) + requireReported(false, true) + + dt.NotifySubscriptionNode("nodeB", false) + requireReported(false, true, false) +} diff --git a/pkg/rtc/types/interfaces.go b/pkg/rtc/types/interfaces.go index 1d7734f1d..2d1fa9519 100644 --- a/pkg/rtc/types/interfaces.go +++ b/pkg/rtc/types/interfaces.go @@ -932,6 +932,8 @@ type DataTrack interface { IsSubscriber(subID livekit.ParticipantID) bool RevokeDisallowedSubscribers(allowedSubscriberIdentities []livekit.ParticipantIdentity) []livekit.ParticipantIdentity + NotifySubscriptionNode(nodeID livekit.NodeID, enabled bool) + AddDataDownTrack(sender DataTrackSender) error DeleteDataDownTrack(subscriberID livekit.ParticipantID) diff --git a/pkg/rtc/types/typesfakes/fake_data_track.go b/pkg/rtc/types/typesfakes/fake_data_track.go index ddd512d92..443402746 100644 --- a/pkg/rtc/types/typesfakes/fake_data_track.go +++ b/pkg/rtc/types/typesfakes/fake_data_track.go @@ -81,6 +81,12 @@ type FakeDataTrack struct { nameReturnsOnCall map[int]struct { result1 string } + NotifySubscriptionNodeStub func(livekit.NodeID, bool) + notifySubscriptionNodeMutex sync.RWMutex + notifySubscriptionNodeArgsForCall []struct { + arg1 livekit.NodeID + arg2 bool + } PubHandleStub func() uint16 pubHandleMutex sync.RWMutex pubHandleArgsForCall []struct { @@ -528,6 +534,39 @@ func (fake *FakeDataTrack) NameReturnsOnCall(i int, result1 string) { }{result1} } +func (fake *FakeDataTrack) NotifySubscriptionNode(arg1 livekit.NodeID, arg2 bool) { + fake.notifySubscriptionNodeMutex.Lock() + fake.notifySubscriptionNodeArgsForCall = append(fake.notifySubscriptionNodeArgsForCall, struct { + arg1 livekit.NodeID + arg2 bool + }{arg1, arg2}) + stub := fake.NotifySubscriptionNodeStub + fake.recordInvocation("NotifySubscriptionNode", []interface{}{arg1, arg2}) + fake.notifySubscriptionNodeMutex.Unlock() + if stub != nil { + fake.NotifySubscriptionNodeStub(arg1, arg2) + } +} + +func (fake *FakeDataTrack) NotifySubscriptionNodeCallCount() int { + fake.notifySubscriptionNodeMutex.RLock() + defer fake.notifySubscriptionNodeMutex.RUnlock() + return len(fake.notifySubscriptionNodeArgsForCall) +} + +func (fake *FakeDataTrack) NotifySubscriptionNodeCalls(stub func(livekit.NodeID, bool)) { + fake.notifySubscriptionNodeMutex.Lock() + defer fake.notifySubscriptionNodeMutex.Unlock() + fake.NotifySubscriptionNodeStub = stub +} + +func (fake *FakeDataTrack) NotifySubscriptionNodeArgsForCall(i int) (livekit.NodeID, bool) { + fake.notifySubscriptionNodeMutex.RLock() + defer fake.notifySubscriptionNodeMutex.RUnlock() + argsForCall := fake.notifySubscriptionNodeArgsForCall[i] + return argsForCall.arg1, argsForCall.arg2 +} + func (fake *FakeDataTrack) PubHandle() uint16 { fake.pubHandleMutex.Lock() ret, specificReturn := fake.pubHandleReturnsOnCall[len(fake.pubHandleArgsForCall)]