stream: make broadcast waiters use SafeSet

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

@@ -110,7 +110,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;
@@ -203,6 +203,7 @@ class BroadcastImpl {
 
     function detach() {
       state.detached = true;
+      self.#waiters.delete(state);
       if (state.resolve) {
         state.resolve({ __proto__: null, done: true, value: undefined });
       }
@@ -262,7 +263,7 @@ class BroadcastImpl {
             const { promise, resolve, reject } = PromiseWithResolvers();
             state.resolve = resolve;
             state.reject = reject;
-            ArrayPrototypePush(self.#waiters, state);
+            self.#waiters.add(state);
             return promise;
           },
 
@@ -313,6 +314,7 @@ class BroadcastImpl {
       consumer.detached = true;
     }
     this.#consumers.clear();
+    this.#waiters.clear();
     this.#cachedMinCursorConsumers = 0;
     const onCancel = this[kOnCancel];
     this[kOnCancel] = null;
@@ -401,6 +403,7 @@ class BroadcastImpl {
         }
       }
     }
+    this.#waiters.clear();
     this.#notifyEndDrained();
   }
 
@@ -422,6 +425,7 @@ class BroadcastImpl {
       consumer.detached = true;
     }
     this.#consumers.clear();
+    this.#waiters.clear();
     this.#cachedMinCursorConsumers = 0;
   }
 
@@ -489,12 +493,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) {
@@ -513,11 +516,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);
         }
       }
     }

test/parallel/test-stream-iter-broadcast-waiter-cleanup.js

@@ -0,0 +1,34 @@
+// Flags: --experimental-stream-iter --expose-gc
+'use strict';
+
+const common = require('../common');
+const assert = require('assert');
+const { broadcast } = require('stream/iter');
+
+async function detachConsumers(shared, count) {
+  for (let i = 0; i < count; i++) {
+    const iterator = shared.push()[Symbol.asyncIterator]();
+    const pending = iterator.next();
+    await iterator.return();
+    await pending;
+  }
+}
+
+async function testDetachedWaitersAreReleased() {
+  const { broadcast: shared } = broadcast();
+
+  await detachConsumers(shared, 100);
+  global.gc();
+  const before = process.memoryUsage().heapUsed;
+
+  await detachConsumers(shared, 50_000);
+  global.gc();
+  const retained = process.memoryUsage().heapUsed - before;
+
+  assert.ok(retained < 8 * 1024 * 1024,
+            `Detached Broadcast waiters retained ${retained} bytes`);
+  assert.strictEqual(shared.consumerCount, 0);
+  shared.cancel();
+}
+
+testDetachedWaitersAreReleased().then(common.mustCall());