Skip to content

Commit 6fb199a

Browse files
jasnelljasnell
committed
stream: ensure nested async iterators are cancellation aware
Signed-off-by: James M Snell <jasnell@gmail.com> Assisted-by: Opencode PR-URL: #66030 Reviewed-By: Filip Skokan <panva.ip@gmail.com> Reviewed-By: Trivikram Kamat <trivikr.dev@gmail.com>
1 parent e898f5c commit 6fb199a

2 files changed

Lines changed: 385 additions & 30 deletions

File tree

‎lib/internal/streams/iter/from.js‎

Lines changed: 214 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,10 @@ const {
1616
DataViewPrototypeGetByteOffset,
1717
FunctionPrototypeCall,
1818
PromisePrototypeThen,
19+
PromiseResolve,
20+
PromiseWithResolvers,
21+
SafePromiseRace,
22+
Symbol,
1923
SymbolAsyncIterator,
2024
SymbolIterator,
2125
TypedArrayPrototypeGetBuffer,
@@ -29,8 +33,10 @@ const { markPromiseAsHandled } = internalBinding('util');
2933
const {
3034
codes: {
3135
ERR_INVALID_ARG_TYPE,
36+
ERR_INVALID_RETURN_VALUE,
3237
},
3338
} = require('internal/errors');
39+
const { lazyDOMException } = require('internal/util');
3440

3541
const {
3642
isAnyArrayBuffer,
@@ -54,6 +60,81 @@ const {
5460
// Bounds peak memory when arrays flow through transforms, which must
5561
// allocate output for the entire batch at once.
5662
const FROM_BATCH_SIZE = 128;
63+
const kNormalizationCancelled = Symbol('kNormalizationCancelled');
64+
65+
function createNormalizationContext() {
66+
return {
67+
__proto__: null,
68+
cancelled: false,
69+
reason: undefined,
70+
resolve: null,
71+
suppressCleanup: false,
72+
};
73+
}
74+
75+
function cancelNormalization(context, reason, suppressCleanup = false) {
76+
if (context.cancelled) return;
77+
context.cancelled = true;
78+
context.reason = reason;
79+
context.suppressCleanup = suppressCleanup;
80+
context.resolve?.(kNormalizationCancelled);
81+
}
82+
83+
function throwIfNormalizationCancelled(context) {
84+
if (context?.cancelled) throw context.reason;
85+
}
86+
87+
async function waitForNormalization(value, context) {
88+
if (context === undefined) return value;
89+
const { promise, resolve } = PromiseWithResolvers();
90+
if (context.cancelled) {
91+
resolve(kNormalizationCancelled);
92+
} else {
93+
context.resolve = resolve;
94+
}
95+
try {
96+
const result = await SafePromiseRace([
97+
PromiseResolve(value),
98+
promise,
99+
]);
100+
throwIfNormalizationCancelled(context);
101+
return result;
102+
} finally {
103+
if (context.resolve === resolve) context.resolve = null;
104+
}
105+
}
106+
107+
function createNormalizationIterator(createIterator) {
108+
const context = createNormalizationContext();
109+
const iterator = createIterator(context);
110+
return {
111+
__proto__: null,
112+
next(value) {
113+
return FunctionPrototypeCall(iterator.next, iterator, value);
114+
},
115+
return(value) {
116+
cancelNormalization(
117+
context, lazyDOMException('Aborted', 'AbortError'));
118+
return FunctionPrototypeCall(iterator.return, iterator, value);
119+
},
120+
throw(error) {
121+
cancelNormalization(context, error, true);
122+
return FunctionPrototypeCall(iterator.throw, iterator, error);
123+
},
124+
[SymbolAsyncIterator]() {
125+
return this;
126+
},
127+
};
128+
}
129+
130+
function createNormalizationSource(createIterator) {
131+
return {
132+
__proto__: null,
133+
[SymbolAsyncIterator]() {
134+
return createNormalizationIterator(createIterator);
135+
},
136+
};
137+
}
57138

58139
// =============================================================================
59140
// Type Guards and Detection
@@ -256,6 +337,101 @@ function* normalizeSyncSource(source) {
256337
}
257338
}
258339

340+
function yieldNormalizationAbortable(source, context) {
341+
if (context === undefined) return source;
342+
return {
343+
__proto__: null,
344+
[SymbolAsyncIterator]() {
345+
const iteratorMethod = source[SymbolAsyncIterator];
346+
const iterator = FunctionPrototypeCall(iteratorMethod, source);
347+
const nextMethod = iterator.next;
348+
let completed = false;
349+
let closed = false;
350+
let reading = false;
351+
352+
async function closeSource(suppressError) {
353+
if (closed) return;
354+
closed = true;
355+
completed = true;
356+
357+
if (suppressError) {
358+
try {
359+
const returnMethod = iterator.return;
360+
if (typeof returnMethod === 'function') {
361+
const cleanup = PromisePrototypeThen(
362+
PromiseResolve(),
363+
() => FunctionPrototypeCall(returnMethod, iterator));
364+
markPromiseAsHandled(cleanup);
365+
}
366+
} catch {
367+
// Cancellation has precedence over source cleanup errors.
368+
}
369+
return;
370+
}
371+
372+
const returnMethod = iterator.return;
373+
if (typeof returnMethod === 'function') {
374+
const result = await FunctionPrototypeCall(returnMethod, iterator);
375+
if ((typeof result !== 'object' && typeof result !== 'function') ||
376+
result === null) {
377+
throw new ERR_INVALID_RETURN_VALUE(
378+
'an object', 'iterator.return()', result);
379+
}
380+
}
381+
}
382+
383+
return {
384+
__proto__: null,
385+
async next() {
386+
if (completed) {
387+
return { __proto__: null, done: true, value: undefined };
388+
}
389+
throwIfNormalizationCancelled(context);
390+
reading = true;
391+
392+
try {
393+
const next = FunctionPrototypeCall(nextMethod, iterator);
394+
const result = await waitForNormalization(next, context);
395+
if ((typeof result !== 'object' && typeof result !== 'function') ||
396+
result === null) {
397+
throw new ERR_INVALID_RETURN_VALUE(
398+
'an object', 'iterator.next()', result);
399+
}
400+
if (result.done) {
401+
reading = false;
402+
throwIfNormalizationCancelled(context);
403+
completed = true;
404+
closed = true;
405+
return { __proto__: null, done: true, value: result.value };
406+
}
407+
const value = result.value;
408+
reading = false;
409+
throwIfNormalizationCancelled(context);
410+
return { __proto__: null, done: false, value };
411+
} catch (error) {
412+
if (context.cancelled) await closeSource(true);
413+
reading = false;
414+
throw error;
415+
}
416+
},
417+
async return(value) {
418+
await closeSource(
419+
context.suppressCleanup || (context.cancelled && reading));
420+
return { __proto__: null, done: true, value };
421+
},
422+
async throw(error) {
423+
await closeSource(
424+
context.suppressCleanup || (context.cancelled && reading));
425+
throw error;
426+
},
427+
[SymbolAsyncIterator]() {
428+
return this;
429+
},
430+
};
431+
},
432+
};
433+
}
434+
259435
// =============================================================================
260436
// Async Normalization (for from and async contexts)
261437
// =============================================================================
@@ -266,11 +442,15 @@ function* normalizeSyncSource(source) {
266442
* and protocol conversions.
267443
* @yields {Uint8Array}
268444
*/
269-
async function* normalizeAsyncValue(value, allowNestedAsyncStreamables = true) {
445+
async function* normalizeAsyncValue(
446+
value, allowNestedAsyncStreamables = true, context) {
447+
throwIfNormalizationCancelled(context);
448+
270449
// Handle promises first
271450
if (isPromise(value)) {
272-
const resolved = await value;
273-
yield* normalizeAsyncValue(resolved, allowNestedAsyncStreamables);
451+
const resolved = await waitForNormalization(value, context);
452+
yield* normalizeAsyncValue(
453+
resolved, allowNestedAsyncStreamables, context);
274454
return;
275455
}
276456

@@ -293,41 +473,47 @@ async function* normalizeAsyncValue(value, allowNestedAsyncStreamables = true) {
293473
if (hasProtocol(value, toAsyncStreamable)) {
294474
const result = FunctionPrototypeCall(value[toAsyncStreamable], value);
295475
if (isPromise(result)) {
296-
yield* normalizeAsyncValue(await result, allowNestedAsyncStreamables);
476+
yield* normalizeAsyncValue(
477+
await waitForNormalization(result, context),
478+
allowNestedAsyncStreamables,
479+
context);
297480
} else {
298-
yield* normalizeAsyncValue(result, allowNestedAsyncStreamables);
481+
yield* normalizeAsyncValue(
482+
result, allowNestedAsyncStreamables, context);
299483
}
300484
return;
301485
}
302486

303487
// Handle ToStreamable protocol
304488
if (hasProtocol(value, toStreamable)) {
305489
const result = FunctionPrototypeCall(value[toStreamable], value);
306-
yield* normalizeAsyncValue(result, allowNestedAsyncStreamables);
490+
yield* normalizeAsyncValue(result, allowNestedAsyncStreamables, context);
307491
return;
308492
}
309493

310494
// Handle arrays (which are also iterable, but check first for efficiency)
311495
if (ArrayIsArray(value)) {
312496
for (let i = 0; i < value.length; i++) {
313-
yield* normalizeAsyncValue(value[i], allowNestedAsyncStreamables);
497+
yield* normalizeAsyncValue(
498+
value[i], allowNestedAsyncStreamables, context);
314499
}
315500
return;
316501
}
317502

318503
// Handle async iterables (check before sync iterables since some objects
319504
// have both)
320505
if (isAsyncIterable(value)) {
321-
for await (const item of value) {
322-
yield* normalizeAsyncValue(item, allowNestedAsyncStreamables);
506+
const iterable = yieldNormalizationAbortable(value, context);
507+
for await (const item of iterable) {
508+
yield* normalizeAsyncValue(item, allowNestedAsyncStreamables, context);
323509
}
324510
return;
325511
}
326512

327513
// Handle sync iterables
328514
if (isSyncIterable(value)) {
329515
for (const item of value) {
330-
yield* normalizeAsyncValue(item, allowNestedAsyncStreamables);
516+
yield* normalizeAsyncValue(item, allowNestedAsyncStreamables, context);
331517
}
332518
return;
333519
}
@@ -346,10 +532,13 @@ async function* normalizeAsyncValue(value, allowNestedAsyncStreamables = true) {
346532
* @param {AsyncIterable|Iterable} source
347533
* @yields {Uint8Array[]}
348534
*/
349-
async function* normalizeAsyncSource(source) {
535+
async function* normalizeAsyncSource(source, context) {
536+
throwIfNormalizationCancelled(context);
537+
350538
// Prefer async iteration if available
351539
if (isAsyncIterable(source)) {
352-
for await (const value of source) {
540+
const iterable = yieldNormalizationAbortable(source, context);
541+
for await (const value of iterable) {
353542
// Fast path 1: value is already a Uint8Array[] batch
354543
if (isUint8ArrayBatch(value)) {
355544
if (value.length > 0) {
@@ -364,7 +553,7 @@ async function* normalizeAsyncSource(source) {
364553
}
365554
// Slow path: normalize the value
366555
let batch = [];
367-
for await (const chunk of normalizeAsyncValue(value)) {
556+
for await (const chunk of normalizeAsyncValue(value, true, context)) {
368557
ArrayPrototypePush(batch, chunk);
369558
if (batch.length === FROM_BATCH_SIZE) {
370559
yield batch;
@@ -383,6 +572,7 @@ async function* normalizeAsyncSource(source) {
383572
let batch = [];
384573

385574
for (const value of source) {
575+
throwIfNormalizationCancelled(context);
386576
// Fast path 1: value is already a Uint8Array[] batch
387577
if (isUint8ArrayBatch(value)) {
388578
// Flush any accumulated batch first
@@ -408,7 +598,7 @@ async function* normalizeAsyncSource(source) {
408598
batch = [];
409599
}
410600
let asyncBatch = [];
411-
for await (const chunk of normalizeAsyncValue(value, false)) {
601+
for await (const chunk of normalizeAsyncValue(value, false, context)) {
412602
ArrayPrototypePush(asyncBatch, chunk);
413603
if (asyncBatch.length === FROM_BATCH_SIZE) {
414604
yield asyncBatch;
@@ -434,6 +624,12 @@ async function* normalizeAsyncSource(source) {
434624
);
435625
}
436626

627+
async function* normalizeAsyncStreamableResult(result, context) {
628+
const resolved = await waitForNormalization(result, context);
629+
const source = resolved?.[kValidatedSource] ? resolved : from(resolved);
630+
yield* yieldNormalizationAbortable(source, context);
631+
}
632+
437633
// =============================================================================
438634
// Public API: from() and fromSync()
439635
// =============================================================================
@@ -610,19 +806,8 @@ function from(input) {
610806
if (result?.[kValidatedSource]) {
611807
return result;
612808
}
613-
return {
614-
__proto__: null,
615-
async *[SymbolAsyncIterator]() {
616-
// The result may be a Promise. Check validated on both the Promise
617-
// itself (if tagged) and the resolved value.
618-
const resolved = await result;
619-
if (resolved?.[kValidatedSource]) {
620-
yield* resolved[SymbolAsyncIterator]();
621-
return;
622-
}
623-
yield* from(resolved)[SymbolAsyncIterator]();
624-
},
625-
};
809+
return createNormalizationSource(
810+
(context) => normalizeAsyncStreamableResult(result, context));
626811
}
627812

628813
// Check toStreamable protocol (takes precedence over iteration protocols)
@@ -640,7 +825,8 @@ function from(input) {
640825
);
641826
}
642827

643-
return normalizeAsyncSource(input);
828+
return createNormalizationIterator(
829+
(context) => normalizeAsyncSource(input, context));
644830
}
645831

646832
// =============================================================================

0 commit comments

Comments
 (0)