diff --git a/callback.go b/callback.go index d33f05db..e3b32cd8 100644 --- a/callback.go +++ b/callback.go @@ -18,6 +18,7 @@ import ( "github.com/pion/webrtc/v4" "github.com/livekit/protocol/livekit" + "github.com/livekit/server-sdk-go/v2/datatrack" ) // ParticipantAttributesChangedFunc is callback for Participant attribute change event. @@ -44,6 +45,8 @@ type ParticipantCallback struct { OnTrackSubscriptionFailed func(sid string, rp *RemoteParticipant) OnTrackPublished func(publication *RemoteTrackPublication, rp *RemoteParticipant) OnTrackUnpublished func(publication *RemoteTrackPublication, rp *RemoteParticipant) + OnDataTrackPublished func(track *datatrack.RemoteTrack, rp *RemoteParticipant) + OnDataTrackUnpublished func(track *datatrack.RemoteTrack, rp *RemoteParticipant) OnDataReceived func(data []byte, params DataReceiveParams) // Deprecated: Use OnDataPacket instead OnDataPacket func(data DataPacket, params DataReceiveParams) OnTranscriptionReceived func(transcriptionSegments []*TranscriptionSegment, p Participant, publication TrackPublication) @@ -66,6 +69,8 @@ func NewParticipantCallback() *ParticipantCallback { OnTrackSubscriptionFailed: func(sid string, rp *RemoteParticipant) {}, OnTrackPublished: func(publication *RemoteTrackPublication, rp *RemoteParticipant) {}, OnTrackUnpublished: func(publication *RemoteTrackPublication, rp *RemoteParticipant) {}, + OnDataTrackPublished: func(track *datatrack.RemoteTrack, rp *RemoteParticipant) {}, + OnDataTrackUnpublished: func(track *datatrack.RemoteTrack, rp *RemoteParticipant) {}, OnDataReceived: func(data []byte, params DataReceiveParams) {}, OnDataPacket: func(data DataPacket, params DataReceiveParams) {}, OnTranscriptionReceived: func(transcriptionSegments []*TranscriptionSegment, p Participant, publication TrackPublication) {}, @@ -113,6 +118,12 @@ func (cb *ParticipantCallback) Merge(other *ParticipantCallback) { if other.OnTrackUnpublished != nil { cb.OnTrackUnpublished = other.OnTrackUnpublished } + if other.OnDataTrackPublished != nil { + cb.OnDataTrackPublished = other.OnDataTrackPublished + } + if other.OnDataTrackUnpublished != nil { + cb.OnDataTrackUnpublished = other.OnDataTrackUnpublished + } if other.OnDataReceived != nil { cb.OnDataReceived = other.OnDataReceived } diff --git a/datatrack/errors.go b/datatrack/errors.go index cf465473..4c98fc4d 100644 --- a/datatrack/errors.go +++ b/datatrack/errors.go @@ -17,12 +17,13 @@ package datatrack import "errors" var ( - ErrNotAllowed = errors.New("data track publishing unauthorized") - ErrDuplicateName = errors.New("track name already taken") - ErrInvalidName = errors.New("track name invalid") - ErrLimitReached = errors.New("data track publication limit reached") - ErrPublishTimeout = errors.New("timed out publishing data track") - ErrDisconnected = errors.New("room disconnected") - ErrUnpublished = errors.New("track unpublished") - ErrQueueFull = errors.New("queue full") + ErrNotAllowed = errors.New("data track publishing unauthorized") + ErrDuplicateName = errors.New("track name already taken") + ErrInvalidName = errors.New("track name invalid") + ErrLimitReached = errors.New("data track publication limit reached") + ErrPublishTimeout = errors.New("timed out publishing data track") + ErrSubscribeTimeout = errors.New("timed out subscribing to data track") + ErrDisconnected = errors.New("room disconnected") + ErrUnpublished = errors.New("track unpublished") + ErrQueueFull = errors.New("queue full") ) diff --git a/datatrack/proto.go b/datatrack/proto.go index b23a65a4..42da96a6 100644 --- a/datatrack/proto.go +++ b/datatrack/proto.go @@ -17,6 +17,7 @@ import ( "fmt" "github.com/livekit/protocol/livekit" + "github.com/livekit/protocol/logger" ) var ( @@ -273,3 +274,41 @@ func publishRejectionFromRequestResponse(msg *livekit.RequestResponse) (rejectio } return rejection, true } + +// subscriptionUpdate asks the SFU to start or stop delivering a track's packets. +type subscriptionUpdate struct { + sid SID + subscribe bool +} + +func (u subscriptionUpdate) toProto() *livekit.UpdateDataSubscription { + return &livekit.UpdateDataSubscription{Updates: []*livekit.UpdateDataSubscription_Update{{ + TrackSid: string(u.sid), + Subscribe: u.subscribe, + }}} +} + +// publicationUpdatesFromProto maps each remote participant to the data tracks it publishes. The +// local participant is skipped and a disconnected participant contributes an empty list. Tracks +// that fail to convert are dropped with a warning. +func publicationUpdatesFromProto(participants []*livekit.ParticipantInfo, localIdentity string, log logger.Logger) map[string][]Info { + updates := make(map[string][]Info, len(participants)) + for _, participant := range participants { + if participant.GetIdentity() == localIdentity { + continue + } + infos := make([]Info, 0, len(participant.GetDataTracks())) + if participant.GetState() != livekit.ParticipantInfo_DISCONNECTED { + for _, msg := range participant.GetDataTracks() { + info, err := infoFromProto(msg) + if err != nil { + log.Warnw("ignoring invalid data track info", err, "participant", participant.GetIdentity()) + continue + } + infos = append(infos, info) + } + } + updates[participant.GetIdentity()] = infos + } + return updates +} diff --git a/datatrack/proto_test.go b/datatrack/proto_test.go index 286368ad..7807086f 100644 --- a/datatrack/proto_test.go +++ b/datatrack/proto_test.go @@ -18,6 +18,7 @@ import ( "testing" "github.com/livekit/protocol/livekit" + "github.com/livekit/protocol/logger" "github.com/stretchr/testify/require" ) @@ -141,3 +142,24 @@ func TestProto_PublishRejectionFromRequestResponse(t *testing.T) { require.Equal(t, trackHandle(1), rejection.handle) require.ErrorIs(t, rejection.err, ErrNotAllowed) } + +func TestProto_SubscriptionUpdateToProto(t *testing.T) { + update := subscriptionUpdate{sid: "DTR_1234", subscribe: true}.toProto() + require.Len(t, update.GetUpdates(), 1) + require.Equal(t, "DTR_1234", update.GetUpdates()[0].GetTrackSid()) + require.True(t, update.GetUpdates()[0].GetSubscribe()) +} + +func TestProto_PublicationUpdatesFromProto(t *testing.T) { + participants := []*livekit.ParticipantInfo{ + {Identity: "local", DataTracks: []*livekit.DataTrackInfo{{PubHandle: 1, Sid: "DTR_0000", Name: "mine"}}}, + {Identity: "publisher", DataTracks: []*livekit.DataTrackInfo{{PubHandle: 1, Sid: "DTR_1234", Name: "track1"}}}, + {Identity: "leaving", State: livekit.ParticipantInfo_DISCONNECTED, DataTracks: []*livekit.DataTrackInfo{{PubHandle: 1, Sid: "DTR_4567", Name: "stale"}}}, + } + + updates := publicationUpdatesFromProto(participants, "local", logger.GetLogger()) + require.Len(t, updates, 2) + require.Len(t, updates["publisher"], 1) + require.Equal(t, SID("DTR_1234"), updates["publisher"][0].SID) + require.Empty(t, updates["leaving"]) +} diff --git a/datatrack/remote.go b/datatrack/remote.go new file mode 100644 index 00000000..3140b690 --- /dev/null +++ b/datatrack/remote.go @@ -0,0 +1,754 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package datatrack + +import ( + "context" + "errors" + "io" + "sync" + "sync/atomic" + "time" + + "github.com/frostbyte73/core" + dtp "github.com/livekit/protocol/datatrack" + "github.com/livekit/protocol/livekit" + "github.com/livekit/protocol/logger" +) + +const ( + defaultBufferSize = 16 + defaultMaxPartialFrames = 1 + packetBufferCount = 16 + subscribeTimeout = 10 * time.Second +) + +// RemoteTransport carries what the remote manager produces: subscription requests to the SFU and +// publication events for the application. +type RemoteTransport interface { + SendUpdateSubscription(req *livekit.UpdateDataSubscription) error + OnTrackPublished(track *RemoteTrack) + OnTrackUnpublished(track *RemoteTrack) +} + +// Decryptor opens end-to-end encrypted frame payloads. +type Decryptor interface { + Decrypt(payload []byte, e2ee E2EEExtension) ([]byte, error) +} + +// SubscribeOptions configure a subscription. +type SubscribeOptions struct { + // BufferSize is the number of received frames buffered for the subscriber; older frames are + // dropped when it is exceeded. Zero is clamped to one. + BufferSize int +} + +// SubscribeOption customizes a subscription. +type SubscribeOption func(*SubscribeOptions) + +// WithBufferSize sets the number of received frames buffered for the subscriber. It has no +// effect when the track already has an active subscription. +func WithBufferSize(frames int) SubscribeOption { + return func(options *SubscribeOptions) { + options.BufferSize = frames + } +} + +// PipelineOptions configure how a remote track's packets are reassembled. +type PipelineOptions struct { + // MaxPartialFrames is the number of frames reassembled concurrently. Higher values tolerate + // more reordering at the cost of buffering. Zero is clamped to one. + MaxPartialFrames int +} + +type RemoteManagerParams struct { + Transport RemoteTransport + // Decryptor opens frames of tracks that use end-to-end encryption; nil disables decryption. + Decryptor Decryptor + Logger logger.Logger +} + +// RemoteManager tracks the data tracks published by remote participants and the local +// participant's subscriptions to them. Methods are safe for concurrent use. +type RemoteManager struct { + params RemoteManagerParams + + // mu guards descriptors and subHandles. Each RemoteTrack guards its own state: the manager + // reaches it through methods, which may take the track's lock while mu is held, and a track + // never calls into the manager while holding its own lock. Methods suffixed Locked expect mu. + mu sync.Mutex + descriptors map[SID]*RemoteTrack + subHandles map[trackHandle]*RemoteTrack +} + +func NewRemoteManager(params RemoteManagerParams) *RemoteManager { + if params.Logger == nil { + params.Logger = logger.GetLogger() + } + return &RemoteManager{ + params: params, + descriptors: make(map[SID]*RemoteTrack), + subHandles: make(map[trackHandle]*RemoteTrack), + } +} + +// HandleParticipantUpdate applies the data tracks listed for each participant in a +// ParticipantUpdate; tracks a participant no longer lists are unpublished. +func (m *RemoteManager) HandleParticipantUpdate(participants []*livekit.ParticipantInfo, localIdentity string) { + m.handlePublicationUpdates(publicationUpdatesFromProto(participants, localIdentity, m.params.Logger)) +} + +// HandleParticipantSnapshot applies a complete list of the room's participants, as carried by a +// JoinResponse. Publishers absent from the list are treated as gone. +func (m *RemoteManager) HandleParticipantSnapshot(participants []*livekit.ParticipantInfo, localIdentity string) { + updates := publicationUpdatesFromProto(participants, localIdentity, m.params.Logger) + m.mu.Lock() + for _, track := range m.descriptors { + if _, present := updates[track.publisherIdentity]; !present { + updates[track.publisherIdentity] = nil + } + } + m.mu.Unlock() + m.handlePublicationUpdates(updates) +} + +func (m *RemoteManager) handlePublicationUpdates(updates map[string][]Info) { + if len(updates) == 0 { + return + } + var published, unpublished []*RemoteTrack + var resubscribe []subscriptionUpdate + + m.mu.Lock() + for publisherIdentity, infos := range updates { + sidsInUpdate := make(map[SID]struct{}, len(infos)) + for _, info := range infos { + sidsInUpdate[info.SID] = struct{}{} + if _, known := m.descriptors[info.SID]; known { + continue + } + if update, reassigned := m.reassignSIDLocked(publisherIdentity, info); reassigned { + if update != nil { + resubscribe = append(resubscribe, *update) + } + continue + } + track := newRemoteTrack(m, info, publisherIdentity) + m.descriptors[info.SID] = track + published = append(published, track) + } + + for sid, track := range m.descriptors { + if track.publisherIdentity != publisherIdentity { + continue + } + if _, present := sidsInUpdate[sid]; !present { + delete(m.descriptors, sid) + if handle, wasActive := track.end(ErrUnpublished); wasActive { + delete(m.subHandles, handle) + } + unpublished = append(unpublished, track) + } + } + } + m.mu.Unlock() + + for _, update := range resubscribe { + m.sendSubscriptionUpdate(update) + } + for _, track := range published { + m.params.Transport.OnTrackPublished(track) + } + for _, track := range unpublished { + m.params.Transport.OnTrackUnpublished(track) + } +} + +// reassignSIDLocked detects a track republished under a new SID after its publisher's full +// reconnect: publisher identity and handle are stable across republications. It reports whether +// the SID was reassigned and, if a subscription must be re-requested under the new SID, the +// update to send. +func (m *RemoteManager) reassignSIDLocked(publisherIdentity string, info Info) (*subscriptionUpdate, bool) { + var track *RemoteTrack + for _, candidate := range m.descriptors { + if candidate.publisherIdentity == publisherIdentity && candidate.pubHandle == info.pubHandle { + track = candidate + break + } + } + if track == nil { + return nil, false + } + oldSID, subscribed, ok := track.reassignSID(info) + if !ok { + return nil, false + } + delete(m.descriptors, oldSID) + m.descriptors[info.SID] = track + if subscribed { + // the SFU does not carry subscriptions across the publisher's full reconnect + return &subscriptionUpdate{sid: info.SID, subscribe: true}, true + } + return nil, true +} + +func schemaEqual(a, b *SchemaID) bool { + if a == nil || b == nil { + return a == b + } + return *a == *b +} + +// HandleSubscriberHandles records the handles the SFU assigned to requested subscriptions, which +// activates pending subscriptions. +func (m *RemoteManager) HandleSubscriberHandles(msg *livekit.DataTrackSubscriberHandles) { + mapping, err := subscriberHandlesFromProto(msg) + if err != nil { + m.params.Logger.Warnw("ignoring invalid data track subscriber handles", err) + return + } + + m.mu.Lock() + defer m.mu.Unlock() + for handle, sid := range mapping { + track, known := m.descriptors[sid] + if !known { + m.params.Logger.Warnw("subscriber handle for unknown data track", nil, "sid", sid) + continue + } + previous, replaced, ok := track.activate(handle) + if !ok { + m.params.Logger.Warnw("subscriber handle for data track without subscription", nil, "sid", sid) + continue + } + if replaced { + delete(m.subHandles, previous) + } + m.subHandles[handle] = track + } +} + +// releaseHandle forgets the handle of a subscription the track ended, unless the SFU has since +// assigned the same handle to the track again. +func (m *RemoteManager) releaseHandle(track *RemoteTrack, handle trackHandle) { + m.mu.Lock() + defer m.mu.Unlock() + if m.subHandles[handle] == track && !track.usesHandle(handle) { + delete(m.subHandles, handle) + } +} + +// HandlePacket routes a packet received on the data channel to its track's subscribers. +func (m *RemoteManager) HandlePacket(data []byte) { + packet, err := parsePacket(data) + if err != nil { + m.params.Logger.Warnw("dropping invalid data track packet", err) + return + } + + handle := trackHandle(packet.Handle) + m.mu.Lock() + track := m.subHandles[handle] + m.mu.Unlock() + if track == nil { + m.params.Logger.Warnw("dropping data track packet without subscription", nil, "handle", packet.Handle) + return + } + track.deliver(handle, packet) +} + +// ResendSubscriptionUpdates re-requests every pending and active subscription after a full +// reconnect. +func (m *RemoteManager) ResendSubscriptionUpdates() { + m.mu.Lock() + var updates []subscriptionUpdate + for _, track := range m.descriptors { + if sid, subscribed := track.subscribedSID(); subscribed { + updates = append(updates, subscriptionUpdate{sid: sid, subscribe: true}) + } + } + m.mu.Unlock() + + for _, update := range updates { + m.sendSubscriptionUpdate(update) + } +} + +// Shutdown ends every subscription and marks every track unpublished. The manager can be used +// again for a new session. +func (m *RemoteManager) Shutdown() { + m.mu.Lock() + for _, track := range m.descriptors { + track.end(ErrDisconnected) + } + clear(m.descriptors) + clear(m.subHandles) + m.mu.Unlock() +} + +func (m *RemoteManager) sendSubscriptionUpdate(update subscriptionUpdate) { + if err := m.params.Transport.SendUpdateSubscription(update.toProto()); err != nil { + m.params.Logger.Warnw("could not send data track subscription update", err, "sid", update.sid, "subscribe", update.subscribe) + } +} + +type subscriptionState int + +const ( + subscriptionNone subscriptionState = iota + subscriptionPending + subscriptionActive +) + +type subscribeResult struct { + stream *Stream + err error +} + +// RemoteTrack is a data track published by a remote participant. Methods are safe for concurrent +// use. +type RemoteTrack struct { + manager *RemoteManager + publisherIdentity string + pubHandle trackHandle + unpublished core.Fuse + maxPartialFrames atomic.Int64 + + // streamList is an immutable snapshot of streams, read by the worker without the lock + streamList atomic.Pointer[[]*Stream] + + // mu guards the fields below. It is never held while calling into the manager. Methods + // suffixed Locked expect it. + mu sync.Mutex + info Info + subscription subscriptionState + waiters []chan subscribeResult + bufferSize int + subHandle trackHandle + // packets feeds the worker goroutine of the active subscription + packets chan *dtp.Packet + streams map[*Stream]struct{} +} + +func newRemoteTrack(manager *RemoteManager, info Info, publisherIdentity string) *RemoteTrack { + track := &RemoteTrack{ + manager: manager, + publisherIdentity: publisherIdentity, + pubHandle: info.pubHandle, + info: info, + streams: make(map[*Stream]struct{}), + } + track.maxPartialFrames.Store(defaultMaxPartialFrames) + track.streamList.Store(&[]*Stream{}) + return track +} + +// Info returns a snapshot of the track's metadata. The SID changes when the publisher completes a +// full reconnect. +func (t *RemoteTrack) Info() Info { + t.mu.Lock() + defer t.mu.Unlock() + return t.info +} + +// PublisherIdentity is the identity of the participant who published the track. +func (t *RemoteTrack) PublisherIdentity() string { + return t.publisherIdentity +} + +// IsPublished reports whether the track is still published. +func (t *RemoteTrack) IsPublished() bool { + return !t.unpublished.IsBroken() +} + +// Unpublished is closed once the track is no longer published, whether by its publisher, the SFU, +// or disconnecting from the room. +func (t *RemoteTrack) Unpublished() <-chan struct{} { + return t.unpublished.Watch() +} + +// SetPipelineOptions configures how the track's packets are reassembled. The options apply to all +// current and future subscriptions and take effect with the next packet. +func (t *RemoteTrack) SetPipelineOptions(options PipelineOptions) { + if options.MaxPartialFrames < 1 { + t.manager.params.Logger.Warnw("zero is not a valid value for MaxPartialFrames, using one", nil) + options.MaxPartialFrames = 1 + } + t.maxPartialFrames.Store(int64(options.MaxPartialFrames)) +} + +// Subscribe starts receiving the track's frames. Only the first subscription talks to the SFU; +// later ones share the pipeline and miss frames delivered before they were made. A 10 second +// deadline applies when ctx has none. +func (t *RemoteTrack) Subscribe(ctx context.Context, opts ...SubscribeOption) (*Stream, error) { + options := SubscribeOptions{BufferSize: defaultBufferSize} + for _, opt := range opts { + opt(&options) + } + if options.BufferSize < 1 { + t.manager.params.Logger.Warnw("zero is not a valid buffer size, using one", nil) + options.BufferSize = 1 + } + + _, hasDeadline := ctx.Deadline() + ctx, cancel := withDefaultTimeout(ctx, subscribeTimeout) + defer cancel() + + t.mu.Lock() + if t.unpublished.IsBroken() { + t.mu.Unlock() + return nil, ErrUnpublished + } + if t.subscription == subscriptionActive { + stream := t.addStreamLocked(options.BufferSize) + t.mu.Unlock() + return stream, nil + } + waiter := make(chan subscribeResult, 1) + t.waiters = append(t.waiters, waiter) + var request *subscriptionUpdate + if t.subscription == subscriptionNone { + t.subscription = subscriptionPending + t.bufferSize = options.BufferSize + request = &subscriptionUpdate{sid: t.info.SID, subscribe: true} + } + t.mu.Unlock() + + if request != nil { + t.manager.sendSubscriptionUpdate(*request) + } + + select { + case res := <-waiter: + return res.stream, res.err + case <-ctx.Done(): + var withdraw *subscriptionUpdate + t.mu.Lock() + t.removeWaiterLocked(waiter) + if t.subscription == subscriptionPending && len(t.waiters) == 0 { + t.subscription = subscriptionNone + withdraw = &subscriptionUpdate{sid: t.info.SID, subscribe: false} + } + t.mu.Unlock() + + if withdraw != nil { + t.manager.sendSubscriptionUpdate(*withdraw) + } + select { + case res := <-waiter: + // the answer came first and was delivered under the lock + if res.stream != nil { + res.stream.Close() + } + default: + } + if !hasDeadline && errors.Is(ctx.Err(), context.DeadlineExceeded) { + return nil, ErrSubscribeTimeout + } + return nil, ctx.Err() + } +} + +func (t *RemoteTrack) removeWaiterLocked(waiter chan subscribeResult) { + for i, w := range t.waiters { + if w == waiter { + t.waiters = append(t.waiters[:i], t.waiters[i+1:]...) + return + } + } +} + +// reassignSID moves the track to the SID it was republished under after its publisher's full +// reconnect. It reports the SID replaced and whether a subscription must be re-requested, or false +// when the info differs in more than the SID. +func (t *RemoteTrack) reassignSID(info Info) (oldSID SID, subscribed bool, ok bool) { + t.mu.Lock() + defer t.mu.Unlock() + // other than the SID, the info should not have changed + if t.info.Name != info.Name || t.info.UsesE2EE != info.UsesE2EE || + !schemaEqual(t.info.Schema, info.Schema) || t.info.FrameEncoding != info.FrameEncoding { + t.manager.params.Logger.Warnw("data track info mismatch, treating as new publication", nil, "sid", t.info.SID) + return "", false, false + } + oldSID = t.info.SID + t.manager.params.Logger.Debugw("data track SID reassigned", "oldSid", oldSID, "newSid", info.SID) + t.info.SID = info.SID + return oldSID, t.subscription != subscriptionNone, true +} + +// subscribedSID reports the track's SID when it has a pending or active subscription. +func (t *RemoteTrack) subscribedSID() (SID, bool) { + t.mu.Lock() + defer t.mu.Unlock() + return t.info.SID, t.subscription != subscriptionNone +} + +// usesHandle reports whether the track's active subscription is routed by handle. +func (t *RemoteTrack) usesHandle(handle trackHandle) bool { + t.mu.Lock() + defer t.mu.Unlock() + return t.subscription == subscriptionActive && t.subHandle == handle +} + +// activate records the handle the SFU assigned to the subscription. A pending subscription +// becomes active and every waiter receives a stream; an active one takes the new handle, as +// happens after a full reconnect, and the handle it replaces is returned. It reports false when +// the track has no subscription. +func (t *RemoteTrack) activate(handle trackHandle) (previous trackHandle, replaced bool, ok bool) { + t.mu.Lock() + defer t.mu.Unlock() + switch t.subscription { + case subscriptionNone: + return 0, false, false + case subscriptionActive: + previous, t.subHandle = t.subHandle, handle + return previous, true, true + } + + var decryptor Decryptor + if t.info.UsesE2EE { + decryptor = t.manager.params.Decryptor + } + t.packets = make(chan *dtp.Packet, packetBufferCount) + go t.runPipeline(t.packets, newRemotePipeline(decryptor, t.manager.params.Logger)) + t.subHandle = handle + t.subscription = subscriptionActive + for _, waiter := range t.waiters { + waiter <- subscribeResult{stream: t.addStreamLocked(t.bufferSize)} + } + t.waiters = nil + return 0, false, true +} + +// deliver hands a packet to the subscription's worker. The packet is dropped when the subscription +// ended, or moved to another handle, since the manager looked the track up. +func (t *RemoteTrack) deliver(handle trackHandle, packet *dtp.Packet) { + t.mu.Lock() + defer t.mu.Unlock() + if t.subscription != subscriptionActive || t.subHandle != handle { + t.manager.params.Logger.Warnw("dropping data track packet without subscription", nil, "handle", packet.Handle) + return + } + // the send happens under the lock so the channel is never closed underneath it + select { + case t.packets <- packet: + default: + t.manager.params.Logger.Debugw("dropping data track packet, pipeline is behind", "sid", t.info.SID) + } +} + +// runPipeline reassembles the subscription's packets on its own goroutine, so tracks never wait on +// one another, and delivers completed frames to every stream. It ends when packets is closed. +func (t *RemoteTrack) runPipeline(packets <-chan *dtp.Packet, pipeline *remotePipeline) { + for packet := range packets { + frame, ok := pipeline.processPacket(packet, int(t.maxPartialFrames.Load())) + if !ok { + continue + } + for _, stream := range *t.streamList.Load() { + stream.push(frame) + } + } +} + +// deactivateLocked stops the worker of the active subscription. +func (t *RemoteTrack) deactivateLocked() { + if t.packets != nil { + close(t.packets) + t.packets = nil + } +} + +func (t *RemoteTrack) refreshStreamListLocked() { + list := make([]*Stream, 0, len(t.streams)) + for stream := range t.streams { + list = append(list, stream) + } + t.streamList.Store(&list) +} + +func (t *RemoteTrack) addStreamLocked(bufferSize int) *Stream { + stream := &Stream{track: t, buf: make([]Frame, bufferSize), ready: make(chan struct{}, 1)} + t.streams[stream] = struct{}{} + t.refreshStreamListLocked() + return stream +} + +// end marks the track unpublished, fails waiters with err, and closes every stream. It reports the +// handle of the active subscription it ended, if there was one. +func (t *RemoteTrack) end(err error) (handle trackHandle, wasActive bool) { + t.mu.Lock() + defer t.mu.Unlock() + handle, wasActive = t.subHandle, t.subscription == subscriptionActive + t.unpublished.Break() + for _, waiter := range t.waiters { + waiter <- subscribeResult{err: err} + } + t.waiters = nil + for stream := range t.streams { + stream.close() + } + clear(t.streams) + t.refreshStreamListLocked() + t.subscription = subscriptionNone + t.deactivateLocked() + return handle, wasActive +} + +// removeStream is the subscriber side of Stream.Close: the last stream to leave ends the SFU +// subscription. +func (t *RemoteTrack) removeStream(stream *Stream) { + var ( + ended bool + handle trackHandle + sid SID + ) + t.mu.Lock() + stream.close() + if _, present := t.streams[stream]; present { + delete(t.streams, stream) + t.refreshStreamListLocked() + if len(t.streams) == 0 && t.subscription == subscriptionActive { + t.subscription = subscriptionNone + t.deactivateLocked() + ended, handle, sid = true, t.subHandle, t.info.SID + } + } + t.mu.Unlock() + + if ended { + t.manager.releaseHandle(t, handle) + t.manager.sendSubscriptionUpdate(subscriptionUpdate{sid: sid, subscribe: false}) + } +} + +// Stream delivers the frames of one subscription. +type Stream struct { + track *RemoteTrack + mu sync.Mutex + buf []Frame + head int + n int + ready chan struct{} + closed bool +} + +// Next blocks until a frame arrives. It returns io.EOF once the stream is closed, the track is +// unpublished, or the room disconnects, and ctx's error when ctx ends first. +func (s *Stream) Next(ctx context.Context) (Frame, error) { + for { + s.mu.Lock() + if s.n > 0 { + frame := s.buf[s.head] + s.buf[s.head] = Frame{} + s.head = (s.head + 1) % len(s.buf) + s.n-- + s.mu.Unlock() + return frame, nil + } + closed := s.closed + s.mu.Unlock() + if closed { + return Frame{}, io.EOF + } + select { + case <-s.ready: + case <-ctx.Done(): + return Frame{}, ctx.Err() + } + } +} + +// Close ends the subscription. It is safe to call more than once. +func (s *Stream) Close() { + s.track.removeStream(s) +} + +// push buffers a frame, dropping the oldest buffered frame when full. +func (s *Stream) push(frame Frame) { + s.mu.Lock() + defer s.mu.Unlock() + if s.closed { + return + } + if s.n == len(s.buf) { + s.buf[s.head] = frame + s.head = (s.head + 1) % len(s.buf) + } else { + s.buf[(s.head+s.n)%len(s.buf)] = frame + s.n++ + } + select { + case s.ready <- struct{}{}: + default: + } +} + +func (s *Stream) close() { + s.mu.Lock() + defer s.mu.Unlock() + if !s.closed { + s.closed = true + close(s.ready) + } +} + +// remotePipeline turns a subscription's packets back into frames. It is owned by one goroutine. +type remotePipeline struct { + decryptor Decryptor + log logger.Logger + depacketizer *depacketizer +} + +func newRemotePipeline(decryptor Decryptor, log logger.Logger) *remotePipeline { + return &remotePipeline{decryptor: decryptor, log: log, depacketizer: newDepacketizer()} +} + +// processPacket reports whether the packet completed a frame. +func (p *remotePipeline) processPacket(packet *dtp.Packet, maxPartialFrames int) (Frame, bool) { + result := p.depacketizer.push(*packet, depacketizerPushOptions{maxPartialFrames: maxPartialFrames}) + + if result.drop != nil { + p.log.Debugw("data track frame dropped", "reason", result.drop.Error()) + } + if result.frame == nil { + return Frame{}, false + } + + frame := Frame{Payload: result.frame.payload, UserTimestamp: result.frame.extensions.UserTimestamp} + if p.decryptor == nil { + return frame, true + } + e2ee := result.frame.extensions.E2EE + if e2ee == nil { + p.log.Errorw("dropping data track frame without E2EE extension", nil) + return Frame{}, false + } + payload, err := p.decryptor.Decrypt(frame.Payload, *e2ee) + if err != nil { + p.log.Errorw("dropping data track frame that failed to decrypt", err) + return Frame{}, false + } + frame.Payload = payload + return frame, true +} + +// withDefaultTimeout applies deadline d to ctx when it has none. The returned cancel must be called. +func withDefaultTimeout(ctx context.Context, d time.Duration) (context.Context, context.CancelFunc) { + if _, ok := ctx.Deadline(); ok { + return ctx, func() {} + } + return context.WithTimeout(ctx, d) +} diff --git a/datatrack/remote_test.go b/datatrack/remote_test.go new file mode 100644 index 00000000..14c95906 --- /dev/null +++ b/datatrack/remote_test.go @@ -0,0 +1,549 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package datatrack + +import ( + "bytes" + "context" + "io" + "testing" + "time" + + dtp "github.com/livekit/protocol/datatrack" + "github.com/livekit/protocol/livekit" + "github.com/livekit/protocol/logger" + "github.com/stretchr/testify/require" +) + +// fakeRemoteTransport records what the manager sends and emits, in order. +type fakeRemoteTransport struct { + subscriptionUpdates chan *livekit.UpdateDataSubscription + published chan *RemoteTrack + unpublished chan *RemoteTrack +} + +func newFakeRemoteTransport() *fakeRemoteTransport { + return &fakeRemoteTransport{ + subscriptionUpdates: make(chan *livekit.UpdateDataSubscription, 16), + published: make(chan *RemoteTrack, 16), + unpublished: make(chan *RemoteTrack, 16), + } +} + +func (f *fakeRemoteTransport) SendUpdateSubscription(req *livekit.UpdateDataSubscription) error { + f.subscriptionUpdates <- req + return nil +} + +func (f *fakeRemoteTransport) OnTrackPublished(track *RemoteTrack) { + f.published <- track +} + +func (f *fakeRemoteTransport) OnTrackUnpublished(track *RemoteTrack) { + f.unpublished <- track +} + +// prefixStrippingDecryptor undoes prefixingEncryptor. +type prefixStrippingDecryptor struct{} + +func (prefixStrippingDecryptor) Decrypt(payload []byte, _ E2EEExtension) ([]byte, error) { + return payload[4:], nil +} + +// expectNoEvent fails if ch delivers anything within the grace period. +func expectNoEvent[T any](t *testing.T, ch <-chan T) { + t.Helper() + select { + case event := <-ch: + t.Fatalf("unexpected event %v", event) + case <-time.After(50 * time.Millisecond): + } +} + +func expectFrame(t *testing.T, stream *Stream) Frame { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond) + defer cancel() + frame, err := stream.Next(ctx) + require.NoError(t, err) + return frame +} + +func expectNoFrame(t *testing.T, stream *Stream) { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + _, err := stream.Next(ctx) + require.ErrorIs(t, err, context.DeadlineExceeded) +} + +func expectClosed(t *testing.T, stream *Stream) { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + _, err := stream.Next(ctx) + require.ErrorIs(t, err, io.EOF) +} + +// publishTrack simulates the SFU announcing a publication and returns the track handed to the +// application. +func publishTrack(t *testing.T, m *RemoteManager, transport *fakeRemoteTransport, publisherIdentity string, info Info) *RemoteTrack { + t.Helper() + m.handlePublicationUpdates(map[string][]Info{publisherIdentity: {info}}) + return expectEvent(t, transport.published) +} + +func expectSubscriptionUpdate(t *testing.T, transport *fakeRemoteTransport) (SID, bool) { + t.Helper() + msg := expectEvent(t, transport.subscriptionUpdates) + require.Len(t, msg.GetUpdates(), 1) + update := msg.GetUpdates()[0] + return SID(update.GetTrackSid()), update.GetSubscribe() +} + +// assignHandle simulates the SFU assigning a subscriber handle. +func assignHandle(m *RemoteManager, handle trackHandle, sid SID) { + m.HandleSubscriberHandles(&livekit.DataTrackSubscriberHandles{ + SubHandles: map[uint32]*livekit.DataTrackSubscriberHandles_PublishedDataTrack{ + uint32(handle): {TrackSid: string(sid)}, + }, + }) +} + +// subscribeAsync calls Subscribe on another goroutine and delivers its outcome on the returned channel. +func subscribeAsync(ctx context.Context, track *RemoteTrack, opts ...SubscribeOption) <-chan subscribeResult { + result := make(chan subscribeResult, 1) + go func() { + stream, err := track.Subscribe(ctx, opts...) + result <- subscribeResult{stream: stream, err: err} + }() + return result +} + +// subscribeAndActivate subscribes, answers the SFU request with handle, and returns the stream. +func subscribeAndActivate(t *testing.T, m *RemoteManager, transport *fakeRemoteTransport, track *RemoteTrack, handle trackHandle) *Stream { + t.Helper() + result := subscribeAsync(context.Background(), track) + + sid, subscribe := expectSubscriptionUpdate(t, transport) + require.True(t, subscribe) + require.Equal(t, track.Info().SID, sid) + + assignHandle(m, handle, sid) + + res := expectEvent(t, result) + require.NoError(t, res.err) + return res.stream +} + +// rawPacket marshals a packet derived from the Rust vector header. +func rawPacket(t *testing.T, marker FrameMarker, handle trackHandle, sequence, frameNumber uint16, payload []byte, extensions Extensions) []byte { + t.Helper() + header := testHeader() + marker.apply(&header) + header.Handle = uint16(handle) + header.SequenceNumber = sequence + header.FrameNumber = frameNumber + extensions.apply(&header) + packet := dtp.Packet{Header: header, Payload: payload} + raw, err := packet.Marshal() + require.NoError(t, err) + return raw +} + +func singlePacket(t *testing.T, handle trackHandle, payload []byte, extensions Extensions) []byte { + return rawPacket(t, FrameMarkerSingle, handle, 0, 0, payload, extensions) +} + +// pushInterleavedTwoFramePair pushes Start(frame1), Start(frame2), Final(frame1), Final(frame2) +// through the manager to exercise the depacketizer's concurrent partial frame handling. +func pushInterleavedTwoFramePair(t *testing.T, m *RemoteManager, handle trackHandle, frameOne, frameOneStart uint16, frameOnePayloads [2][]byte, frameTwo, frameTwoStart uint16, frameTwoPayloads [2][]byte) { + t.Helper() + push := func(frameNumber, sequence uint16, marker FrameMarker, payload []byte) { + m.HandlePacket(rawPacket(t, marker, handle, sequence, frameNumber, payload, Extensions{})) + } + push(frameOne, frameOneStart, FrameMarkerStart, frameOnePayloads[0]) + push(frameTwo, frameTwoStart, FrameMarkerStart, frameTwoPayloads[0]) + push(frameOne, frameOneStart+1, FrameMarkerFinal, frameOnePayloads[1]) + push(frameTwo, frameTwoStart+1, FrameMarkerFinal, frameTwoPayloads[1]) +} + +func TestRemotePipeline_ProcessPacket(t *testing.T) { + const payloadLen = 1024 + pipeline := newRemotePipeline(nil, logger.GetLogger()) + + header := testHeader() + FrameMarkerSingle.apply(&header) + packet := &dtp.Packet{Header: header, Payload: bytes.Repeat([]byte{0xab}, payloadLen)} + + frame, ok := pipeline.processPacket(packet, defaultMaxPartialFrames) + require.True(t, ok, "should return a frame") + require.Len(t, frame.Payload, payloadLen) +} + +func TestRemoteManager_Shutdown(t *testing.T) { + transport := newFakeRemoteTransport() + m := NewRemoteManager(RemoteManagerParams{Transport: transport}) + + track := publishTrack(t, m, transport, "id", Info{SID: "DTR_1234", pubHandle: 1, Name: "test"}) + pending := subscribeAsync(context.Background(), track) + expectSubscriptionUpdate(t, transport) + + m.Shutdown() + + require.ErrorIs(t, expectEvent(t, pending).err, ErrDisconnected) + require.False(t, track.IsPublished()) + expectEvent(t, track.Unpublished()) +} + +func TestRemoteManager_Subscribe(t *testing.T) { + publisherIdentity, trackName, trackSID := "publisher", "track", SID("DTR_1234") + subHandle := trackHandle(0x1234) + + transport := newFakeRemoteTransport() + m := NewRemoteManager(RemoteManagerParams{Transport: transport}) + + // Simulate track published + track := publishTrack(t, m, transport, publisherIdentity, Info{SID: trackSID, pubHandle: 1, Name: trackName}) + require.True(t, track.IsPublished()) + require.Equal(t, trackName, track.Info().Name) + require.Equal(t, trackSID, track.Info().SID) + require.Equal(t, publisherIdentity, track.PublisherIdentity()) + + result := subscribeAsync(context.Background(), track) + + sid, subscribe := expectSubscriptionUpdate(t, transport) + require.True(t, subscribe) + require.Equal(t, trackSID, sid) + time.Sleep(20 * time.Millisecond) + + // Simulate SFU reply + assignHandle(m, subHandle, trackSID) + + res := expectEvent(t, result) + require.NoError(t, res.err) + require.NotNil(t, res.stream) +} + +func TestRemoteManager_TrackPublicationAddAndRemove(t *testing.T) { + transport := newFakeRemoteTransport() + m := NewRemoteManager(RemoteManagerParams{Transport: transport}) + + trackSID := SID("DTR_1234") + + // Simulate track published + track := publishTrack(t, m, transport, "identity1", Info{SID: trackSID, pubHandle: 1, Name: "test"}) + require.Equal(t, trackSID, track.Info().SID) + require.Equal(t, "test", track.Info().Name) + require.True(t, track.IsPublished()) + + // Simulate track unpublished + m.handlePublicationUpdates(map[string][]Info{"identity1": nil}) + + expectEvent(t, track.Unpublished()) + require.False(t, track.IsPublished()) + + unpublished := expectEvent(t, transport.unpublished) + require.Equal(t, trackSID, unpublished.Info().SID) +} + +func TestRemoteManager_SfuPublicationUpdatesIdempotent(t *testing.T) { + transport := newFakeRemoteTransport() + m := NewRemoteManager(RemoteManagerParams{Transport: transport}) + + info := Info{SID: "DTR_1234", pubHandle: 1, Name: "test"} + + // Simulate three identical publication updates + for range 3 { + m.handlePublicationUpdates(map[string][]Info{"identity1": {info}}) + } + + expectEvent(t, transport.published) + + // No second publication should appear + m.Shutdown() + expectNoEvent(t, transport.published) +} + +func TestRemoteManager_SidReassignmentDoesNotRepublish(t *testing.T) { + transport := newFakeRemoteTransport() + m := NewRemoteManager(RemoteManagerParams{Transport: transport}) + + pubHandle := trackHandle(7) + oldSID, newSID := SID("DTR_1234"), SID("DTR_5678") + + // Simulate track published + track := publishTrack(t, m, transport, "id", Info{SID: oldSID, pubHandle: pubHandle, Name: "test"}) + require.Equal(t, oldSID, track.Info().SID) + + // Simulate publisher full reconnect: same track, new SID + m.handlePublicationUpdates(map[string][]Info{"id": {{SID: newSID, pubHandle: pubHandle, Name: "test"}}}) + + // No publish/unpublish should appear + m.Shutdown() + expectNoEvent(t, transport.published) + expectNoEvent(t, transport.unpublished) + require.Equal(t, newSID, track.Info().SID) +} + +func TestRemoteManager_SidReassignmentResubscribesActiveSubscription(t *testing.T) { + transport := newFakeRemoteTransport() + m := NewRemoteManager(RemoteManagerParams{Transport: transport}) + + pubHandle := trackHandle(7) + oldSID, newSID := SID("DTR_1234"), SID("DTR_5678") + oldSubHandle, newSubHandle := trackHandle(0x1001), trackHandle(0x1002) + + // Simulate track published + track := publishTrack(t, m, transport, "id", Info{SID: oldSID, pubHandle: pubHandle, Name: "test"}) + + // Subscribe to the track + stream := subscribeAndActivate(t, m, transport, track, oldSubHandle) + + // Simulate publisher full reconnect: same track, new SID + m.handlePublicationUpdates(map[string][]Info{"id": {{SID: newSID, pubHandle: pubHandle, Name: "test"}}}) + + // Manager should re-subscribe under the new SID + sid, subscribe := expectSubscriptionUpdate(t, transport) + require.True(t, subscribe) + require.Equal(t, newSID, sid) + require.Equal(t, newSID, track.Info().SID) + require.True(t, track.IsPublished()) + + // Simulate SFU assigning a new subscriber handle + assignHandle(m, newSubHandle, newSID) + + // Frames received on the new handle reach the existing subscriber + m.HandlePacket(singlePacket(t, newSubHandle, []byte{1, 2, 3, 4, 5}, Extensions{})) + + frame := expectFrame(t, stream) + require.Equal(t, []byte{1, 2, 3, 4, 5}, frame.Payload) +} + +func TestRemoteManager_SubscribeReceivesFrame(t *testing.T) { + transport := newFakeRemoteTransport() + m := NewRemoteManager(RemoteManagerParams{Transport: transport}) + + trackSID, subHandle := SID("DTR_1234"), trackHandle(0x1234) + + // Simulate track published + track := publishTrack(t, m, transport, "id", Info{SID: trackSID, pubHandle: 1, Name: "test"}) + + // Subscribe to the track + stream := subscribeAndActivate(t, m, transport, track, subHandle) + + // Simulate receiving a single-frame packet + m.HandlePacket(singlePacket(t, subHandle, []byte{1, 2, 3, 4, 5}, Extensions{})) + + frame := expectFrame(t, stream) + require.Equal(t, []byte{1, 2, 3, 4, 5}, frame.Payload) +} + +func TestRemoteManager_SubscribeWithE2EE(t *testing.T) { + transport := newFakeRemoteTransport() + m := NewRemoteManager(RemoteManagerParams{Transport: transport, Decryptor: prefixStrippingDecryptor{}}) + + trackSID, subHandle := SID("DTR_1234"), trackHandle(0x1234) + + // Simulate track published (with e2ee) + track := publishTrack(t, m, transport, "id", Info{SID: trackSID, pubHandle: 1, Name: "test", UsesE2EE: true}) + + // Subscribe to the track + stream := subscribeAndActivate(t, m, transport, track, subHandle) + + // Simulate receiving an encrypted single-frame packet + payload := []byte{0xde, 0xad, 0xbe, 0xef, 1, 2, 3, 4, 5} + m.HandlePacket(singlePacket(t, subHandle, payload, Extensions{E2EE: &E2EEExtension{}})) + + // Payload should have fake encryption prefix stripped by decryptor + frame := expectFrame(t, stream) + require.Equal(t, []byte{1, 2, 3, 4, 5}, frame.Payload) +} + +func TestRemoteManager_SubscribeFanOutToMultipleSubscribers(t *testing.T) { + transport := newFakeRemoteTransport() + m := NewRemoteManager(RemoteManagerParams{Transport: transport}) + + trackSID, subHandle := SID("DTR_1234"), trackHandle(0x1234) + + // Simulate track published + track := publishTrack(t, m, transport, "id", Info{SID: trackSID, pubHandle: 1, Name: "test"}) + + // First subscriber triggers SFU interaction + stream1 := subscribeAndActivate(t, m, transport, track, subHandle) + + // Additional subscribers attach directly (no further SFU interaction) + stream2, err := track.Subscribe(context.Background()) + require.NoError(t, err) + stream3, err := track.Subscribe(context.Background()) + require.NoError(t, err) + expectNoEvent(t, transport.subscriptionUpdates) + + // Simulate receiving a single-frame packet + m.HandlePacket(singlePacket(t, subHandle, []byte{1, 2, 3, 4, 5}, Extensions{})) + + // All subscribers should receive the same frame + for _, stream := range []*Stream{stream1, stream2, stream3} { + frame := expectFrame(t, stream) + require.Equal(t, []byte{1, 2, 3, 4, 5}, frame.Payload) + } +} + +func TestRemoteManager_SubscribeUnknownTrackFails(t *testing.T) { + transport := newFakeRemoteTransport() + m := NewRemoteManager(RemoteManagerParams{Transport: transport}) + + // A track that is no longer published cannot be subscribed to + track := publishTrack(t, m, transport, "id", Info{SID: "DTR_1234", pubHandle: 1, Name: "test"}) + m.handlePublicationUpdates(map[string][]Info{"id": nil}) + expectEvent(t, transport.unpublished) + + _, err := track.Subscribe(context.Background()) + require.ErrorIs(t, err, ErrUnpublished) +} + +func TestRemoteManager_UnpublishTerminatesPendingSubscription(t *testing.T) { + transport := newFakeRemoteTransport() + m := NewRemoteManager(RemoteManagerParams{Transport: transport}) + + trackSID := SID("DTR_1234") + + // Simulate track published + track := publishTrack(t, m, transport, "id", Info{SID: trackSID, pubHandle: 1, Name: "test"}) + + // Subscribe (enters pending state) + result := subscribeAsync(context.Background(), track) + _, subscribe := expectSubscriptionUpdate(t, transport) + require.True(t, subscribe) + + // Simulate track unpublished before SFU assigns a handle + m.handlePublicationUpdates(map[string][]Info{"id": nil}) + + require.ErrorIs(t, expectEvent(t, result).err, ErrUnpublished) + + unpublished := expectEvent(t, transport.unpublished) + require.Equal(t, trackSID, unpublished.Info().SID) +} + +func TestRemoteManager_UnpublishTerminatesActiveSubscription(t *testing.T) { + transport := newFakeRemoteTransport() + m := NewRemoteManager(RemoteManagerParams{Transport: transport}) + + trackSID, subHandle := SID("DTR_1234"), trackHandle(0x1234) + + // Simulate track published + track := publishTrack(t, m, transport, "id", Info{SID: trackSID, pubHandle: 1, Name: "test"}) + + // Subscribe to the track + stream := subscribeAndActivate(t, m, transport, track, subHandle) + + // Simulate track unpublished while subscription is active + m.handlePublicationUpdates(map[string][]Info{"id": nil}) + + expectClosed(t, stream) + + unpublished := expectEvent(t, transport.unpublished) + require.Equal(t, trackSID, unpublished.Info().SID) +} + +func TestRemoteManager_AllSubscribersDroppedTerminatesSfuSubscription(t *testing.T) { + transport := newFakeRemoteTransport() + m := NewRemoteManager(RemoteManagerParams{Transport: transport}) + + trackSID, subHandle := SID("DTR_1234"), trackHandle(0x1234) + + // Simulate track published + track := publishTrack(t, m, transport, "id", Info{SID: trackSID, pubHandle: 1, Name: "test"}) + + // Subscribe to the track + stream := subscribeAndActivate(t, m, transport, track, subHandle) + + // Close the only subscriber + stream.Close() + + // Manager should request SFU to unsubscribe + sid, subscribe := expectSubscriptionUpdate(t, transport) + require.False(t, subscribe) + require.Equal(t, trackSID, sid) +} + +// Should depacketize multiple interleaved partial frames when MaxPartialFrames is set before subscribe. +func TestRemoteManager_MaxPartialFramesSetBeforeSubscribe(t *testing.T) { + transport := newFakeRemoteTransport() + m := NewRemoteManager(RemoteManagerParams{Transport: transport}) + + subHandle := trackHandle(0x1234) + track := publishTrack(t, m, transport, "id", Info{SID: "DTR_1234", pubHandle: 1, Name: "test"}) + + // Configure the track BEFORE any subscribe + track.SetPipelineOptions(PipelineOptions{MaxPartialFrames: 3}) + + stream := subscribeAndActivate(t, m, transport, track, subHandle) + + // Two interleaved partial frames: Start(1), Start(2), Final(1), Final(2). With the default + // MaxPartialFrames of 1 frame 1 would be evicted by frame 2; with 3 both frames coexist and emerge. + pushInterleavedTwoFramePair(t, m, subHandle, 1, 0, [2][]byte{{0xa1}, {0xa2}}, 2, 100, [2][]byte{{0xb1}, {0xb2}}) + + require.Equal(t, []byte{0xa1, 0xa2}, expectFrame(t, stream).Payload) + require.Equal(t, []byte{0xb1, 0xb2}, expectFrame(t, stream).Payload) +} + +// Should pick up MaxPartialFrames live on an already-active subscription. +func TestRemoteManager_MaxPartialFramesSetLive(t *testing.T) { + transport := newFakeRemoteTransport() + m := NewRemoteManager(RemoteManagerParams{Transport: transport}) + + subHandle := trackHandle(0x1234) + track := publishTrack(t, m, transport, "id", Info{SID: "DTR_1234", pubHandle: 1, Name: "test"}) + + stream := subscribeAndActivate(t, m, transport, track, subHandle) + + // Subscription is now active; flip the cap on the live pipeline + track.SetPipelineOptions(PipelineOptions{MaxPartialFrames: 3}) + + pushInterleavedTwoFramePair(t, m, subHandle, 1, 0, [2][]byte{{0xa1}, {0xa2}}, 2, 100, [2][]byte{{0xb1}, {0xb2}}) + + require.Equal(t, []byte{0xa1, 0xa2}, expectFrame(t, stream).Payload) + require.Equal(t, []byte{0xb1, 0xb2}, expectFrame(t, stream).Payload) +} + +// Should drop the older partial frame by default (no MaxPartialFrames set). +func TestRemoteManager_DefaultDropsOlderPartialFrame(t *testing.T) { + transport := newFakeRemoteTransport() + m := NewRemoteManager(RemoteManagerParams{Transport: transport}) + + subHandle := trackHandle(0x1234) + track := publishTrack(t, m, transport, "id", Info{SID: "DTR_1234", pubHandle: 1, Name: "test"}) + + stream := subscribeAndActivate(t, m, transport, track, subHandle) + + // Default cap of 1: Start(2) evicts Start(1), so Final(1) is unknown and only frame 2 makes it through + pushInterleavedTwoFramePair(t, m, subHandle, 1, 0, [2][]byte{{0xa1}, {0xa2}}, 2, 100, [2][]byte{{0xb1}, {0xb2}}) + + require.Equal(t, []byte{0xb1, 0xb2}, expectFrame(t, stream).Payload) + expectNoFrame(t, stream) +} + +func TestStream_DropsOldestWhenFull(t *testing.T) { + stream := &Stream{buf: make([]Frame, 2), ready: make(chan struct{}, 1)} + for i := byte(1); i <= 3; i++ { + stream.push(Frame{Payload: []byte{i}}) + } + require.Equal(t, []byte{2}, expectFrame(t, stream).Payload) + require.Equal(t, []byte{3}, expectFrame(t, stream).Payload) + expectNoFrame(t, stream) + stream.close() + expectClosed(t, stream) +} diff --git a/examples/datatrack/subscriber/main.go b/examples/datatrack/subscriber/main.go new file mode 100644 index 00000000..d28285e0 --- /dev/null +++ b/examples/datatrack/subscriber/main.go @@ -0,0 +1,83 @@ +package main + +import ( + "context" + "flag" + "fmt" + "os/signal" + "syscall" + + "github.com/livekit/protocol/logger" + lksdk "github.com/livekit/server-sdk-go/v2" + "github.com/livekit/server-sdk-go/v2/datatrack" +) + +var host, apiKey, apiSecret, roomName, identity string + +func init() { + flag.StringVar(&host, "host", "", "livekit server host") + flag.StringVar(&apiKey, "api-key", "", "livekit api key") + flag.StringVar(&apiSecret, "api-secret", "", "livekit api secret") + flag.StringVar(&roomName, "room-name", "", "room name") + flag.StringVar(&identity, "identity", "subscriber", "participant identity") +} + +func main() { + logger.InitFromConfig(&logger.Config{Level: "info"}, "datatrack-subscriber") + lksdk.SetLogger(logger.GetLogger()) + flag.Parse() + if host == "" || apiKey == "" || apiSecret == "" || roomName == "" { + fmt.Println("invalid arguments.") + return + } + + ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM) + defer stop() + + // Subscribe to any published data tracks + callback := &lksdk.RoomCallback{ + ParticipantCallback: lksdk.ParticipantCallback{ + OnDataTrackPublished: func(track *datatrack.RemoteTrack, rp *lksdk.RemoteParticipant) { + subscribe(ctx, track) + }, + }, + } + + room, err := lksdk.ConnectToRoom(host, lksdk.ConnectInfo{ + APIKey: apiKey, + APISecret: apiSecret, + RoomName: roomName, + ParticipantIdentity: identity, + }, callback) + if err != nil { + panic(err) + } + defer room.Disconnect() + + <-ctx.Done() +} + +// subscribe subscribes to the given data track and logs received frames. +func subscribe(ctx context.Context, track *datatrack.RemoteTrack) { + logger.Infow("subscribing", "track", track.Info().Name, "publisher", track.PublisherIdentity()) + + stream, err := track.Subscribe(ctx) + if err != nil { + logger.Warnw("failed to subscribe", err) + return + } + defer stream.Close() + + for { + frame, err := stream.Next(ctx) + if err != nil { + break + } + logger.Infow("received frame", "bytes", len(frame.Payload)) + + if latency, ok := frame.DurationSinceTimestamp(); ok { + logger.Infow("latency", "duration", latency) + } + } + logger.Infow("unsubscribed") +} diff --git a/remotedatatrack.go b/remotedatatrack.go new file mode 100644 index 00000000..19961544 --- /dev/null +++ b/remotedatatrack.go @@ -0,0 +1,52 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package lksdk + +import ( + "github.com/livekit/protocol/livekit" + "github.com/livekit/server-sdk-go/v2/datatrack" +) + +// remoteDataTrackTransport connects the remote data track manager to the engine and the room's +// callbacks. +type remoteDataTrackTransport struct { + room *Room +} + +func (t remoteDataTrackTransport) SendUpdateSubscription(req *livekit.UpdateDataSubscription) error { + return t.room.engine.SendUpdateDataSubscription(req) +} + +// OnTrackPublished runs the callbacks on their own goroutine so they may block, for example on +// Subscribe, without stalling signal handling. +func (t remoteDataTrackTransport) OnTrackPublished(track *datatrack.RemoteTrack) { + rp := t.room.GetParticipantByIdentity(track.PublisherIdentity()) + go func() { + if rp != nil { + rp.Callback.OnDataTrackPublished(track, rp) + } + t.room.callback.OnDataTrackPublished(track, rp) + }() +} + +func (t remoteDataTrackTransport) OnTrackUnpublished(track *datatrack.RemoteTrack) { + rp := t.room.GetParticipantByIdentity(track.PublisherIdentity()) + go func() { + if rp != nil { + rp.Callback.OnDataTrackUnpublished(track, rp) + } + t.room.callback.OnDataTrackUnpublished(track, rp) + }() +} diff --git a/room.go b/room.go index 58d21234..ab9d101d 100644 --- a/room.go +++ b/room.go @@ -323,6 +323,7 @@ type Room struct { name string LocalParticipant *LocalParticipant localDataTracks *datatrack.LocalManager + remoteDataTracks *datatrack.RemoteManager callback *RoomCallback sidReady chan struct{} disconnectReason livekit.DisconnectReason @@ -369,6 +370,7 @@ func NewRoom(callback *RoomCallback) *Room { r.LocalParticipant = newLocalParticipant(r.engine, r.callback, r.serverInfo, r.log) r.localDataTracks = datatrack.NewLocalManager(datatrack.LocalManagerParams{Transport: localDataTrackTransport{engine: r.engine}, Logger: r.log}) r.LocalParticipant.dataTracks = r.localDataTracks + r.remoteDataTracks = datatrack.NewRemoteManager(datatrack.RemoteManagerParams{Transport: remoteDataTrackTransport{room: r}, Logger: r.log}) return r } @@ -790,6 +792,7 @@ func (r *Room) cleanup() { r.engine.Close() r.LocalParticipant.closeTracks() r.localDataTracks.Shutdown() + r.remoteDataTracks.Shutdown() r.setSid("", true) r.byteStreamHandlers.Clear() @@ -1064,6 +1067,7 @@ func (r *Room) OnRoomJoined( r.clearParticipantDefers(livekit.ParticipantID(pi.Sid), pi) // no need to run participant defers here, since we are connected for the first time } + r.remoteDataTracks.HandleParticipantSnapshot(otherParticipants, r.LocalParticipant.Identity()) } func (r *Room) OnDisconnected(reason livekit.DisconnectReason) { @@ -1110,6 +1114,8 @@ func (r *Room) OnRestarted( r.LocalParticipant.republishTracks() r.localDataTracks.RepublishTracks() + r.remoteDataTracks.HandleParticipantSnapshot(otherParticipants, r.LocalParticipant.Identity()) + r.remoteDataTracks.ResendSubscriptionUpdates() r.callback.OnReconnected() } @@ -1186,6 +1192,7 @@ func (r *Room) OnParticipantUpdate(participants []*livekit.ParticipantInfo) { r.runParticipantDefers(newSid, rp) } } + r.remoteDataTracks.HandleParticipantUpdate(participants, r.LocalParticipant.Identity()) } func (r *Room) OnParticipantDisconnect(rp *RemoteParticipant, reason livekit.DisconnectReason) { @@ -1301,6 +1308,7 @@ func (r *Room) OnRoomMoved(moved *livekit.RoomMovedResponse) { infos = append(infos, moved.Participant) infos = append(infos, moved.OtherParticipants...) r.OnParticipantUpdate(infos) + r.remoteDataTracks.HandleParticipantSnapshot(moved.OtherParticipants, r.LocalParticipant.Identity()) } func (r *Room) OnTrackRemoteMuted(msg *livekit.MuteTrackRequest) { @@ -1425,9 +1433,11 @@ func (r *Room) OnUnpublishDataTrackResponse(unpublishDataTrackResponse *livekit. } func (r *Room) OnDataTrackSubscriberHandles(dataTrackSubscriberHandles *livekit.DataTrackSubscriberHandles) { + r.remoteDataTracks.HandleSubscriberHandles(dataTrackSubscriberHandles) } func (r *Room) OnDataTrackPacket(data []byte) { + r.remoteDataTracks.HandlePacket(data) } func (r *Room) OnStreamHeader(streamHeader *livekit.DataStream_Header, participantIdentity string) {