Track data track demand

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()

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)
+}

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)
 

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)]