diff --git a/.changeset/cancel-incoming-stream-on-early-exit.md b/.changeset/cancel-incoming-stream-on-early-exit.md new file mode 100644 index 0000000000..2335477952 --- /dev/null +++ b/.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. diff --git a/src/room/data-stream/incoming/IncomingDataStreamManager.test.ts b/src/room/data-stream/incoming/IncomingDataStreamManager.test.ts index b2d602b281..808e64a0b7 100644 --- a/src/room/data-stream/incoming/IncomingDataStreamManager.test.ts +++ b/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 = []; + 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'); + }); }); }); diff --git a/src/room/data-stream/incoming/StreamReader.ts b/src/room/data-stream/incoming/StreamReader.ts index ffbec375e7..28f63cee0e 100644 --- a/src/room/data-stream/incoming/StreamReader.ts +++ b/src/room/data-stream/incoming/StreamReader.ts @@ -78,7 +78,12 @@ export class ByteStreamReader extends BaseStreamReader { // 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 { return { done: false, value: result.value.content }; } } catch (err) { - cleanup(); + cleanup(err); throw err; } }, @@ -177,7 +182,12 @@ export class TextStreamReader extends BaseStreamReader { 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 { }; } } catch (err) { - cleanup(); + cleanup(err); throw err; } },