stream: fixup writer to terminate on consumer return/throw

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

@@ -425,6 +425,13 @@ class PushQueue {
   // ===========================================================================
 
   async read() {
+    if (this.#consumerState === 'returned') {
+      return { __proto__: null, done: true, value: undefined };
+    }
+    if (this.#consumerState === 'thrown') {
+      throw this.#error;
+    }
+
     // If there's data in the buffer, return it immediately
     if (this.#slots.length > 0) {
       const result = this.#drain();
@@ -458,16 +465,9 @@ class PushQueue {
   consumerReturn() {
     if (this.#consumerState !== 'active') return;
     this.#consumerState = 'returned';
-    this.#cleanup();
+    const error = new ERR_INVALID_STATE.TypeError('Stream closed by consumer');
+    this.#terminateWriterFromConsumer(error);
     this.#resolvePendingReads();
-    this.#rejectPendingWrites(
-      new ERR_INVALID_STATE.TypeError('Stream closed by consumer'));
-    // If closing, reject the pending end promise
-    if (this.#writerState === 'closing' && this.#pendingEnd) {
-      this.#pendingEnd.reject(
-        new ERR_INVALID_STATE.TypeError('Stream closed by consumer'));
-      this.#pendingEnd = null;
-    }
     // Resolve pending drains with false - no more data will be consumed
     this.#resolvePendingDrains(false);
   }
@@ -476,13 +476,8 @@ class PushQueue {
     if (this.#consumerState !== 'active') return;
     this.#consumerState = 'thrown';
     this.#error = error;
-    this.#cleanup();
+    this.#terminateWriterFromConsumer(error);
     this.#rejectPendingReads(error);
-    this.#rejectPendingWrites(error);
-    if (this.#writerState === 'closing' && this.#pendingEnd) {
-      this.#pendingEnd.reject(error);
-      this.#pendingEnd = null;
-    }
     // Reject pending drains - the consumer errored
     this.#rejectPendingDrains(error);
   }
@@ -516,9 +511,30 @@ class PushQueue {
     return size;
   }
 
+  #terminateWriterFromConsumer(error) {
+    this.#slots.clear();
+    this.#bufferedBytes = 0;
+    if (this.#writerState === 'open' || this.#writerState === 'closing') {
+      this.#writerState = 'errored';
+      this.#error = error;
+    }
+    this.#cleanup();
+    this.#rejectPendingWrites(error);
+    if (this.#pendingEnd) {
+      this.#pendingEnd.reject(error);
+      this.#pendingEnd = null;
+    }
+  }
+
   #resolvePendingReads() {
     while (this.#pendingReads.length > 0) {
-      if (this.#slots.length > 0) {
+      if (this.#consumerState === 'returned') {
+        const pending = this.#pendingReads.shift();
+        pending.resolve({ __proto__: null, done: true, value: undefined });
+      } else if (this.#consumerState === 'thrown') {
+        const pending = this.#pendingReads.shift();
+        pending.reject(this.#error);
+      } else if (this.#slots.length > 0) {
         const pending = this.#pendingReads.shift();
         const result = this.#drain();
         this.#resolvePendingWrites();
@@ -533,9 +549,6 @@ class PushQueue {
       } else if (this.#writerState === 'errored') {
         const pending = this.#pendingReads.shift();
         pending.reject(this.#error);
-      } else if (this.#consumerState === 'returned') {
-        const pending = this.#pendingReads.shift();
-        pending.resolve({ __proto__: null, done: true, value: undefined });
       } else {
         break;
       }

test/parallel/test-stream-iter-push-writer.js

@@ -424,6 +424,20 @@ async function testConsumerReturnResolvesPendingRead() {
   assert.strictEqual(readResult.done, true);
 }
 
+async function testEndRejectsAfterConsumerReturn() {
+  const { writer, readable } = push();
+  writer.writeSync('data');
+  const iter = readable[Symbol.asyncIterator]();
+
+  await iter.return();
+
+  await assert.rejects(
+    writer.end({ signal: AbortSignal.timeout(common.platformTimeout(100)) }),
+    { code: 'ERR_INVALID_STATE' },
+  );
+  assert.strictEqual((await iter.next()).done, true);
+}
+
 // iterator.throw() rejects a pending read with the thrown error
 async function testConsumerThrowRejectsPendingRead() {
   const { readable } = push();
@@ -599,6 +613,7 @@ Promise.all([
   testFailRejectsFutureReadWithFalsyReason(),
   testFailRejectsPendingReadWithFalsyReason(),
   testConsumerReturnResolvesPendingRead(),
+  testEndRejectsAfterConsumerReturn(),
   testConsumerThrowRejectsPendingRead(),
   testEndRejectsPendingWrites(),
   testEndIdempotentWhenClosed(),