11import type { DataStream_Chunk } from '@livekit/protocol' ;
22import type { BaseStreamInfo , ByteStreamInfo , TextStreamInfo } from './types' ;
3- import { bigIntToNumber } from './utils' ;
3+ import { Future , bigIntToNumber } from './utils' ;
4+
5+ export type BaseStreamReaderReadAllOpts = {
6+ /** An AbortSignal can be used to terminate reads early. */
7+ signal ?: AbortSignal ;
8+ } ;
49
510abstract class BaseStreamReader < T extends BaseStreamInfo > {
611 protected reader : ReadableStream < DataStream_Chunk > ;
@@ -26,7 +31,7 @@ abstract class BaseStreamReader<T extends BaseStreamInfo> {
2631
2732 onProgress ?: ( progress : number | undefined ) => void ;
2833
29- abstract readAll ( ) : Promise < string | Array < Uint8Array > > ;
34+ abstract readAll ( opts ?: BaseStreamReaderReadAllOpts ) : Promise < string | Array < Uint8Array > > ;
3035}
3136
3237export class ByteStreamReader extends BaseStreamReader < ByteStreamInfo > {
@@ -40,35 +45,77 @@ export class ByteStreamReader extends BaseStreamReader<ByteStreamInfo> {
4045
4146 onProgress ?: ( progress : number | undefined ) => void ;
4247
48+ signal ?: AbortSignal ;
49+
4350 [ Symbol . asyncIterator ] ( ) {
4451 const reader = this . reader . getReader ( ) ;
4552
53+ let rejectingSignalFuture = new Future < never > ( ) ;
54+ let activeSignal : AbortSignal | null = null ;
55+ let onAbort : ( ( ) => void ) | null = null ;
56+ if ( this . signal ) {
57+ const signal = this . signal ;
58+ onAbort = ( ) => {
59+ rejectingSignalFuture . reject ?.( signal . reason ) ;
60+ } ;
61+ signal . addEventListener ( 'abort' , onAbort ) ;
62+ activeSignal = signal ;
63+ }
64+
65+ const cleanup = ( ) => {
66+ reader . releaseLock ( ) ;
67+
68+ if ( activeSignal && onAbort ) {
69+ activeSignal . removeEventListener ( 'abort' , onAbort ) ;
70+ }
71+
72+ this . signal = undefined ;
73+ } ;
74+
4675 return {
4776 next : async ( ) : Promise < IteratorResult < Uint8Array > > => {
4877 try {
49- const { done, value } = await reader . read ( ) ;
78+ const { done, value } = await Promise . race ( [
79+ reader . read ( ) ,
80+ rejectingSignalFuture . promise ,
81+ ] ) ;
5082 if ( done ) {
5183 return { done : true , value : undefined as any } ;
5284 } else {
5385 this . handleChunkReceived ( value ) ;
5486 return { done : false , value : value . content } ;
5587 }
56- } catch ( error ) {
57- // TODO handle errors
58- return { done : true , value : undefined } ;
88+ } catch ( err ) {
89+ cleanup ( ) ;
90+ throw err ;
5991 }
6092 } ,
6193
94+ // note: `return` runs only for premature exits, see:
95+ // https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Iteration_protocols#errors_during_iteration
6296 async return ( ) : Promise < IteratorResult < Uint8Array > > {
63- reader . releaseLock ( ) ;
97+ cleanup ( ) ;
6498 return { done : true , value : undefined } ;
6599 } ,
66100 } ;
67101 }
68102
69- async readAll ( ) : Promise < Array < Uint8Array > > {
103+ /**
104+ * Injects an AbortSignal, which if aborted, will terminate the currently active
105+ * stream iteration operation.
106+ *
107+ * Note that when using AbortSignal.timeout(...), the timeout applies across
108+ * the whole iteration operation, not just one individual chunk read.
109+ */
110+ withAbortSignal ( signal : AbortSignal ) {
111+ this . signal = signal ;
112+ return this ;
113+ }
114+
115+ async readAll ( opts : BaseStreamReaderReadAllOpts = { } ) : Promise < Array < Uint8Array > > {
70116 let chunks : Set < Uint8Array > = new Set ( ) ;
71- for await ( const chunk of this ) {
117+ const iterator = opts . signal ? this . withAbortSignal ( opts . signal ) : this ;
118+ for await ( const chunk of iterator ) {
72119 chunks . add ( chunk ) ;
73120 }
74121 return Array . from ( chunks ) ;
@@ -81,6 +128,8 @@ export class ByteStreamReader extends BaseStreamReader<ByteStreamInfo> {
81128export class TextStreamReader extends BaseStreamReader < TextStreamInfo > {
82129 private receivedChunks : Map < number , DataStream_Chunk > ;
83130
131+ signal ?: AbortSignal ;
132+
84133 /**
85134 * A TextStreamReader instance can be used as an AsyncIterator that returns the entire string
86135 * that has been received up to the current point in time.
@@ -123,10 +172,35 @@ export class TextStreamReader extends BaseStreamReader<TextStreamInfo> {
123172 const reader = this . reader . getReader ( ) ;
124173 const decoder = new TextDecoder ( ) ;
125174
175+ let rejectingSignalFuture = new Future < never > ( ) ;
176+ let activeSignal : AbortSignal | null = null ;
177+ let onAbort : ( ( ) => void ) | null = null ;
178+ if ( this . signal ) {
179+ const signal = this . signal ;
180+ onAbort = ( ) => {
181+ rejectingSignalFuture . reject ?.( signal . reason ) ;
182+ } ;
183+ signal . addEventListener ( 'abort' , onAbort ) ;
184+ activeSignal = signal ;
185+ }
186+
187+ const cleanup = ( ) => {
188+ reader . releaseLock ( ) ;
189+
190+ if ( activeSignal && onAbort ) {
191+ activeSignal . removeEventListener ( 'abort' , onAbort ) ;
192+ }
193+
194+ this . signal = undefined ;
195+ } ;
196+
126197 return {
127198 next : async ( ) : Promise < IteratorResult < string > > => {
128199 try {
129- const { done, value } = await reader . read ( ) ;
200+ const { done, value } = await Promise . race ( [
201+ reader . read ( ) ,
202+ rejectingSignalFuture . promise ,
203+ ] ) ;
130204 if ( done ) {
131205 return { done : true , value : undefined } ;
132206 } else {
@@ -137,22 +211,37 @@ export class TextStreamReader extends BaseStreamReader<TextStreamInfo> {
137211 value : decoder . decode ( value . content ) ,
138212 } ;
139213 }
140- } catch ( error ) {
141- // TODO handle errors
142- return { done : true , value : undefined } ;
214+ } catch ( err ) {
215+ cleanup ( ) ;
216+ throw err ;
143217 }
144218 } ,
145219
220+ // note: `return` runs only for premature exits, see:
221+ // https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Iteration_protocols#errors_during_iteration
146222 async return ( ) : Promise < IteratorResult < string > > {
147- reader . releaseLock ( ) ;
223+ cleanup ( ) ;
148224 return { done : true , value : undefined } ;
149225 } ,
150226 } ;
151227 }
152228
153- async readAll ( ) : Promise < string > {
229+ /**
230+ * Injects an AbortSignal, which if aborted, will terminate the currently active
231+ * stream iteration operation.
232+ *
233+ * Note that when using AbortSignal.timeout(...), the timeout applies across
234+ * the whole iteration operation, not just one individual chunk read.
235+ */
236+ withAbortSignal ( signal : AbortSignal ) {
237+ this . signal = signal ;
238+ return this ;
239+ }
240+
241+ async readAll ( opts : BaseStreamReaderReadAllOpts = { } ) : Promise < string > {
154242 let finalString : string = '' ;
155- for await ( const chunk of this ) {
243+ const iterator = opts . signal ? this . withAbortSignal ( opts . signal ) : this ;
244+ for await ( const chunk of iterator ) {
156245 finalString += chunk ;
157246 }
158247 return finalString ;
0 commit comments