stream: have broadcast clean factory listeners on cancel

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>

Author: jasnell

lib/internal/streams/iter/broadcast.js

@@ -88,6 +88,7 @@ const kOnEndDrained = Symbol('kOnEndDrained');
 const kOnCancel = Symbol('kOnCancel');
 const kPendingWriteRemoved = Symbol('kPendingWriteRemoved');
 const kNoBroadcastError = Symbol('kNoBroadcastError');
+const kSetFactorySignal = Symbol('kSetFactorySignal');
 
 function raceEndWithSignal(promise, signal) {
   if (!signal) return promise;
@@ -120,6 +121,7 @@ class BroadcastImpl {
   #writer = null;
   #cachedMinCursor = 0;
   #cachedMinCursorConsumers = 0;
+  #abortHandler;
   /** Cumulative byte size of buffered entries */
   #bufferedBytes = 0;
 
@@ -134,6 +136,11 @@ class BroadcastImpl {
     this.#writer = writer;
   }
 
+  [kSetFactorySignal](signal) {
+    this.#abortHandler = () => this.cancel(signal.reason);
+    onSignalAbort(signal, this.#abortHandler);
+  }
+
   get backpressurePolicy() {
     return this.#options.backpressure;
   }
@@ -317,6 +324,7 @@ class BroadcastImpl {
     this.#consumers.clear();
     this.#waiters.clear();
     this.#cachedMinCursorConsumers = 0;
+    this.#cleanupFactorySignal();
     const onCancel = this[kOnCancel];
     this[kOnCancel] = null;
     onCancel?.(reason);
@@ -428,6 +436,7 @@ class BroadcastImpl {
     this.#consumers.clear();
     this.#waiters.clear();
     this.#cachedMinCursorConsumers = 0;
+    this.#cleanupFactorySignal();
   }
 
   /**
@@ -447,10 +456,18 @@ class BroadcastImpl {
 
   #notifyEndDrained() {
     if (this.#ended && this.#consumers.size === 0) {
+      this.#cleanupFactorySignal();
       this[kOnEndDrained]?.();
     }
   }
 
+  #cleanupFactorySignal() {
+    if (this.#abortHandler !== undefined) {
+      this.#options.signal.removeEventListener('abort', this.#abortHandler);
+      this.#abortHandler = undefined;
+    }
+  }
+
   #recomputeMinCursor() {
     const { minCursor, minCursorConsumers } = getMinCursor(
       this.#consumers, this.#bufferStart + this.#buffer.length);
@@ -857,10 +874,6 @@ function wireBroadcastWriteSignal(entry, signal, resolve, reject, self) {
   signal.addEventListener('abort', onAbort, { __proto__: null, once: true });
 }
 
-function onBroadcastCancel(broadcastImpl, signal) {
-  onSignalAbort(signal, () => broadcastImpl.cancel(signal.reason));
-}
-
 // =============================================================================
 // Public API
 // =============================================================================
@@ -894,7 +907,7 @@ function broadcast(options = { __proto__: null }) {
   broadcastImpl.setWriter(writer);
 
   if (signal) {
-    onBroadcastCancel(broadcastImpl, signal);
+    broadcastImpl[kSetFactorySignal](signal);
   }
 
   return { __proto__: null, writer, broadcast: broadcastImpl };

lib/internal/streams/iter/duplex.js

@@ -45,8 +45,10 @@ function duplex(options = { __proto__: null }) {
     backpressure: b?.backpressure ?? backpressure,
   });
 
-  const channelA = createDuplexChannel(aWriter, aReadable);
-  const channelB = createDuplexChannel(bWriter, bReadable);
+  let cleanupSignal;
+  const onClose = () => cleanupSignal?.();
+  const channelA = createDuplexChannel(aWriter, aReadable, onClose);
+  const channelB = createDuplexChannel(bWriter, bReadable, onClose);
 
   // Signal handler: fail both writers with the abort reason so consumers
   // see the error. This is an error-path shutdown, not a clean close.
@@ -55,6 +57,11 @@ function duplex(options = { __proto__: null }) {
       const reason = signal.reason;
       aWriter.fail(reason);
       bWriter.fail(reason);
+      cleanupSignal?.();
+    };
+    cleanupSignal = () => {
+      signal.removeEventListener('abort', abortBoth);
+      cleanupSignal = undefined;
     };
     if (signal.aborted) {
       abortBoth();
@@ -67,7 +74,7 @@ function duplex(options = { __proto__: null }) {
   return [channelA, channelB];
 }
 
-function createDuplexChannel(writer, readable) {
+function createDuplexChannel(writer, readable, onClose) {
   // A push readable has one shared consumer state. Keeping an iterator from
   // creation lets close() terminate that state even if no caller has iterated.
   const closeIterator = readable[SymbolAsyncIterator]();
@@ -78,7 +85,7 @@ function createDuplexChannel(writer, readable) {
     get writer() { return writer; },
     get readable() { return readable; },
     close() {
-      closePromise ??= closeDuplexChannel(writer, closeIterator);
+      closePromise ??= closeDuplexChannel(writer, closeIterator, onClose);
       return closePromise;
     },
     [SymbolAsyncDispose]() {
@@ -87,19 +94,23 @@ function createDuplexChannel(writer, readable) {
   };
 }
 
-async function closeDuplexChannel(writer, closeIterator) {
-  const result = writer.endSync();
-  const endPromise = result < 0 ? writer.end() : undefined;
-  const returnPromise = closeIterator.return();
+async function closeDuplexChannel(writer, closeIterator, onClose) {
+  try {
+    const result = writer.endSync();
+    const endPromise = result < 0 ? writer.end() : undefined;
+    const returnPromise = closeIterator.return();
 
-  if (endPromise !== undefined) {
-    try {
-      await SafePromiseAllReturnVoid([endPromise, returnPromise]);
-    } catch (error) {
-      if (!isConsumerReturnError(error)) throw error;
+    if (endPromise !== undefined) {
+      try {
+        await SafePromiseAllReturnVoid([endPromise, returnPromise]);
+      } catch (error) {
+        if (!isConsumerReturnError(error)) throw error;
+      }
+    } else {
+      await returnPromise;
     }
-  } else {
-    await returnPromise;
+  } finally {
+    onClose();
   }
 }
 

lib/internal/streams/iter/share.js

@@ -71,6 +71,7 @@ const {
 
 const kNoShareError = Symbol('kNoShareError');
 const kShareCancelled = Symbol('kShareCancelled');
+const kSetFactorySignal = Symbol('kSetFactorySignal');
 
 class ShareImpl {
   #source;
@@ -89,6 +90,7 @@ class ShareImpl {
   #cancelError = kNoShareError;
   #cachedMinCursor = 0;
   #cachedMinCursorConsumers = 0;
+  #abortHandler;
   /** Cumulative byte size of buffered entries */
   #bufferedBytes = 0;
 
@@ -104,6 +106,11 @@ class ShareImpl {
     return this.#consumers.size;
   }
 
+  [kSetFactorySignal](signal) {
+    this.#abortHandler = () => this.cancel(signal.reason);
+    onSignalAbort(signal, this.#abortHandler);
+  }
+
   pull(...args) {
     const parsed = parsePullArgs(args);
     const { transforms } = parsed;
@@ -303,6 +310,7 @@ class ShareImpl {
     this.#consumers.clear();
     this.#buffer.clear();
     this.#bufferedBytes = 0;
+    this.#cleanupFactorySignal();
 
     for (let i = 0; i < this.#pullWaiters.length; i++) {
       this.#pullWaiters[i]();
@@ -424,6 +432,7 @@ class ShareImpl {
         this.#sourceError = error;
         this.#sourceExhausted = true;
       } finally {
+        if (this.#sourceExhausted) this.#cleanupFactorySignal();
         this.#pulling = false;
         for (let i = 0; i < this.#pullWaiters.length; i++) {
           this.#pullWaiters[i]();
@@ -470,6 +479,13 @@ class ShareImpl {
     this.#cachedMinCursorConsumers = minCursorConsumers;
   }
 
+  #cleanupFactorySignal() {
+    if (this.#abortHandler !== undefined) {
+      this.#options.signal.removeEventListener('abort', this.#abortHandler);
+      this.#abortHandler = undefined;
+    }
+  }
+
   #deleteConsumerFromMin(consumer) {
     if (consumer.cursor === this.#cachedMinCursor) {
       this.#cachedMinCursorConsumers--;
@@ -768,10 +784,6 @@ class SyncShareImpl {
   }
 }
 
-function onShareCancel(shareImpl, signal) {
-  onSignalAbort(signal, () => shareImpl.cancel(signal.reason));
-}
-
 // =============================================================================
 // Public API
 // =============================================================================
@@ -800,7 +812,7 @@ function share(source, options = { __proto__: null }) {
   const shareImpl = new ShareImpl(normalized, opts);
 
   if (signal) {
-    onShareCancel(shareImpl, signal);
+    shareImpl[kSetFactorySignal](signal);
   }
 
   return shareImpl;

test/parallel/test-stream-iter-factory-signal.js

@@ -0,0 +1,84 @@
+// Flags: --experimental-stream-iter
+'use strict';
+
+const common = require('../common');
+const assert = require('assert');
+const { getEventListeners } = require('events');
+const {
+  broadcast,
+  duplex,
+  push,
+  share,
+  text,
+} = require('stream/iter');
+
+function abortListenerCount(signal) {
+  return getEventListeners(signal, 'abort').length;
+}
+
+async function testPushSignalLifetime() {
+  const controller = new AbortController();
+  const { writer, readable } = push({ signal: controller.signal });
+  const iterator = readable[Symbol.asyncIterator]();
+
+  assert.strictEqual(abortListenerCount(controller.signal), 1);
+  assert.strictEqual(writer.writeSync('x'), true);
+  const ending = writer.end();
+  assert.strictEqual(abortListenerCount(controller.signal), 1);
+
+  assert.strictEqual((await iterator.next()).done, false);
+  assert.strictEqual((await iterator.next()).done, true);
+  assert.strictEqual(await ending, 1);
+  assert.strictEqual(abortListenerCount(controller.signal), 0);
+}
+
+async function testBroadcastSignalLifetime() {
+  const controller = new AbortController();
+  const { writer, broadcast: shared } = broadcast({
+    signal: controller.signal,
+  });
+  const iterator = shared.push()[Symbol.asyncIterator]();
+
+  assert.strictEqual(abortListenerCount(controller.signal), 1);
+  assert.strictEqual(writer.writeSync('x'), true);
+  const ending = writer.end();
+  assert.strictEqual(abortListenerCount(controller.signal), 1);
+
+  assert.strictEqual((await iterator.next()).done, false);
+  assert.strictEqual((await iterator.next()).done, true);
+  assert.strictEqual(await ending, 1);
+  assert.strictEqual(abortListenerCount(controller.signal), 0);
+}
+
+async function testShareSignalLifetime() {
+  const controller = new AbortController();
+  const shared = share('x', { signal: controller.signal });
+  const iterator = shared.pull()[Symbol.asyncIterator]();
+
+  assert.strictEqual(abortListenerCount(controller.signal), 1);
+  assert.strictEqual((await iterator.next()).done, false);
+  assert.strictEqual(abortListenerCount(controller.signal), 1);
+  assert.strictEqual((await iterator.next()).done, true);
+  assert.strictEqual(abortListenerCount(controller.signal), 0);
+}
+
+async function testDuplexSignalLifetime() {
+  const controller = new AbortController();
+  const [channelA, channelB] = duplex({ signal: controller.signal });
+
+  assert.strictEqual(abortListenerCount(controller.signal), 1);
+  await channelA.writer.write('x');
+  const closing = channelA.close();
+  assert.strictEqual(abortListenerCount(controller.signal), 1);
+
+  assert.strictEqual(await text(channelB.readable), 'x');
+  await closing;
+  assert.strictEqual(abortListenerCount(controller.signal), 0);
+}
+
+Promise.all([
+  testPushSignalLifetime(),
+  testBroadcastSignalLifetime(),
+  testShareSignalLifetime(),
+  testDuplexSignalLifetime(),
+]).then(common.mustCall());