stream: keep writer async context for parked transform writes

When a TransformStream write parks on backpressure, the reader that
clears the backpressure delivers the cached write continuation. The
continuation was registered in the reader's async context. Thus
transform() for the parked chunk, and for every later queued chunk,
ran in the reader's context.

In the spec, a parked write is a reaction to
[[backpressureChangePromise]], and the writable controller reacts to
the sink write promise. Both reactions are registered at write time,
so they run in the writer's context.

Store the async context frame when the write parks, and register the
continuation in that frame. This gives the v24.20.0 behavior for
AsyncLocalStorage again. If the reader runs in the same frame, the
frame does not change.

Fixes: #66906
Refs: #65143
Refs: #65625
Signed-off-by: Timer <timer@vercel.com>

Author: Timer

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',
+  ]);
+});