stream: drain pending writes before broadcast end

Project: nodejs/node | Language: JavaScript | Concepts: concurrency, error-handling, data-structures | Difficulty: mid
Link: https://github.com/nodejs/node/pull/65334

The problem

When calling end() on a broadcast stream writer, writes that were still pending (due to backpressure) could be dropped or left unprocessed. This meant that data written just before end() might never reach consumers, especially if the stream was backpressured or slow. Additionally, end() could resolve before all consumers had finished reading, leading to incomplete data delivery.

Why it's non-trivial

A naive fix would simply call end() and mark the stream as closed, but this would not account for writes that are still buffered or pending due to backpressure. Handling these edge cases requires careful coordination to ensure:

• All pending writes are fully processed and delivered to consumers before ending.
• The end() promise only resolves after all consumers have reached the end of the stream.
• Correct state transitions and error handling, especially with asynchronous and abortable operations.

The approach

1. Track and wait for pending writes before ending

The writer now tracks its state more granularly ('open', 'closing', 'closed', 'errored'). When end() is called, it transitions to 'closing' and only signals the broadcast end after all pending writes are drained:

+  end(options) {
+    const signal = getWriterSignal(options);
+    if (this.#state === 'errored') return PromiseReject(this.#error);
+    if (this.#state === 'closed') return PromiseResolve(this.#totalBytes);
+    if (signal?.aborted) return PromiseReject(signal.reason);
+
+    const endPromise = this.#getEndPromise();
+    if (this.#state === 'open') {
+      this.#state = 'closing';
+      this.#resolvePendingDrains(false);
+      this.#finishEndIfReady();
+    }
+
+    return raceEndWithSignal(endPromise, signal);
+  }

What this does:

• Checks for error/closed state or aborted signal.
• If still open, transitions to 'closing', resolves any pending drains, and checks if it's ready to finish ending (i.e., no pending writes).
• Returns a promise that resolves only when the stream is truly ended.

2. Only signal end after pending writes are drained

A helper method ensures that kEnd is only called when there are no pending writes left:

+  #finishEndIfReady() {
+    if (this.#state === 'closing' && this.#pendingWrites.length === 0) {
+      this.#broadcast[kEnd]();
+    }
+  }

What this does:

• When in 'closing' state, if there are no pending writes, signals the broadcast to end.

3. Notify when all consumers have finished

A new callback is added to the broadcast implementation to notify the writer when all consumers have finished reading, allowing the writer to fully close and resolve the end() promise:

+    this.#broadcast[kOnEndDrained] = () => this.#endDrained();
...
+  #endDrained() {
+    if (this.#state !== 'closing') return;
+    this.#state = 'closed';
+    this.#pendingEnd?.resolve(this.#totalBytes);
+  }

What this does:

• When all consumers are drained, transitions the writer to 'closed' and resolves the end() promise.

4. Ensure pending writes removed on abort and after draining

When a pending write is removed (e.g., due to abort), the writer checks if it’s now ready to finish ending:

+  [kPendingWriteRemoved]() {
+    this.#finishEndIfReady();
+  }

What this does:

• Ensures that aborting or resolving a pending write can trigger end logic if appropriate.

5. Tests for edge cases

New tests ensure that:

• end() waits for pending writes to drain.
• endSync() returns -1 if ending is not yet complete.
• Aborted writes do not block ending.
• All consumers receive all data before end() resolves.

Example test:

+async function testEndDrainsPendingWrite() {
+  ...
+  await writer.write(chunk1);
+  const pendingWrite = writer.write(chunk2);
+  const endPromise = writer.end();
+  ...
+  await pendingWrite;
+  ...
+  assert.strictEqual(await endPromise, 16385);
+}

From the review thread

No notable review discussion on this PR.

Takeaways

• Always drain all pending operations before signaling completion in concurrent systems.
  Ending a stream should wait for all pending writes to be processed to avoid data loss.

• Promise resolution should reflect true completion, not just initiation.
  A call like end() should only resolve when all consumers have finished, not just when the end is requested.

• Explicit state management avoids subtle concurrency bugs.
  Tracking states like 'open', 'closing', 'closed', and 'errored' makes it easier to reason about transitions and avoid race conditions.