Skip to content

Commit 55ef82a

Browse files
jasnelladuh95
jasnell
authored and
aduh95
committed
stream: make classic reads wake on return/throw/abort
Signed-off-by: James M Snell <[email protected]> Assisted-by: Opencode PR-URL: #66079 Reviewed-By: Trivikram Kamat <[email protected]>
1 parent 4e8df7c commit 55ef82a

3 files changed

Lines changed: 138 additions & 39 deletions

File tree

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

Lines changed: 78 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -186,8 +186,10 @@ async function normalizeBatch(raw) {
186186
// (Uint8Array subclass) and are yielded directly.
187187
const nop = () => {};
188188

189-
async function* createBatchedAsyncIterator(stream, normalize) {
189+
function createBatchedAsyncIterator(stream, normalize) {
190190
let callback = nop;
191+
let canceled = false;
192+
let started = false;
191193

192194
function next(resolve) {
193195
if (this === stream) {
@@ -198,54 +200,91 @@ async function* createBatchedAsyncIterator(stream, normalize) {
198200
}
199201
}
200202

201-
stream.on('readable', next);
202-
203-
let error;
204-
const cleanup = eos(stream, { writable: false }, (err) => {
205-
error = err ? aggregateTwoErrors(error, err) : null;
203+
function wake() {
206204
callback();
207205
callback = nop;
208-
});
206+
}
209207

210-
try {
211-
while (true) {
212-
const chunk = stream.destroyed ? null : stream.read();
213-
if (chunk !== null) {
214-
const batch = [chunk];
215-
while (batch.length < MAX_DRAIN_BATCH &&
216-
stream._readableState?.length > 0) {
217-
const c = stream.read();
218-
if (c === null) break;
219-
ArrayPrototypePush(batch, c);
220-
}
221-
if (normalize !== null) {
222-
const result = await normalize(batch);
223-
if (result !== null) {
224-
yield result;
208+
async function* generate() {
209+
stream.on('readable', next);
210+
211+
let error;
212+
const cleanup = eos(stream, { writable: false }, (err) => {
213+
error = err ? aggregateTwoErrors(error, err) : null;
214+
stream.removeListener('readable', next);
215+
wake();
216+
});
217+
const cleanupAll = () => {
218+
stream.removeListener('close', cleanupAll);
219+
cleanup();
220+
};
221+
stream.on('close', cleanupAll);
222+
223+
try {
224+
while (!canceled) {
225+
const chunk = stream.destroyed ? null : stream.read();
226+
if (chunk !== null) {
227+
const batch = [chunk];
228+
while (batch.length < MAX_DRAIN_BATCH &&
229+
stream._readableState?.length > 0) {
230+
const c = stream.read();
231+
if (c === null) break;
232+
ArrayPrototypePush(batch, c);
233+
}
234+
if (normalize !== null) {
235+
const result = await normalize(batch);
236+
if (result !== null) {
237+
yield result;
238+
}
239+
} else {
240+
yield batch;
225241
}
242+
} else if (error) {
243+
throw error;
244+
} else if (error === null) {
245+
return;
226246
} else {
227-
yield batch;
247+
await new Promise(next);
228248
}
229-
} else if (error) {
230-
throw error;
231-
} else if (error === null) {
232-
return;
233-
} else {
234-
await new Promise(next);
235249
}
236-
}
237-
} catch (err) {
238-
error = aggregateTwoErrors(error, err);
239-
throw error;
240-
} finally {
241-
if (error === undefined ||
242-
(stream._readableState?.autoDestroy)) {
243-
destroyImpl.destroyer(stream, null);
244-
} else {
250+
} catch (err) {
251+
error = aggregateTwoErrors(error, err);
252+
throw error;
253+
} finally {
245254
stream.removeListener('readable', next);
246-
cleanup();
255+
if (error === undefined ||
256+
(stream._readableState?.autoDestroy)) {
257+
destroyImpl.destroyer(stream, null);
258+
if (stream._readableState?.closeEmitted) cleanupAll();
259+
} else {
260+
cleanupAll();
261+
}
247262
}
248263
}
264+
265+
const iterator = generate();
266+
const iteratorNext = iterator.next;
267+
const iteratorReturn = iterator.return;
268+
const iteratorThrow = iterator.throw;
269+
270+
iterator.next = function(value) {
271+
started = true;
272+
return FunctionPrototypeCall(iteratorNext, iterator, value);
273+
};
274+
iterator.return = function(value) {
275+
canceled = true;
276+
wake();
277+
if (!started) stream.destroy();
278+
return FunctionPrototypeCall(iteratorReturn, iterator, value);
279+
};
280+
iterator.throw = function(reason) {
281+
canceled = true;
282+
wake();
283+
if (!started) stream.destroy();
284+
return FunctionPrototypeCall(iteratorThrow, iterator, reason);
285+
};
286+
287+
return iterator;
249288
}
250289

251290
/**

‎test/parallel/test-stream-iter-readable-interop.js‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
const common = require('../common');
88
const assert = require('assert');
99
const { Readable } = require('stream');
10+
const { setImmediate } = require('timers/promises');
1011
const {
1112
from,
1213
pull,
@@ -588,6 +589,38 @@ async function testAbortSignal() {
588589
assert.ok(chunks.length >= 2);
589590
}
590591

592+
async function testReturnWhileReadIsPending() {
593+
const readable = new Readable({ read() {} });
594+
const iterator = from(readable)[Symbol.asyncIterator]();
595+
const pending = iterator.next();
596+
await setImmediate();
597+
598+
assert.deepStrictEqual(await iterator.return('stopped'), {
599+
value: 'stopped',
600+
done: true,
601+
});
602+
assert.deepStrictEqual(await pending, { value: undefined, done: true });
603+
await setImmediate();
604+
assert.strictEqual(readable.destroyed, true);
605+
assert.strictEqual(readable.listenerCount('readable'), 0);
606+
}
607+
608+
async function testAbortWhileReadIsPending() {
609+
const readable = new Readable({ read() {} });
610+
const controller = new AbortController();
611+
const iterator = pull(readable, {
612+
signal: controller.signal,
613+
})[Symbol.asyncIterator]();
614+
const pending = iterator.next();
615+
await setImmediate();
616+
617+
controller.abort();
618+
await assert.rejects(pending, { name: 'AbortError' });
619+
await setImmediate();
620+
assert.strictEqual(readable.destroyed, true);
621+
assert.strictEqual(readable.listenerCount('readable'), 0);
622+
}
623+
591624
// =============================================================================
592625
// kValidatedSource identity - from() returns same object for validated sources
593626
// =============================================================================
@@ -637,4 +670,6 @@ Promise.all([
637670
testDuplexStream(),
638671
testSetEncodingDynamic(),
639672
testAbortSignal(),
673+
testReturnWhileReadIsPending(),
674+
testAbortWhileReadIsPending(),
640675
]).then(common.mustCall());

‎test/parallel/test-stream-iter-to-readable.js‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -157,6 +157,30 @@ async function testFalsyThenableCleanupError() {
157157
assert.strictEqual(result.reason, reason);
158158
}
159159

160+
async function testSynchronousIteratorReturn() {
161+
let returnCalled = false;
162+
const source = {
163+
__proto__: null,
164+
[Symbol.asyncIterator]() {
165+
return {
166+
__proto__: null,
167+
next() { return kNeverResolves; },
168+
return() {
169+
returnCalled = true;
170+
return { __proto__: null, done: true };
171+
},
172+
};
173+
},
174+
};
175+
const readable = toReadable(source);
176+
const { promise, resolve } = Promise.withResolvers();
177+
readable.once('close', resolve);
178+
179+
readable.destroy();
180+
await promise;
181+
assert.strictEqual(returnCalled, true);
182+
}
183+
160184
async function testFalsyCleanupGetterErrors() {
161185
for (const [symbol, create] of [
162186
[Symbol.asyncIterator, toReadable],
@@ -722,6 +746,7 @@ Promise.all([
722746
testErrorAsync(),
723747
testFalsyErrorAsync(),
724748
testFalsyThenableCleanupError(),
749+
testSynchronousIteratorReturn(),
725750
testFalsyCleanupGetterErrors(),
726751
testEmptyAsync(),
727752
testEmptyBatchAsync(),

0 commit comments

Comments
 (0)