Port data streams over to rust livekit-ffi version

.changeset/quiet-ducks-yell.md

@@ -0,0 +1,5 @@
+---
+'@livekit/rtc-node': patch
+---
+
+Convert data streams to use livekit-ffi exposed data streams interface

packages/livekit-rtc/scripts/run-e2e.mjs

@@ -16,6 +16,7 @@ const PORT = 7880;
 const URL = `ws://${HOST}:${PORT}`;
 const API_KEY = 'devkey';
 const API_SECRET = 'secret';
+const TEST_FILES = ['src/tests/e2e.test.ts', 'src/tests/e2e_data_streams.test.ts'];
 
 async function tcpReady(host, port, timeoutMs) {
   const deadline = Date.now() + timeoutMs;
@@ -77,7 +78,7 @@ process.on('SIGTERM', () => onSignal('SIGTERM'));
 try {
   await tcpReady(HOST, PORT, 15_000);
 
-  const args = ['exec', 'vitest', 'run', 'src/tests/e2e.test.ts', ...process.argv.slice(2)];
+  const args = ['exec', 'vitest', 'run', ...TEST_FILES, ...process.argv.slice(2)];
   testProc = spawn('pnpm', args, {
     env: {
       ...process.env,

packages/livekit-rtc/src/audio_stream_room_lifecycle.test.ts

@@ -146,12 +146,15 @@ function makeEndTrackingStream(): { stream: AudioStreamSource; endCount: () => n
 
 function makeLocalParticipant(identity: string): LocalParticipant {
   // Bypass the FFI-touching constructor; set only the fields the lifecycle
-  // paths read (identity getter + trackPublications map).
+  // paths read (identity getter, trackPublications map, and the open data
+  // stream writers that disconnect cleanup disposes).
   const p = Object.create(LocalParticipant.prototype) as LocalParticipant;
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   (p as any).info = { identity };
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   (p as any).trackPublications = new Map();
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  (p as any).openStreamWriters = new Set();
   return p;
 }
 

packages/livekit-rtc/src/data_streams/stream_reader.ts

@@ -3,7 +3,6 @@
 // SPDX-License-Identifier: Apache-2.0
 import type { DataStream_Chunk } from '@livekit/rtc-ffi-bindings';
 import { log } from '../log.js';
-import { bigIntToNumber } from '../utils.js';
 import type { BaseStreamInfo, ByteStreamInfo, TextStreamInfo } from './types.js';
 
 abstract class BaseStreamReader<T extends BaseStreamInfo> {
@@ -15,6 +14,12 @@ abstract class BaseStreamReader<T extends BaseStreamInfo> {
 
   protected bytesReceived: number;
 
+  private closed = false;
+
+  // The reader held by an in-progress iteration. Cancelling has to go through
+  // it, since it holds the stream's lock.
+  private activeReader: ReadableStreamDefaultReader<DataStream_Chunk> | null = null;
+
   get info() {
     return this._info;
   }
@@ -26,7 +31,68 @@ abstract class BaseStreamReader<T extends BaseStreamInfo> {
     this.bytesReceived = 0;
   }
 
-  protected abstract handleChunkReceived(chunk: DataStream_Chunk): void;
+  protected handleChunkReceived(chunk: DataStream_Chunk) {
+    this.bytesReceived += chunk.content!.byteLength;
+    const currentProgress = this.totalByteSize
+      ? this.bytesReceived / this.totalByteSize
+      : undefined;
+    this.onProgress?.(currentProgress);
+  }
+
+  /** Takes the stream's reader for an iteration, remembering it so that
+   * close() can cancel a read that is still in flight. */
+  protected acquireReader(): ReadableStreamDefaultReader<DataStream_Chunk> {
+    const reader = this.reader.getReader();
+    this.activeReader = reader;
+    return reader;
+  }
+
+  /** Releases the stream lock after an iteration ended on its own — because a
+   * read reported end-of-stream, or because the stream itself errored. Either
+   * way the stream is in a terminal state and the FFI subscription behind it
+   * was already dropped by that terminal event, so there is nothing to cancel.
+   *
+   * An error raised *around* a read rather than by it — a chunk that failed to
+   * process — leaves the stream live, so it has to go through close() instead. */
+  protected finishIteration(reader: ReadableStreamDefaultReader<DataStream_Chunk>) {
+    this.closed = true;
+    if (this.activeReader === reader) {
+      this.activeReader = null;
+    }
+    reader.releaseLock();
+  }
+
+  /**
+   * Stops receiving this stream and releases the resources behind it.
+   *
+   * A reader is subscribed to FFI events from the moment it is handed to a
+   * stream handler, and only unsubscribes once a read consumes the stream's
+   * end-of-stream event. So a reader that is abandoned part-way through, or
+   * that a handler never reads at all, has to be closed here or its
+   * subscription lives for the rest of the process. Reading to completion — or
+   * breaking out of a `for await` — closes the reader for you, and closing an
+   * already-closed reader is a no-op.
+   */
+  async close(): Promise<void> {
+    if (this.closed) {
+      return;
+    }
+    this.closed = true;
+    const active = this.activeReader;
+    this.activeReader = null;
+    try {
+      if (active) {
+        await active.cancel();
+        active.releaseLock();
+      } else {
+        await this.reader.cancel();
+      }
+    } catch (error: unknown) {
+      // The stream was already closed or errored (e.g. the room disconnected
+      // mid-stream), in which case its subscription is already gone.
+      log.debug('error closing stream reader: %s', error);
+    }
+  }
 
   onProgress?: (progress: number | undefined) => void;
 
@@ -37,41 +103,48 @@ abstract class BaseStreamReader<T extends BaseStreamInfo> {
  * A class to read chunks from a ReadableStream and provide them in a structured format.
  */
 export class ByteStreamReader extends BaseStreamReader<ByteStreamInfo> {
-  protected handleChunkReceived(chunk: DataStream_Chunk) {
-    this.bytesReceived += chunk.content!.byteLength;
-    const currentProgress = this.totalByteSize
-      ? this.bytesReceived / this.totalByteSize
-      : undefined;
-    this.onProgress?.(currentProgress);
-  }
-
   [Symbol.asyncIterator]() {
-    const reader = this.reader.getReader();
+    const reader = this.acquireReader();
 
     return {
       next: async (): Promise<IteratorResult<Uint8Array>> => {
+        let read: Awaited<ReturnType<typeof reader.read>>;
         try {
-          const { done, value } = await reader.read();
-          if (done) {
-            // Release the lock when the stream is exhausted so the
-            // underlying ReadableStream can be garbage-collected.
-            reader.releaseLock();
-            return { done: true, value: undefined as unknown };
-          } else {
-            this.handleChunkReceived(value);
-            return { done: false, value: value.content! };
-          }
+          read = await reader.read();
         } catch (error: unknown) {
-          // Release the lock on error so it doesn't stay held when the
-          // consumer never calls return() (e.g. breaking out of for-await).
-          reader.releaseLock();
-          log.error('error processing stream update: %s', error);
+          // The stream errored, which is terminal: release the lock so it
+          // doesn't stay held when the consumer never calls return() (a
+          // rejecting next() doesn't trigger it).
+          this.finishIteration(reader);
+          log.error('error reading stream: %s', error);
+          // Propagate abnormal termination (e.g. remote abort, payload over
+          // the receiver's size limit) instead of presenting the truncated
+          // payload as a clean EOF.
+          throw error;
+        }
+
+        if (read.done) {
+          // Release the lock when the stream is exhausted so the
+          // underlying ReadableStream can be garbage-collected.
+          this.finishIteration(reader);
           return { done: true, value: undefined as unknown };
         }
+
+        try {
+          this.handleChunkReceived(read.value);
+          return { done: false, value: read.value.content! };
+        } catch (error: unknown) {
+          // The chunk arrived fine but handling it threw (e.g. a consumer's
+          // onProgress callback). The stream is still live, so it has to be
+          // cancelled to release its FFI subscription.
+          await this.close();
+          log.error('error processing stream update: %s', error);
+          throw error;
+        }
       },
 
-      return(): IteratorResult<Uint8Array> {
-        reader.releaseLock();
+      return: async (): Promise<IteratorResult<Uint8Array>> => {
+        await this.close();
         return { done: true, value: undefined };
       },
     };
@@ -90,77 +163,54 @@ export class ByteStreamReader extends BaseStreamReader<ByteStreamInfo> {
  * A class to read chunks from a ReadableStream and provide them in a structured format.
  */
 export class TextStreamReader extends BaseStreamReader<TextStreamInfo> {
-  private receivedChunks: Map<number, DataStream_Chunk>;
-
-  /**
-   * A TextStreamReader instance can be used as an AsyncIterator that returns the entire string
-   * that has been received up to the current point in time.
-   */
-  constructor(
-    info: TextStreamInfo,
-    stream: ReadableStream<DataStream_Chunk>,
-    totalChunkCount?: number,
-  ) {
-    super(info, stream, totalChunkCount);
-    this.receivedChunks = new Map();
-  }
-
-  protected handleChunkReceived(chunk: DataStream_Chunk) {
-    const index = bigIntToNumber(chunk.chunkIndex!);
-    const previousChunkAtIndex = this.receivedChunks.get(index!);
-    if (previousChunkAtIndex && previousChunkAtIndex.version! > chunk.version!) {
-      // we have a newer version already, dropping the old one
-      return;
-    }
-    this.receivedChunks.set(index, chunk);
-    const currentProgress = this.totalByteSize
-      ? this.receivedChunks.size / this.totalByteSize
-      : undefined;
-    this.onProgress?.(currentProgress);
-  }
-
   /**
    * Async iterator implementation to allow usage of `for await...of` syntax.
    * Yields structured chunks from the stream.
    *
    */
   [Symbol.asyncIterator]() {
-    const reader = this.reader.getReader();
+    const reader = this.acquireReader();
     const decoder = new TextDecoder();
-    const receivedChunks = this.receivedChunks;
 
     return {
       next: async (): Promise<IteratorResult<string>> => {
+        let read: Awaited<ReturnType<typeof reader.read>>;
         try {
-          const { done, value } = await reader.read();
-          if (done) {
-            // Release the lock when the stream is exhausted so the
-            // underlying ReadableStream can be garbage-collected.
-            reader.releaseLock();
-            // Clear received chunks so the buffered data can be GC'd.
-            receivedChunks.clear();
-            return { done: true, value: undefined };
-          } else {
-            this.handleChunkReceived(value);
-            return {
-              done: false,
-              value: decoder.decode(value.content!),
-            };
-          }
+          read = await reader.read();
         } catch (error: unknown) {
-          // Release the lock on error so it doesn't stay held when the
-          // consumer never calls return() (e.g. breaking out of for-await).
-          reader.releaseLock();
-          receivedChunks.clear();
-          log.error('error processing stream update: %s', error);
+          // The stream errored, which is terminal: release the lock so it
+          // doesn't stay held when the consumer never calls return() (a
+          // rejecting next() doesn't trigger it).
+          this.finishIteration(reader);
+          log.error('error reading stream: %s', error);
+          // Propagate abnormal termination (e.g. remote abort, payload over
+          // the receiver's size limit) instead of presenting the truncated
+          // payload as a clean EOF.
+          throw error;
+        }
+
+        if (read.done) {
+          // Release the lock when the stream is exhausted so the
+          // underlying ReadableStream can be garbage-collected.
+          this.finishIteration(reader);
           return { done: true, value: undefined };
         }
+
+        try {
+          this.handleChunkReceived(read.value);
+          return { done: false, value: decoder.decode(read.value.content!) };
+        } catch (error: unknown) {
+          // The chunk arrived fine but handling it threw (e.g. a consumer's
+          // onProgress callback). The stream is still live, so it has to be
+          // cancelled to release its FFI subscription.
+          await this.close();
+          log.error('error processing stream update: %s', error);
+          throw error;
+        }
       },
 
-      return(): IteratorResult<string> {
-        reader.releaseLock();
-        // Clear received chunks so the buffered data can be GC'd.
-        receivedChunks.clear();
+      return: async (): Promise<IteratorResult<string>> => {
+        await this.close();
         return { done: true, value: undefined };
       },
     };

packages/livekit-rtc/src/data_streams/types.ts

@@ -42,6 +42,12 @@ export interface TextStreamOptions extends DataStreamOptions {
 export interface ByteStreamOptions extends DataStreamOptions {
   name?: string;
   onProgress?: (progress: number) => void;
+  /**
+   * Whether the payload may be compressed on the wire. Defaults to true;
+   * compression is only applied when it actually reduces the payload size
+   * and every recipient supports it.
+   */
+  compress?: boolean;
 }
 
 export type ByteStreamHandler = (