Add local data track flush method (#1925)

Author: 1egoman

.changeset/afraid-cooks-visit.md

@@ -0,0 +1,5 @@
+---
+'livekit-client': patch
+---
+
+Add local data track flush method

src/room/Room.ts

@@ -290,8 +290,10 @@ class Room extends (EventEmitter as new () => TypedEmitter<RoomEventCallbacks>)
       .on('trackUnpublished', (event) => {
         this.emit(RoomEvent.LocalDataTrackUnpublished, event.sid);
       })
-      .on('packetAvailable', ({ bytes }) => {
-        this.engine.sendLossyBytes(bytes, DataChannelKind.DATA_TRACK_LOSSY, 'wait');
+      .on('packetAvailable', ({ handle, bytes }) => {
+        this.engine
+          .sendLossyBytes(bytes, DataChannelKind.DATA_TRACK_LOSSY, 'wait')
+          .finally(() => this.outgoingDataTrackManager.handlePacketSendComplete(handle));
       });
 
     this.disconnectLock = new Mutex();
@@ -1730,6 +1732,8 @@ class Room extends (EventEmitter as new () => TypedEmitter<RoomEventCallbacks>)
     this.bufferedEvents = [];
     this.transcriptionReceivedTimes.clear();
     this.incomingDataStreamManager.clearControllers();
+    this.incomingDataTrackManager.reset();
+    this.outgoingDataTrackManager.reset();
     if (this.state === ConnectionState.Disconnected) {
       return;
     }

src/room/data-track/LocalDataTrack.ts

@@ -1,4 +1,5 @@
 import log, { LoggerNames, type StructuredLogger, getLogger } from '../../logger';
+import { Future } from '../utils';
 import { type DataTrackFrame, DataTrackFrameInternal } from './frame';
 import type { DataTrackHandle } from './handle';
 import type OutgoingDataTrackManager from './outgoing/OutgoingDataTrackManager';
@@ -28,14 +29,35 @@ export default class LocalDataTrack implements ILocalTrack, IDataTrack {
 
   protected log: StructuredLogger = log;
 
+  /** Resolves once the data track has sent all pending packets the rtc data channel buffer. */
+  protected flushedFuture = new Future<void, never>();
+
   /** @internal */
   constructor(options: DataTrackOptions, manager: OutgoingDataTrackManager) {
     this.options = options;
     this.manager = manager;
 
     this.log = getLogger(LoggerNames.DataTracks);
+
+    this.manager.on('packetsFlushed', this.handleManagerPacketsFlushed);
+    this.manager.on('reset', this.handleManagerReset);
   }
 
+  private handleManagerReset = () => {
+    // When the associated manager resets, mark any in flight flushes as complete
+    // There's nothing actionable a user can do to get these to complete so no
+    // error is being thrown.
+    this.handleManagerPacketsFlushed();
+
+    this.manager.off('packetsFlushed', this.handleManagerReset);
+    this.manager.off('reset', this.handleManagerReset);
+  };
+
+  private handleManagerPacketsFlushed = () => {
+    this.flushedFuture.resolve?.();
+    this.flushedFuture = new Future();
+  };
+
   /** @internal */
   static withExplicitHandle(
     options: DataTrackOptions,
@@ -104,6 +126,35 @@ export default class LocalDataTrack implements ILocalTrack, IDataTrack {
     }
   }
 
+  /**
+   * When called, waits for all in flight packets to be sent before resolving.
+   *
+   * Use this to:
+   *
+   * 1. Send frames exactly in order:
+   * ```ts
+   * await track.tryPush(/* ... *\/);
+   * await track.flush();
+   * await track.tryPush(/* ... *\/);
+   * await track.flush();
+   * // ... etc ...
+   * ```
+   *
+   * 2. Wait for frames to all be delivered before unpublishing a local data track:
+   *
+   * ```ts
+   * await track.tryPush(/* ... *\/);
+   * await track.tryPush(/* ... *\/);
+   * await track.tryPush(/* ... *\/);
+   * // ... etc ...
+   * await track.flush();
+   * await track.unpublish();
+   * ```
+   **/
+  async flush(): Promise<void> {
+    return this.flushedFuture.promise;
+  }
+
   /**
    * Unpublish the track from the SFU. Once this is called, any further calls to {@link tryPush}
    * will fail.

src/room/data-track/handle.ts

@@ -66,4 +66,8 @@ export class DataTrackHandleAllocator {
     }
     return this.value;
   }
+
+  reset() {
+    this.value = 0;
+  }
 }

src/room/data-track/incoming/IncomingDataTrackManager.test.ts

@@ -327,7 +327,7 @@ describe('DataTrackIncomingManager', () => {
       expect(sfuUpdateSubscriptionEvent.subscribe).toStrictEqual(false);
 
       // 7. Make sure shutting down the manager doesn't throw errors
-      manager.shutdown();
+      manager.reset();
     });
 
     it('should NOT terminate the sfu subscription if the abortsignal is triggered on one of two active subscriptions', async () => {
@@ -675,7 +675,7 @@ describe('DataTrackIncomingManager', () => {
       );
 
       // 4. Shutdown the manager, and make sure it doesn't throw
-      manager.shutdown();
+      manager.reset();
 
       // 5. Make sure the trackUnpublished event fires for the descriptor
       const trackUnpublishedEvent = await managerEvents.waitFor('trackUnpublished');

src/room/data-track/incoming/IncomingDataTrackManager.ts

@@ -92,7 +92,7 @@ export default class IncomingDataTrackManager extends (EventEmitter as new () =>
    *
    * This is an index that allows track descriptors to be looked up
    * by subscriber handle in O(1) time, to make routing incoming packets
-   * a (hot code path) faster.
+   * (a hot code path) faster.
    */
   private subscriptionHandles = new Map<DataTrackHandle, DataTrackSid>();
 
@@ -626,8 +626,9 @@ export default class IncomingDataTrackManager extends (EventEmitter as new () =>
     }
   }
 
-  /** Shutdown the manager, ending any subscriptions. */
-  shutdown() {
+  /** Resets the manager, ending any subscriptions, and getting it ready for the next room
+   * connection.  */
+  reset() {
     for (const descriptor of this.descriptors.values()) {
       this.emit('trackUnpublished', {
         sid: descriptor.info.sid,
@@ -643,5 +644,6 @@ export default class IncomingDataTrackManager extends (EventEmitter as new () =>
       }
     }
     this.descriptors.clear();
+    this.subscriptionHandles.clear();
   }
 }