stream: address stream/iter review feedback

Signed-off-by: James M Snell <jasnell@gmail.com>
Assisted-by: Opencode
PR-URL: #65652
Reviewed-By: Trivikram Kamat <trivikr.dev@gmail.com>
Reviewed-By: Matteo Collina <matteo.collina@gmail.com>

Author: jasnell

lib/internal/streams/iter/consumers.js

@@ -18,6 +18,7 @@ const {
   Promise,
   PromisePrototypeThen,
   SafePromiseAllReturnVoid,
+  SafeSet,
   Symbol,
   SymbolAsyncIterator,
   TypedArrayPrototypeGetBuffer,
@@ -440,6 +441,7 @@ function merge(...args) {
       // async tick per batch. Each source has at most one pending .next()
       // at a time. Every batch from every source is preserved.
       const ready = [];
+      const pendingPulls = new SafeSet();
       let activeCount = normalized.length;
       let waitResolve = null;
       let onAbort;
@@ -461,6 +463,7 @@ function merge(...args) {
       // Called when a source's .next() settles. Pushes the result into
       // the ready queue and wakes the consumer if it's waiting.
       const onSettled = (iterator, result) => {
+        pendingPulls.delete(iterator);
         if (stopped) return;
         if (result.done) {
           activeCount--;
@@ -478,7 +481,8 @@ function merge(...args) {
         }
       };
 
-      const onRejected = (reason) => {
+      const onRejected = (iterator, reason) => {
+        pendingPulls.delete(iterator);
         if (stopped) return;
         ArrayPrototypePush(ready, {
           __proto__: null,
@@ -496,14 +500,14 @@ function merge(...args) {
       for (let i = 0; i < normalized.length; i++) {
         const iterator = normalized[i][SymbolAsyncIterator]();
         ArrayPrototypePush(iterators, iterator);
+        pendingPulls.add(iterator);
         PromisePrototypeThen(
           iterator.next(),
           (r) => onSettled(iterator, r),
-          onRejected,
+          (reason) => onRejected(iterator, reason),
         );
       }
 
-      let completed = false;
       let primaryError = kNoMergeError;
       try {
         while (activeCount > 0 || ready.length > 0) {
@@ -516,10 +520,11 @@ function merge(...args) {
               throw item.reason;
             }
             yield item.value;
+            pendingPulls.add(item.iterator);
             PromisePrototypeThen(
               item.iterator.next(),
               (r) => onSettled(item.iterator, r),
-              onRejected,
+              (reason) => onRejected(item.iterator, reason),
             );
           }
 
@@ -534,7 +539,6 @@ function merge(...args) {
             });
           }
         }
-        completed = true;
       } catch (err) {
         primaryError = err;
       } finally {
@@ -548,20 +552,20 @@ function merge(...args) {
         await cleanupIterators(
           iterators,
           primaryError,
-          !completed,
+          pendingPulls,
         );
       }
     },
   };
 }
 
-async function cleanupIterators(iterators, primaryError, skipAwaitCleanup) {
+async function cleanupIterators(iterators, primaryError, pendingPulls) {
   let cleanupError = kNoMergeError;
   await SafePromiseAllReturnVoid(iterators, async (iterator) => {
     if (iterator.return) {
       try {
         const result = iterator.return();
-        if (skipAwaitCleanup) {
+        if (pendingPulls.has(iterator)) {
           markPromiseAsHandled(result);
         } else {
           await result;

lib/internal/streams/iter/duplex.js

@@ -12,6 +12,7 @@ const {
 } = primordials;
 
 const {
+  isConsumerReturnError,
   push,
 } = require('internal/streams/iter/push');
 const {
@@ -99,7 +100,11 @@ async function closeDuplexChannel(writer, closeIterator) {
   const returnPromise = closeIterator.return();
 
   if (endPromise !== undefined) {
-    await SafePromiseAllReturnVoid([endPromise, returnPromise]);
+    try {
+      await SafePromiseAllReturnVoid([endPromise, returnPromise]);
+    } catch (error) {
+      if (!isConsumerReturnError(error)) throw error;
+    }
   } else {
     await returnPromise;
   }

lib/internal/streams/iter/push.js

@@ -11,6 +11,7 @@ const {
   PromiseReject,
   PromiseResolve,
   PromiseWithResolvers,
+  SafeWeakSet,
   Symbol,
   SymbolAsyncDispose,
   SymbolAsyncIterator,
@@ -54,6 +55,11 @@ const {
 } = require('internal/streams/iter/ringbuffer');
 
 const kNoFailReason = Symbol('kNoFailReason');
+const consumerReturnErrors = new SafeWeakSet();
+
+function isConsumerReturnError(error) {
+  return consumerReturnErrors.has(error);
+}
 
 function raceEndWithSignal(promise, signal) {
   if (!signal) return promise;
@@ -466,6 +472,7 @@ class PushQueue {
     if (this.#consumerState !== 'active') return;
     this.#consumerState = 'returned';
     const error = new ERR_INVALID_STATE.TypeError('Stream closed by consumer');
+    consumerReturnErrors.add(error);
     this.#terminateWriterFromConsumer(error);
     this.#resolvePendingReads();
     // Resolve pending drains with false - no more data will be consumed
@@ -779,5 +786,6 @@ function push(...args) {
 }
 
 module.exports = {
+  isConsumerReturnError,
   push,
 };

test/parallel/test-stream-iter-consumers-merge.js

@@ -447,6 +447,27 @@ async function testMergeBreakWithCleanupError() {
   );
 }
 
+async function testMergeMultiSourceBreakWithCleanupError() {
+  async function* failingReturnSource() {
+    try {
+      yield [new TextEncoder().encode('data')];
+    } finally {
+      await Promise.resolve();
+      throwInFinally('async cleanup on break');
+    }
+  }
+
+  await assert.rejects(
+    async () => {
+      // eslint-disable-next-line no-unused-vars
+      for await (const _ of merge(failingReturnSource(), from('other'))) {
+        break;
+      }
+    },
+    { message: 'async cleanup on break' },
+  );
+}
+
 Promise.all([
   testMergeTwoSources(),
   testMergeSingleSource(),
@@ -468,4 +489,5 @@ Promise.all([
   testMergeCleanupErrorOnly(),
   testMergePrimaryErrorPrecedesCleanupError(),
   testMergeBreakWithCleanupError(),
+  testMergeMultiSourceBreakWithCleanupError(),
 ]).then(common.mustCall());

test/parallel/test-stream-iter-duplex.js

@@ -77,6 +77,20 @@ async function testChannelClose() {
   assert.strictEqual(batches.length, 0);
 }
 
+async function testConcurrentChannelClose() {
+  const [channelA, channelB] = duplex();
+
+  const results = await Promise.allSettled([
+    channelA.close(),
+    channelB.close(),
+  ]);
+
+  assert.deepStrictEqual(results, [
+    { status: 'fulfilled', value: undefined },
+    { status: 'fulfilled', value: undefined },
+  ]);
+}
+
 async function testWithOptions() {
   const [channelA, channelB] = duplex({
     budget: 16384,
@@ -223,6 +237,7 @@ Promise.all([
   testBidirectional(),
   testMultipleWrites(),
   testChannelClose(),
+  testConcurrentChannelClose(),
   testWithOptions(),
   testPerChannelOptions(),
   testAbortSignal(),