stream: ensure iterator cleanup on done, reject, etc

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/broadcast.js

@@ -52,7 +52,7 @@ const {
 } = require('internal/streams/iter/from');
 
 const {
-  pull: pullWithTransforms,
+  pullWithConsumerCleanup,
 } = require('internal/streams/iter/pull');
 
 const {
@@ -157,12 +157,7 @@ class BroadcastImpl {
     // When no transforms, return rawConsumer directly (controller elided
     // per PULL-02 optimization -- no transforms means no signal recipient).
     if (transforms.length > 0 || signal) {
-      const pullArgs = [...transforms];
-      if (signal) {
-        ArrayPrototypePush(pullArgs,
-                           { __proto__: null, signal });
-      }
-      return pullWithTransforms(rawConsumer, ...pullArgs);
+      return pullWithConsumerCleanup(rawConsumer, transforms, signal);
     }
     return rawConsumer;
   }

lib/internal/streams/iter/pull.js

@@ -12,6 +12,7 @@ const {
   ArrayIsArray,
   ArrayPrototypePush,
   ArrayPrototypeSlice,
+  FunctionPrototypeCall,
   PromisePrototypeThen,
   PromiseResolve,
   SymbolAsyncIterator,
@@ -896,6 +897,76 @@ function pull(source, ...args) {
   };
 }
 
+// Keep ownership of a bonded consumer outside the transform pipeline so it can
+// be detached even when the pipeline never starts or terminates early.
+function pullWithConsumerCleanup(source, transforms, signal) {
+  const sourceIterator = source[SymbolAsyncIterator]();
+  const pipelineSource = {
+    __proto__: null,
+    [SymbolAsyncIterator]() {
+      return sourceIterator;
+    },
+  };
+  const pipeline = signal === undefined ?
+    pull(pipelineSource, ...transforms) :
+    pull(pipelineSource, ...transforms, { __proto__: null, signal });
+  let sourceClosed = false;
+  let abortHandler;
+
+  function closeSource(method, value) {
+    if (sourceClosed) return;
+    sourceClosed = true;
+    if (abortHandler !== undefined) {
+      signal.removeEventListener('abort', abortHandler);
+    }
+    const close = sourceIterator[method] ?? sourceIterator.return;
+    if (typeof close === 'function') {
+      const result = FunctionPrototypeCall(close, sourceIterator, value);
+      PromisePrototypeThen(PromiseResolve(result), undefined, () => {});
+    }
+  }
+
+  if (signal !== undefined) {
+    abortHandler = () => closeSource('throw', signal.reason);
+    signal.addEventListener('abort', abortHandler,
+                            { __proto__: null, once: true });
+    if (signal.aborted) abortHandler();
+  }
+
+  return {
+    __proto__: null,
+    [SymbolAsyncIterator]() {
+      const iterator = pipeline[SymbolAsyncIterator]();
+      return {
+        __proto__: null,
+        next(value) {
+          return PromisePrototypeThen(
+            iterator.next(value),
+            (result) => {
+              if (result.done) closeSource('return');
+              return result;
+            },
+            (error) => {
+              closeSource('throw', error);
+              throw error;
+            });
+        },
+        return(value) {
+          closeSource('return', value);
+          return iterator.return(value);
+        },
+        throw(error) {
+          closeSource('throw', error);
+          return iterator.throw(error);
+        },
+        [SymbolAsyncIterator]() {
+          return this;
+        },
+      };
+    },
+  };
+}
+
 // =============================================================================
 // Public API: pipeTo() and pipeToSync()
 // =============================================================================
@@ -1112,4 +1183,5 @@ module.exports = {
   pipeToSync,
   pull,
   pullSync,
+  pullWithConsumerCleanup,
 };

lib/internal/streams/iter/push.js

@@ -46,7 +46,7 @@ const {
 } = require('internal/streams/iter/utils');
 
 const {
-  pull: pullWithTransforms,
+  pullWithConsumerCleanup,
 } = require('internal/streams/iter/pull');
 
 const {
@@ -756,12 +756,8 @@ function push(...args) {
   // Apply transforms lazily if provided
   let readable;
   if (transforms.length > 0) {
-    if (options.signal) {
-      readable = pullWithTransforms(
-        rawReadable, ...transforms, { __proto__: null, signal: options.signal });
-    } else {
-      readable = pullWithTransforms(rawReadable, ...transforms);
-    }
+    readable = pullWithConsumerCleanup(
+      rawReadable, transforms, options.signal);
   } else {
     readable = rawReadable;
   }

lib/internal/streams/iter/share.js

@@ -30,8 +30,8 @@ const {
 } = require('internal/streams/iter/from');
 
 const {
-  pull: pullWithTransforms,
   pullSync: pullSyncWithTransforms,
+  pullWithConsumerCleanup,
 } = require('internal/streams/iter/pull');
 
 const {
@@ -111,13 +111,7 @@ class ShareImpl {
     const rawConsumer = this.#createRawConsumer();
 
     if (transforms.length > 0 || signal) {
-      if (signal) {
-        return pullWithTransforms(
-          rawConsumer,
-          ...transforms,
-          { __proto__: null, signal });
-      }
-      return pullWithTransforms(rawConsumer, ...transforms);
+      return pullWithConsumerCleanup(rawConsumer, transforms, signal);
     }
     return rawConsumer;
   }

test/parallel/test-stream-iter-pull-async.js

@@ -3,7 +3,15 @@
 
 const common = require('../common');
 const assert = require('assert');
-const { pull, from, text, tap } = require('stream/iter');
+const {
+  broadcast,
+  from,
+  pull,
+  push,
+  share,
+  tap,
+  text,
+} = require('stream/iter');
 
 async function testPullIdentity() {
   const data = await text(pull(from('hello-async')));
@@ -199,6 +207,32 @@ async function testPullReturnWhileSourceNextPending() {
   await next;
 }
 
+async function testTransformedConsumerReturnBeforeNext() {
+  const identity = (chunks) => chunks;
+  const pushed = push(identity);
+  const { broadcast: bc } = broadcast();
+  const broadcastConsumer = bc.push(identity);
+  const shared = share(from('shared'));
+  const sharedConsumer = shared.pull(identity);
+
+  assert.strictEqual(bc.consumerCount, 1);
+  assert.strictEqual(shared.consumerCount, 1);
+
+  const cases = [
+    [pushed.readable, common.mustCall(
+      () => assert.strictEqual(pushed.writer.canWrite, null))],
+    [broadcastConsumer, common.mustCall(
+      () => assert.strictEqual(bc.consumerCount, 0))],
+    [sharedConsumer, common.mustCall(
+      () => assert.strictEqual(shared.consumerCount, 0))],
+  ];
+
+  for (const [readable, verify] of cases) {
+    await readable[Symbol.asyncIterator]().return();
+    verify();
+  }
+}
+
 async function testPullSignalAbortWithTransformWhileSourceNextPending() {
   const source = {
     [Symbol.asyncIterator]() {
@@ -447,6 +481,7 @@ async function testTransformOptionsNotShared() {
     testPullSignalAbortMidIteration(),
     testPullSignalAbortWhileSourceNextPending(),
     testPullReturnWhileSourceNextPending(),
+    testTransformedConsumerReturnBeforeNext(),
     testPullSignalAbortWithTransformWhileSourceNextPending(),
     testPullConsumerBreakCleanup(),
     testPullTransformReturnsPromise(),