Skip to content

Commit ab5ed72

Browse files
trivikraduh95
trivikr
authored and
aduh95
committed
stream: reject iter consumers on abort
Make stream/iter async consumers observe abort signals while waiting for a pending async iterator read. This lets bytes(), text(), arrayBuffer(), and array() reject promptly with the abort reason instead of waiting for another batch. Move the shared abort-aware iterator wrapper to stream/iter utils so pull and consumers use the same helper. Signed-off-by: Kamat, Trivikram <16024985+trivikr@users.noreply.github.com> Assisted-by: openai:gpt-5.5 PR-URL: #64066 Fixes: #64065 Reviewed-By: James M Snell <jasnell@gmail.com>
1 parent d3fa77c commit ab5ed72

4 files changed

Lines changed: 143 additions & 84 deletions

File tree

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

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@ const {
5252

5353
const {
5454
concatBytes,
55+
yieldAbortable,
5556
} = require('internal/streams/iter/utils');
5657

5758
const {
@@ -125,7 +126,9 @@ async function collectAsync(source, signal, limit) {
125126
signal?.throwIfAborted();
126127

127128
// Normalize source via from() - accepts strings, ArrayBuffers, protocols, etc.
128-
const normalized = from(source);
129+
const abortableSource = signal && isAsyncIterable(source) ?
130+
yieldAbortable(source, signal) : source;
131+
const normalized = from(abortableSource);
129132
const chunks = [];
130133

131134
// Fast path: no signal and no limit
@@ -140,8 +143,9 @@ async function collectAsync(source, signal, limit) {
140143

141144
// Slow path: with signal or limit checks
142145
let totalBytes = 0;
146+
const iterable = signal ? yieldAbortable(normalized, signal) : normalized;
143147

144-
for await (const batch of normalized) {
148+
for await (const batch of iterable) {
145149
signal?.throwIfAborted();
146150
for (let i = 0; i < batch.length; i++) {
147151
const chunk = batch[i];

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

Lines changed: 1 addition & 82 deletions
Original file line numberDiff line numberDiff line change
@@ -14,19 +14,12 @@ const {
1414
ArrayPrototypeSlice,
1515
PromisePrototypeThen,
1616
PromiseResolve,
17-
PromiseWithResolvers,
18-
SafePromisePrototypeFinally,
19-
SafePromiseRace,
2017
SymbolAsyncIterator,
2118
SymbolIterator,
2219
TypedArrayPrototypeGetByteLength,
2320
Uint8Array,
2421
} = primordials;
2522

26-
const {
27-
markPromiseAsHandled,
28-
} = internalBinding('util');
29-
3023
const {
3124
codes: {
3225
ERR_INVALID_ARG_TYPE,
@@ -60,6 +53,7 @@ const {
6053
parsePullArgs,
6154
toUint8Array,
6255
wrapError,
56+
yieldAbortable,
6357
} = require('internal/streams/iter/utils');
6458

6559
const {
@@ -690,81 +684,6 @@ async function* applyValidatedStatefulAsyncTransform(source, transform, options)
690684
options.signal?.throwIfAborted();
691685
}
692686

693-
function getOnAbort(reject, signal) {
694-
return () => reject(signal.reason);
695-
}
696-
697-
/**
698-
* Read one item from an async iterator, rejecting early if the signal aborts.
699-
* @param {AsyncIterator} iterator - The iterator to read from.
700-
* @param {AbortSignal|undefined} signal - Optional abort signal.
701-
* @returns {Promise<IteratorResult<Uint8Array[]>>|IteratorResult<Uint8Array[]>}
702-
*/
703-
function abortableNext(iterator, signal) {
704-
if (signal === undefined) {
705-
return iterator.next();
706-
}
707-
708-
signal.throwIfAborted();
709-
710-
const next = iterator.next();
711-
const { promise, reject } = PromiseWithResolvers();
712-
const onAbort = getOnAbort(reject, signal);
713-
signal.addEventListener('abort', onAbort, { __proto__: null, once: true });
714-
if (signal.aborted) {
715-
onAbort();
716-
}
717-
718-
return SafePromisePrototypeFinally(SafePromiseRace([next, promise]), () => {
719-
signal.removeEventListener('abort', onAbort);
720-
});
721-
}
722-
723-
/**
724-
* Wrap an async source so each pending read is abort-aware.
725-
* @param {AsyncIterable<Uint8Array[]>} source - The source to read from.
726-
* @param {AbortSignal|undefined} signal - Optional abort signal.
727-
* @returns {AsyncIterable<Uint8Array[]>}
728-
*/
729-
function yieldAbortable(source, signal) {
730-
if (signal === undefined) {
731-
return source;
732-
}
733-
734-
return {
735-
__proto__: null,
736-
async *[SymbolAsyncIterator]() {
737-
const iterator = source[SymbolAsyncIterator]();
738-
let completed = false;
739-
let aborted = false;
740-
741-
try {
742-
while (true) {
743-
const { done, value } = await abortableNext(iterator, signal);
744-
if (done) {
745-
completed = true;
746-
return;
747-
}
748-
signal.throwIfAborted();
749-
yield value;
750-
}
751-
} catch (error) {
752-
aborted = signal.aborted;
753-
throw error;
754-
} finally {
755-
if (!completed && typeof iterator.return === 'function') {
756-
const result = iterator.return();
757-
if (aborted) {
758-
markPromiseAsHandled(result);
759-
} else {
760-
await result;
761-
}
762-
}
763-
}
764-
},
765-
};
766-
}
767-
768687
/**
769688
* Create an async pipeline from source through transforms.
770689
* @yields {Uint8Array[]}

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

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,14 +8,22 @@ const {
88
MathMin,
99
NumberMAX_SAFE_INTEGER,
1010
PromiseResolve,
11+
PromiseWithResolvers,
12+
SafePromisePrototypeFinally,
13+
SafePromiseRace,
1114
String,
15+
SymbolAsyncIterator,
1216
TypedArrayPrototypeGetBuffer,
1317
TypedArrayPrototypeGetByteLength,
1418
TypedArrayPrototypeGetByteOffset,
1519
TypedArrayPrototypeSet,
1620
Uint8Array,
1721
} = primordials;
1822

23+
const {
24+
markPromiseAsHandled,
25+
} = internalBinding('util');
26+
1927
const { TextEncoder } = require('internal/encoding');
2028
const {
2129
codes: {
@@ -69,6 +77,81 @@ function onSignalAbort(signal, handler) {
6977
}
7078
}
7179

80+
function getOnAbort(reject, signal) {
81+
return () => reject(signal.reason);
82+
}
83+
84+
/**
85+
* Read one item from an async iterator, rejecting early if the signal aborts.
86+
* @param {AsyncIterator} iterator - The iterator to read from.
87+
* @param {AbortSignal|undefined} signal - Optional abort signal.
88+
* @returns {Promise<IteratorResult<Uint8Array[]>>|IteratorResult<Uint8Array[]>}
89+
*/
90+
function abortableNext(iterator, signal) {
91+
if (signal === undefined) {
92+
return iterator.next();
93+
}
94+
95+
signal.throwIfAborted();
96+
97+
const next = iterator.next();
98+
const { promise, reject } = PromiseWithResolvers();
99+
const onAbort = getOnAbort(reject, signal);
100+
signal.addEventListener('abort', onAbort, { __proto__: null, once: true });
101+
if (signal.aborted) {
102+
onAbort();
103+
}
104+
105+
return SafePromisePrototypeFinally(SafePromiseRace([next, promise]), () => {
106+
signal.removeEventListener('abort', onAbort);
107+
});
108+
}
109+
110+
/**
111+
* Wrap an async source so each pending read is abort-aware.
112+
* @param {AsyncIterable<Uint8Array[]>} source - The source to read from.
113+
* @param {AbortSignal|undefined} signal - Optional abort signal.
114+
* @returns {AsyncIterable<Uint8Array[]>}
115+
*/
116+
function yieldAbortable(source, signal) {
117+
if (signal === undefined) {
118+
return source;
119+
}
120+
121+
return {
122+
__proto__: null,
123+
async *[SymbolAsyncIterator]() {
124+
const iterator = source[SymbolAsyncIterator]();
125+
let completed = false;
126+
let aborted = false;
127+
128+
try {
129+
while (true) {
130+
const { done, value } = await abortableNext(iterator, signal);
131+
if (done) {
132+
completed = true;
133+
return;
134+
}
135+
signal.throwIfAborted();
136+
yield value;
137+
}
138+
} catch (error) {
139+
aborted = signal.aborted;
140+
throw error;
141+
} finally {
142+
if (!completed && typeof iterator.return === 'function') {
143+
const result = iterator.return();
144+
if (aborted) {
145+
markPromiseAsHandled(result);
146+
} else {
147+
await result;
148+
}
149+
}
150+
}
151+
},
152+
};
153+
}
154+
72155
/**
73156
* Compute the minimum cursor across a set of consumers and count how many
74157
* consumers are at that cursor.
@@ -301,4 +384,5 @@ module.exports = {
301384
toUint8Array,
302385
validateBackpressure,
303386
wrapError,
387+
yieldAbortable,
304388
};

‎test/parallel/test-stream-iter-consumers-bytes.js‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ const {
1414
arrayBufferSync,
1515
array,
1616
arraySync,
17+
toAsyncStreamable,
1718
} = require('stream/iter');
1819

1920
// =============================================================================
@@ -51,6 +52,55 @@ async function testBytesAsyncAbort() {
5152
);
5253
}
5354

55+
async function testAsyncConsumersAbortPendingNext() {
56+
const consumers = [
57+
['bytes', bytes],
58+
['text', text],
59+
['arrayBuffer', arrayBuffer],
60+
['array', array],
61+
];
62+
63+
for (const [name, consumer] of consumers) {
64+
const ac = new AbortController();
65+
const reason = new Error(`${name} boom`);
66+
67+
async function* never() {
68+
await new Promise(() => {});
69+
yield [];
70+
}
71+
72+
const promise = consumer(never(), { __proto__: null, signal: ac.signal });
73+
ac.abort(reason);
74+
75+
await assert.rejects(promise, reason);
76+
}
77+
}
78+
79+
async function testAsyncConsumersAbortPendingNormalization() {
80+
const consumers = [
81+
['bytes', bytes],
82+
['text', text],
83+
['arrayBuffer', arrayBuffer],
84+
['array', array],
85+
];
86+
87+
for (const [name, consumer] of consumers) {
88+
const ac = new AbortController();
89+
const reason = new Error(`${name} normalization boom`);
90+
const source = {
91+
__proto__: null,
92+
[toAsyncStreamable]() {
93+
return new Promise(() => {});
94+
},
95+
};
96+
97+
const promise = consumer(source, { __proto__: null, signal: ac.signal });
98+
ac.abort(reason);
99+
100+
await assert.rejects(promise, reason);
101+
}
102+
}
103+
54104
async function testBytesEmpty() {
55105
const data = await bytes(from([]));
56106
assert.ok(data instanceof Uint8Array);
@@ -203,6 +253,8 @@ Promise.all([
203253
testBytesAsync(),
204254
testBytesAsyncLimit(),
205255
testBytesAsyncAbort(),
256+
testAsyncConsumersAbortPendingNext(),
257+
testAsyncConsumersAbortPendingNormalization(),
206258
testBytesEmpty(),
207259
testArrayBufferSyncBasic(),
208260
testArrayBufferAsync(),

0 commit comments

Comments
 (0)