stream: more iter streams fixes

doc/api/stream_iter.md

@@ -1374,7 +1374,8 @@ added:
 
 * `input` {AsyncIterable|Iterable|BroadcastChannel}
 * `options` {Object} Same as `broadcast()`.
-* Returns: {Object} `{ writer, broadcast }`
+* Returns: {BroadcastChannel|Object} A `broadcastProtocol` input returns its
+  {BroadcastChannel} directly. Other inputs return `{ writer, broadcast }`.
 
 Create a {BroadcastChannel} from an existing source. The source is consumed
 automatically and pushed to all subscribers.
@@ -1875,7 +1876,7 @@ class MessageBus {
 }
 
 const bus = new MessageBus();
-const { broadcast } = Broadcast.from(bus);
+const broadcast = Broadcast.from(bus);
 const consumer = broadcast.push();
 bus.send('hello');
 bus.close();
@@ -1911,7 +1912,7 @@ class MessageBus {
 }
 
 const bus = new MessageBus();
-const { broadcast } = Broadcast.from(bus);
+const broadcast = Broadcast.from(bus);
 const consumer = broadcast.push();
 bus.send('hello');
 bus.close();

lib/internal/streams/iter/broadcast.js

@@ -10,6 +10,7 @@ const {
   ArrayIsArray,
   ArrayPrototypePush,
   ArrayPrototypeShift,
+  FunctionPrototypeCall,
   PromisePrototypeThen,
   PromiseReject,
   PromiseResolve,
@@ -24,6 +25,10 @@ const {
 } = primordials;
 
 const { lazyDOMException } = require('internal/util');
+const {
+  AbortController,
+  abortSignal,
+} = require('internal/abort_controller');
 
 const {
   codes: {
@@ -56,13 +61,14 @@ const {
   kResolvedPromise,
   convertChunks,
   createBatchEntry,
+  getProtocolMethod,
   getWriterSignal,
   getMinCursor,
-  hasProtocol,
   onSignalAbort,
   parsePullArgs,
   toWriterUint8Array,
   validateBatchEntry,
+  yieldAbortable,
 } = require('internal/streams/iter/utils');
 const {
   converters,
@@ -79,8 +85,10 @@ const kAbort = Symbol('kAbort');
 const kCanWrite = Symbol('kCanWrite');
 const kOnBufferDrained = Symbol('kOnBufferDrained');
 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;
@@ -104,7 +112,7 @@ class BroadcastImpl {
   #buffer = new RingBuffer();
   #bufferStart = 0;
   #consumers = new SafeSet();
-  #waiters = [];  // Consumers with pending resolve (subset of #consumers)
+  #waiters = new SafeSet();  // Consumers with pending resolve
   #ended = false;
   #error;
   #errored = false;
@@ -113,19 +121,26 @@ class BroadcastImpl {
   #writer = null;
   #cachedMinCursor = 0;
   #cachedMinCursorConsumers = 0;
+  #abortHandler;
   /** Cumulative byte size of buffered entries */
   #bufferedBytes = 0;
 
   constructor(options) {
     this.#options = options;
     this[kOnBufferDrained] = null;
     this[kOnEndDrained] = null;
+    this[kOnCancel] = null;
   }
 
   setWriter(writer) {
     this.#writer = writer;
   }
 
+  [kSetFactorySignal](signal) {
+    this.#abortHandler = () => this.cancel(signal.reason);
+    onSignalAbort(signal, this.#abortHandler);
+  }
+
   get backpressurePolicy() {
     return this.#options.backpressure;
   }
@@ -196,6 +211,7 @@ class BroadcastImpl {
 
     function detach() {
       state.detached = true;
+      self.#waiters.delete(state);
       if (state.resolve) {
         state.resolve({ __proto__: null, done: true, value: undefined });
       }
@@ -255,7 +271,7 @@ class BroadcastImpl {
             const { promise, resolve, reject } = PromiseWithResolvers();
             state.resolve = resolve;
             state.reject = reject;
-            ArrayPrototypePush(self.#waiters, state);
+            self.#waiters.add(state);
             return promise;
           },
 
@@ -306,7 +322,12 @@ class BroadcastImpl {
       consumer.detached = true;
     }
     this.#consumers.clear();
+    this.#waiters.clear();
     this.#cachedMinCursorConsumers = 0;
+    this.#cleanupFactorySignal();
+    const onCancel = this[kOnCancel];
+    this[kOnCancel] = null;
+    onCancel?.(reason);
   }
 
   [SymbolDispose]() {
@@ -391,6 +412,7 @@ class BroadcastImpl {
         }
       }
     }
+    this.#waiters.clear();
     this.#notifyEndDrained();
   }
 
@@ -412,7 +434,9 @@ class BroadcastImpl {
       consumer.detached = true;
     }
     this.#consumers.clear();
+    this.#waiters.clear();
     this.#cachedMinCursorConsumers = 0;
+    this.#cleanupFactorySignal();
   }
 
   /**
@@ -432,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);
@@ -444,6 +476,8 @@ class BroadcastImpl {
   }
 
   #tryTrimBuffer() {
+    // Retain buffered data for consumers that attach while none are active.
+    if (this.#consumers.size === 0) return;
     if (this.#cachedMinCursorConsumers === 0) {
       this.#recomputeMinCursor();
     }
@@ -477,12 +511,11 @@ class BroadcastImpl {
 
   #notifyConsumers() {
     const waiters = this.#waiters;
-    if (waiters.length === 0) return;
+    if (waiters.size === 0) return;
     // Swap out the waiters list so consumers that re-wait during
     // resolve don't get processed twice in this cycle.
-    this.#waiters = [];
-    for (let i = 0; i < waiters.length; i++) {
-      const consumer = waiters[i];
+    this.#waiters = new SafeSet();
+    for (const consumer of waiters) {
       if (consumer.resolve) {
         const bufferIndex = consumer.cursor - this.#bufferStart;
         if (bufferIndex < this.#buffer.length) {
@@ -501,11 +534,11 @@ class BroadcastImpl {
           if (consumer.detached && this.#deleteConsumer(consumer)) {
             this.#tryTrimBuffer();
           } else if (this.#promotePending(consumer)) {
-            ArrayPrototypePush(this.#waiters, consumer);
+            this.#waiters.add(consumer);
           }
         } else {
           // Still waiting -- put back
-          ArrayPrototypePush(this.#waiters, consumer);
+          this.#waiters.add(consumer);
         }
       }
     }
@@ -841,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
 // =============================================================================
@@ -878,26 +907,23 @@ function broadcast(options = { __proto__: null }) {
   broadcastImpl.setWriter(writer);
 
   if (signal) {
-    onBroadcastCancel(broadcastImpl, signal);
+    broadcastImpl[kSetFactorySignal](signal);
   }
 
   return { __proto__: null, writer, broadcast: broadcastImpl };
 }
 
-function isBroadcastable(value) {
-  return hasProtocol(value, broadcastProtocol);
-}
-
 const Broadcast = {
   __proto__: null,
   from(input, options) {
-    if (isBroadcastable(input)) {
-      const bc = input[broadcastProtocol](options);
+    const protocol = getProtocolMethod(input, broadcastProtocol);
+    if (protocol !== undefined) {
+      const bc = FunctionPrototypeCall(protocol, input, options);
       if (bc === null || typeof bc !== 'object') {
         throw new ERR_INVALID_RETURN_VALUE(
           'an object', '[Symbol.for(\'Stream.broadcastProtocol\')]', bc);
       }
-      return { __proto__: null, writer: { __proto__: null }, broadcast: bc };
+      return bc;
     }
 
     const source = from(input);
@@ -913,13 +939,23 @@ const Broadcast = {
     });
     const result = broadcast(options);
     const { signal } = options;
+    const controller = new AbortController();
+    if (signal?.aborted) {
+      abortSignal(controller.signal, signal.reason);
+    }
+    const onCancel = (reason) => {
+      if (!controller.signal.aborted) {
+        abortSignal(controller.signal, reason);
+      }
+    };
+    result.broadcast[kOnCancel] = onCancel;
 
     const pump = async () => {
       const w = result.writer;
       try {
         if (isAsyncIterable(source)) {
-          for await (const chunks of source) {
-            signal?.throwIfAborted();
+          for await (const chunks of yieldAbortable(source, controller.signal)) {
+            controller.signal.throwIfAborted();
             if (ArrayIsArray(chunks)) {
               if (!w.writevSync(chunks)) {
                 await w.writev(chunks, signal ? { signal } : undefined);
@@ -930,7 +966,7 @@ const Broadcast = {
           }
         } else if (isSyncIterable(source)) {
           for (const chunks of source) {
-            signal?.throwIfAborted();
+            controller.signal.throwIfAborted();
             if (ArrayIsArray(chunks)) {
               if (!w.writevSync(chunks)) {
                 await w.writev(chunks, signal ? { signal } : undefined);
@@ -944,7 +980,13 @@ const Broadcast = {
           await w.end(signal ? { signal } : undefined);
         }
       } catch (error) {
-        w.fail(error);
+        if (!controller.signal.aborted) {
+          w.fail(error);
+        }
+      } finally {
+        if (result.broadcast[kOnCancel] === onCancel) {
+          result.broadcast[kOnCancel] = null;
+        }
       }
     };
     PromisePrototypeThen(pump(), undefined, () => {});

lib/internal/streams/iter/consumers.js

@@ -15,6 +15,7 @@ const {
   ArrayPrototypePush,
   ArrayPrototypeShift,
   ArrayPrototypeSlice,
+  FunctionPrototypeCall,
   Promise,
   PromisePrototypeThen,
   SafePromiseAllReturnVoid,
@@ -53,6 +54,7 @@ const {
 const {
   concatBytes,
   createBatchEntry,
+  getProtocolMethod,
   validateBatchEntry,
   yieldAbortable,
 } = require('internal/streams/iter/utils');
@@ -391,14 +393,9 @@ function ondrain(drainable) {
     return null;
   }
 
-  if (
-    !(drainableProtocol in drainable) ||
-    typeof drainable[drainableProtocol] !== 'function'
-  ) {
-    return null;
-  }
-
-  return drainable[drainableProtocol]();
+  const protocol = getProtocolMethod(drainable, drainableProtocol);
+  return protocol === undefined ?
+    null : FunctionPrototypeCall(protocol, drainable);
 }
 
 // =============================================================================