stream: keep writer async context for parked transform writes

lib/internal/webstreams/transformstream.js

@@ -81,6 +81,7 @@ const {
 } = require('internal/webstreams/writablestream');
 
 const assert = require('internal/assert');
+const AsyncContextFrame = require('internal/async_context_frame');
 
 const kSkipThrow = Symbol('kSkipThrow');
 
@@ -264,6 +265,7 @@ class TransformStreamState {
   pullPending = false;
   pendingWriteParked = false;
   pendingWriteChunk = undefined;
+  pendingWriteFrame = undefined;
   writeContinuation = undefined;
 }
 ObjectSetPrototypeOf(TransformStreamState.prototype, null);
@@ -505,7 +507,18 @@ function transformStreamSetBackpressure(stream, backpressure) {
         state.readable[kState].controller[kState].pullFulfilled);
     }
   } else if (state.pendingWriteParked) {
-    PromisePrototypeThen(kResolvedPromise, state.writeContinuation);
+    // The spec's reaction to [[backpressureChangePromise]] is registered
+    // when the write parks, so it runs in the writer's async context, not
+    // in the context of the reader that clears the backpressure.
+    const frame = state.pendingWriteFrame;
+    const priorFrame = AsyncContextFrame.current();
+    if (frame === priorFrame) {
+      PromisePrototypeThen(kResolvedPromise, state.writeContinuation);
+    } else {
+      AsyncContextFrame.set(frame);
+      PromisePrototypeThen(kResolvedPromise, state.writeContinuation);
+      AsyncContextFrame.set(priorFrame);
+    }
   }
 }
 
@@ -662,10 +675,12 @@ function transformStreamDefaultSinkWriteAlgorithm(chunk, writableController) {
     assert(!state.pendingWriteParked);
     state.pendingWriteParked = true;
     state.pendingWriteChunk = chunk;
+    state.pendingWriteFrame = AsyncContextFrame.current();
     state.writeContinuation ??= () => {
       const pendingChunk = state.pendingWriteChunk;
       state.pendingWriteParked = false;
       state.pendingWriteChunk = undefined;
+      state.pendingWriteFrame = undefined;
       const writableState = state.writable[kState];
       const writableControllerState = writableState.controller[kState];
       if (writableState.state === 'erroring') {

test/parallel/test-whatwg-transformstream-parked-write-async-context.js

@@ -0,0 +1,74 @@
+'use strict';
+
+require('../common');
+const assert = require('node:assert');
+const { AsyncLocalStorage } = require('node:async_hooks');
+const { test } = require('node:test');
+const { setImmediate } = require('node:timers/promises');
+
+// In the spec, a transform sink write that arrives while the stream has
+// backpressure is a reaction to [[backpressureChangePromise]], and the
+// writable controller reacts to the sink write promise. Both reactions are
+// registered at write time, so transform() runs in the writer's async
+// context, also when the reader that clears the backpressure runs in a
+// different context.
+
+function recordingTransform(als, seen) {
+  return new TransformStream({
+    transform(chunk, controller) {
+      seen.push(`${chunk}:${als.getStore()}`);
+      controller.enqueue(chunk);
+    },
+    flush() {
+      seen.push(`flush:${als.getStore()}`);
+    },
+  });
+}
+
+test('pipeThrough() keeps the writer context after a parked write', async () => {
+  const als = new AsyncLocalStorage();
+  const seen = [];
+  const readable = als.run('writer', () =>
+    new ReadableStream({
+      start(controller) {
+        for (const chunk of 'abcde') controller.enqueue(chunk);
+        controller.close();
+      },
+    }).pipeThrough(recordingTransform(als, seen)));
+
+  // Let the internal pipe write the first chunk, which parks.
+  await setImmediate();
+
+  // Read outside als.run().
+  await Array.fromAsync(readable);
+
+  assert.deepStrictEqual(seen, [
+    'a:writer', 'b:writer', 'c:writer', 'd:writer', 'e:writer', 'flush:writer',
+  ]);
+});
+
+test('writer.write() keeps the writer context after a parked write', async () => {
+  const als = new AsyncLocalStorage();
+  const seen = [];
+  const { readable, writable } = recordingTransform(als, seen);
+
+  // Let start() settle, so that the first write parks on backpressure.
+  await setImmediate();
+  const done = als.run('writer', () => {
+    const writer = writable.getWriter();
+    return Promise.all([
+      writer.write('a'),
+      writer.write('b'),
+      writer.write('c'),
+      writer.close(),
+    ]);
+  });
+
+  // Read outside als.run().
+  await Array.fromAsync(readable);
+  await done;
+
+  assert.deepStrictEqual(seen, [
+    'a:writer', 'b:writer', 'c:writer', 'flush:writer',
+  ]);
+});