Cancel incoming data streams when a reader stops early

.changeset/cancel-incoming-stream-on-early-exit.md

@@ -0,0 +1,5 @@
+---
+'livekit-client': patch
+---
+
+Cancel an incoming data stream when its reader stops early (an aborted `signal`, a `break` out of `for await`, or a read error), so the stream is unregistered and later packets for it are dropped instead of buffering until the trailer arrives. A reader can no longer resume a stream after leaving its iteration early.

src/room/data-stream/incoming/IncomingDataStreamManager.test.ts

@@ -1884,5 +1884,58 @@ describe('IncomingDataStreamManager', () => {
       expect(readers).toHaveLength(2);
       expect(await readers[1].readAll()).toEqual('fresh stream');
     });
+
+    it('should treat packets for a stream whose read was aborted as an unknown stream', async () => {
+      const manager = new IncomingDataStreamManager();
+      manager.setConnected(true);
+
+      const readers: Array<TextStreamReader> = [];
+      manager.registerTextStreamHandler('my-topic', (reader) => readers.push(reader));
+
+      const streamId = crypto.randomUUID();
+      manager.handleDataStreamPacket(
+        headerPacket(streamId, 'textHeader', { compression: DataStream_CompressionType.NONE }),
+        Encryption_Type.NONE,
+      );
+      manager.handleDataStreamPacket(
+        chunkPacket(streamId, 0, new TextEncoder().encode('first')),
+        Encryption_Type.NONE,
+      );
+      expect(readers).toHaveLength(1);
+
+      // The consumer gives up on the stream part way through.
+      const abortController = new AbortController();
+      const readAllPromise = readers[0].readAll({ signal: abortController.signal });
+      abortController.abort(new Error('consumer aborted'));
+      await expect(readAllPromise).rejects.toThrow('consumer aborted');
+      // The cancel propagates back up the pipe chain to the source stream asynchronously.
+      await new Promise((resolve) => setTimeout(resolve, 0));
+
+      // Packets still in flight for the aborted stream are dropped, just like packets for a stream
+      // that was never opened, rather than buffering in a stream nobody reads.
+      expect(() =>
+        manager.handleDataStreamPacket(
+          chunkPacket(streamId, 1, new TextEncoder().encode('second')),
+          Encryption_Type.NONE,
+        ),
+      ).not.toThrow();
+      expect(readers).toHaveLength(1);
+
+      // No state is left behind for the stream id, even though no trailer has arrived to close it: a
+      // header reusing it opens a brand new stream instead of failing with `AlreadyOpened`, and that
+      // stream reads cleanly from chunk 0.
+      manager.handleDataStreamPacket(
+        headerPacket(streamId, 'textHeader', { compression: DataStream_CompressionType.NONE }),
+        Encryption_Type.NONE,
+      );
+      manager.handleDataStreamPacket(
+        chunkPacket(streamId, 0, new TextEncoder().encode('fresh stream')),
+        Encryption_Type.NONE,
+      );
+      manager.handleDataStreamPacket(trailerPacket(streamId), Encryption_Type.NONE);
+
+      expect(readers).toHaveLength(2);
+      expect(await readers[1].readAll()).toEqual('fresh stream');
+    });
   });
 });

src/room/data-stream/incoming/StreamReader.ts

@@ -78,7 +78,12 @@ export class ByteStreamReader extends BaseStreamReader<ByteStreamInfo> {
     // already propagated through reader.read() to the consumer.
     reader.closed.catch(() => {});
 
-    const cleanup = () => {
+    // Runs when iteration stops early (abort, `break`/`return`, or an error). The stream is
+    // cancelled rather than just unlocked, so the manager stops buffering chunks nobody will read
+    // and unregisters the stream - later packets for it are then dropped.
+    const cleanup = (reason?: unknown) => {
+      // The stream may already be errored, in which case there is nothing left to cancel.
+      reader.cancel(reason).catch(() => {});
       reader.releaseLock();
       this.signal = undefined;
     };
@@ -117,7 +122,7 @@ export class ByteStreamReader extends BaseStreamReader<ByteStreamInfo> {
             return { done: false, value: result.value.content };
           }
         } catch (err) {
-          cleanup();
+          cleanup(err);
           throw err;
         }
       },
@@ -177,7 +182,12 @@ export class TextStreamReader extends BaseStreamReader<TextStreamInfo> {
     const decoder = new TextDecoder('utf-8');
     const signal = this.signal;
 
-    const cleanup = () => {
+    // Runs when iteration stops early (abort, `break`/`return`, or an error). The stream is
+    // cancelled rather than just unlocked, so the manager stops buffering chunks nobody will read
+    // and unregisters the stream - later packets for it are then dropped.
+    const cleanup = (reason?: unknown) => {
+      // The stream may already be errored, in which case there is nothing left to cancel.
+      reader.cancel(reason).catch(() => {});
       reader.releaseLock();
       this.signal = undefined;
     };
@@ -229,7 +239,7 @@ export class TextStreamReader extends BaseStreamReader<TextStreamInfo> {
             };
           }
         } catch (err) {
-          cleanup();
+          cleanup(err);
           throw err;
         }
       },