Ensure subscribe data track handles are unique (#4162)

pkg/rtc/datadowntrack.go

@@ -16,7 +16,6 @@ package rtc
 
 import (
 	"fmt"
-	"math/rand"
 	"time"
 
 	"github.com/livekit/livekit-server/pkg/rtc/datatrack"
@@ -29,13 +28,13 @@ type DataDownTrackParams struct {
 	Logger           logger.Logger
 	SubscriberID     livekit.ParticipantID
 	PublishDataTrack types.DataTrack
+	Handle           uint16
 	Transport        types.DataTrackTransport
 }
 
 type DataDownTrack struct {
 	params    DataDownTrackParams
 	dti       *livekit.DataTrackInfo
-	handle    uint16
 	createdAt int64
 }
 
@@ -43,7 +42,6 @@ func NewDataDownTrack(params DataDownTrackParams, dti *livekit.DataTrackInfo) (*
 	d := &DataDownTrack{
 		params:    params,
 		dti:       dti,
-		handle:    uint16(rand.Intn(256)),
 		createdAt: time.Now().UnixNano(),
 	}
 
@@ -62,7 +60,7 @@ func (d *DataDownTrack) Close() {
 }
 
 func (d *DataDownTrack) Handle() uint16 {
-	return d.handle
+	return d.params.Handle
 }
 
 func (d *DataDownTrack) PublishDataTrack() types.DataTrack {
@@ -88,14 +86,14 @@ func (d *DataDownTrack) SubscriberID() livekit.ParticipantID {
 
 func (d *DataDownTrack) WritePacket(data []byte, packet *datatrack.Packet) {
 	forwardedPacket := *packet
-	forwardedPacket.Handle = d.handle
+	forwardedPacket.Handle = d.params.Handle
 	buf, err := forwardedPacket.Marshal()
 	if err != nil {
 		d.params.Logger.Warnw("could not marshal data track message", err)
 		return
 	}
 	if err := d.params.Transport.SendDataTrackMessage(buf); err != nil {
-		d.params.Logger.Warnw("could not send data track message", err, "handle", d.handle)
+		d.params.Logger.Warnw("could not send data track message", err, "handle", d.params.Handle)
 	}
 }
 

pkg/rtc/datatrack.go

@@ -95,6 +95,7 @@ func (d *DataTrack) AddSubscriber(sub types.LocalParticipant) (types.DataDownTra
 			Logger:           sub.GetLogger().WithValues("trackID", d.ID()),
 			SubscriberID:     sub.ID(),
 			PublishDataTrack: d,
+			Handle:           sub.GetNextSubscribedDataTrackHandle(),
 			Transport:        sub.GetDataTrackTransport(),
 		},
 		d.dti,

pkg/rtc/participant.go

@@ -18,6 +18,7 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"math/rand"
 	"os"
 	"slices"
 	"strings"
@@ -271,6 +272,8 @@ type ParticipantImpl struct {
 	*UpDataTrackManager
 	*SubscriptionManager
 
+	nextSubscribedDataTrackHandle uint16
+
 	icQueue [2]atomic.Pointer[webrtc.ICECandidate]
 
 	requireBroadcast bool
@@ -358,10 +361,11 @@ func NewParticipant(params ParticipantParams) (*ParticipantImpl, error) {
 			joiningMessageFirstSeqs:       make(map[livekit.ParticipantID]uint32),
 			joiningMessageLastWrittenSeqs: make(map[livekit.ParticipantID]uint32),
 		},
-		rpcPendingAcks:      make(map[string]*utils.DataChannelRpcPendingAckHandler),
-		rpcPendingResponses: make(map[string]*utils.DataChannelRpcPendingResponseHandler),
-		onClose:             make(map[string]func(types.LocalParticipant)),
-		telemetryGuard:      &telemetry.ReferenceGuard{},
+		rpcPendingAcks:                make(map[string]*utils.DataChannelRpcPendingAckHandler),
+		rpcPendingResponses:           make(map[string]*utils.DataChannelRpcPendingResponseHandler),
+		onClose:                       make(map[string]func(types.LocalParticipant)),
+		telemetryGuard:                &telemetry.ReferenceGuard{},
+		nextSubscribedDataTrackHandle: uint16(rand.Intn(256)),
 	}
 	p.setupSignalling()
 

pkg/rtc/participant_data_track.go

@@ -139,3 +139,15 @@ func (p *ParticipantImpl) onReceivedDataTrackMessage(data []byte) {
 
 	p.listener().OnDataTrackMessage(p, data, &packet)
 }
+
+func (p *ParticipantImpl) GetNextSubscribedDataTrackHandle() uint16 {
+	p.lock.Lock()
+	defer p.lock.Unlock()
+
+	p.nextSubscribedDataTrackHandle++
+	if p.nextSubscribedDataTrackHandle == 0 {
+		p.nextSubscribedDataTrackHandle++
+	}
+
+	return p.nextSubscribedDataTrackHandle
+}

pkg/rtc/types/interfaces.go

@@ -543,6 +543,8 @@ type LocalParticipant interface {
 	GetDataTrackTransport() DataTrackTransport
 
 	ClearParticipantListener()
+
+	GetNextSubscribedDataTrackHandle() uint16
 }
 
 // ---------------------------------------------

pkg/rtc/types/typesfakes/fake_local_participant.go

@@ -376,6 +376,16 @@ type FakeLocalParticipant struct {
 	getLoggerResolverReturnsOnCall map[int]struct {
 		result1 logger.DeferredFieldResolver
 	}
+	GetNextSubscribedDataTrackHandleStub        func() uint16
+	getNextSubscribedDataTrackHandleMutex       sync.RWMutex
+	getNextSubscribedDataTrackHandleArgsForCall []struct {
+	}
+	getNextSubscribedDataTrackHandleReturns struct {
+		result1 uint16
+	}
+	getNextSubscribedDataTrackHandleReturnsOnCall map[int]struct {
+		result1 uint16
+	}
 	GetPacerStub        func() pacer.Pacer
 	getPacerMutex       sync.RWMutex
 	getPacerArgsForCall []struct {
@@ -3299,6 +3309,59 @@ func (fake *FakeLocalParticipant) GetLoggerResolverReturnsOnCall(i int, result1
 	}{result1}
 }
 
+func (fake *FakeLocalParticipant) GetNextSubscribedDataTrackHandle() uint16 {
+	fake.getNextSubscribedDataTrackHandleMutex.Lock()
+	ret, specificReturn := fake.getNextSubscribedDataTrackHandleReturnsOnCall[len(fake.getNextSubscribedDataTrackHandleArgsForCall)]
+	fake.getNextSubscribedDataTrackHandleArgsForCall = append(fake.getNextSubscribedDataTrackHandleArgsForCall, struct {
+	}{})
+	stub := fake.GetNextSubscribedDataTrackHandleStub
+	fakeReturns := fake.getNextSubscribedDataTrackHandleReturns
+	fake.recordInvocation("GetNextSubscribedDataTrackHandle", []interface{}{})
+	fake.getNextSubscribedDataTrackHandleMutex.Unlock()
+	if stub != nil {
+		return stub()
+	}
+	if specificReturn {
+		return ret.result1
+	}
+	return fakeReturns.result1
+}
+
+func (fake *FakeLocalParticipant) GetNextSubscribedDataTrackHandleCallCount() int {
+	fake.getNextSubscribedDataTrackHandleMutex.RLock()
+	defer fake.getNextSubscribedDataTrackHandleMutex.RUnlock()
+	return len(fake.getNextSubscribedDataTrackHandleArgsForCall)
+}
+
+func (fake *FakeLocalParticipant) GetNextSubscribedDataTrackHandleCalls(stub func() uint16) {
+	fake.getNextSubscribedDataTrackHandleMutex.Lock()
+	defer fake.getNextSubscribedDataTrackHandleMutex.Unlock()
+	fake.GetNextSubscribedDataTrackHandleStub = stub
+}
+
+func (fake *FakeLocalParticipant) GetNextSubscribedDataTrackHandleReturns(result1 uint16) {
+	fake.getNextSubscribedDataTrackHandleMutex.Lock()
+	defer fake.getNextSubscribedDataTrackHandleMutex.Unlock()
+	fake.GetNextSubscribedDataTrackHandleStub = nil
+	fake.getNextSubscribedDataTrackHandleReturns = struct {
+		result1 uint16
+	}{result1}
+}
+
+func (fake *FakeLocalParticipant) GetNextSubscribedDataTrackHandleReturnsOnCall(i int, result1 uint16) {
+	fake.getNextSubscribedDataTrackHandleMutex.Lock()
+	defer fake.getNextSubscribedDataTrackHandleMutex.Unlock()
+	fake.GetNextSubscribedDataTrackHandleStub = nil
+	if fake.getNextSubscribedDataTrackHandleReturnsOnCall == nil {
+		fake.getNextSubscribedDataTrackHandleReturnsOnCall = make(map[int]struct {
+			result1 uint16
+		})
+	}
+	fake.getNextSubscribedDataTrackHandleReturnsOnCall[i] = struct {
+		result1 uint16
+	}{result1}
+}
+
 func (fake *FakeLocalParticipant) GetPacer() pacer.Pacer {
 	fake.getPacerMutex.Lock()
 	ret, specificReturn := fake.getPacerReturnsOnCall[len(fake.getPacerArgsForCall)]

test/scenarios.go

@@ -348,13 +348,11 @@ func publishDataTracksForClients(t *testing.T, clients ...*testclient.RTCClient)
 	var writers []testclient.TrackWriter
 	for i := range clients {
 		c := clients[i]
-		dtw, err := c.PublishDataTrack()
-		require.NoError(t, err)
-		writers = append(writers, dtw)
-
-		dtw, err = c.PublishDataTrack()
-		require.NoError(t, err)
-		writers = append(writers, dtw)
+		for range 2 {
+			dtw, err := c.PublishDataTrack()
+			require.NoError(t, err)
+			writers = append(writers, dtw)
+		}
 	}
 	return writers
 }