stream: keep async context of parked writes

A TransformStream write that parks on backpressure was resumed from
the read that lifted the backpressure, so transform(), the writes
queued behind it, and flush() ran in the reader's async context
instead of the writer's. Capture the async context frame when the
write parks and restore it while the continuation runs, matching the
context the spec's [[backpressureChangePromise]] reaction had.

Fixes: https://github.com/nodejs/node/issues/66906

Signed-off-by: Matteo Collina <hello@matteocollina.com>
PR-URL: https://github.com/nodejs/node/pull/66907
Fixes: https://github.com/nodejs/node/issues/66906
Reviewed-By: Antoine du Hamel <duhamelantoine1995@gmail.com>
Reviewed-By: James M Snell <jasnell@gmail.com>

lib/internal/webstreams/transformstream.js

@@ -26,6 +26,8 @@ const {
   DOMException,
 } = internalBinding('messaging');
 
+const AsyncContextFrame = require('internal/async_context_frame');
+
 const {
   customInspectSymbol: kInspect,
   kEmptyObject,
@@ -264,6 +266,7 @@ class TransformStreamState {
   pullPending = false;
   pendingWriteParked = false;
   pendingWriteChunk = undefined;
+  pendingWriteContextFrame = undefined;
   writeContinuation = undefined;
 }
 ObjectSetPrototypeOf(TransformStreamState.prototype, null);
@@ -658,30 +661,39 @@ function transformStreamDefaultSinkWriteAlgorithm(chunk, writableController) {
     // them before invoking the write algorithm), replacing the promise
     // record the old code allocated and resolved per parked chunk. The
     // writable dispatches a single write at a time, so one pending slot
-    // suffices.
+    // suffices. The old reaction was registered here, so the continuation
+    // runs in the async context captured at park time rather than in the
+    // context of the read that lifts the backpressure.
     assert(!state.pendingWriteParked);
     state.pendingWriteParked = true;
     state.pendingWriteChunk = chunk;
+    state.pendingWriteContextFrame = AsyncContextFrame.current();
     state.writeContinuation ??= () => {
       const pendingChunk = state.pendingWriteChunk;
+      const prior = AsyncContextFrame.exchange(state.pendingWriteContextFrame);
       state.pendingWriteParked = false;
       state.pendingWriteChunk = undefined;
-      const writableState = state.writable[kState];
-      const writableControllerState = writableState.controller[kState];
-      if (writableState.state === 'erroring') {
-        const error = writableState.storedError;
+      state.pendingWriteContextFrame = undefined;
+      try {
+        const writableState = state.writable[kState];
+        const writableControllerState = writableState.controller[kState];
+        if (writableState.state === 'erroring') {
+          const error = writableState.storedError;
+          PromisePrototypeThen(
+            kResolvedPromise,
+            () => writableControllerState.writeRejected(error));
+          return;
+        }
+        assert(writableState.state === 'writable');
         PromisePrototypeThen(
-          kResolvedPromise,
-          () => writableControllerState.writeRejected(error));
-        return;
+          transformStreamDefaultControllerPerformTransform(
+            controller,
+            pendingChunk),
+          writableControllerState.writeFulfilled,
+          writableControllerState.writeRejected);
+      } finally {
+        AsyncContextFrame.set(prior);
       }
-      assert(writableState.state === 'writable');
-      PromisePrototypeThen(
-        transformStreamDefaultControllerPerformTransform(
-          controller,
-          pendingChunk),
-        writableControllerState.writeFulfilled,
-        writableControllerState.writeRejected);
     };
     return kParkedAlgorithmResult;
   }

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

@@ -0,0 +1,58 @@
+'use strict';
+
+// A write parked on backpressure must run transform() and the writes queued
+// behind it in the async context of the writer, not in the context of the
+// read that lifts the backpressure.
+// Refs: https://github.com/nodejs/node/issues/66906
+
+const common = require('../common');
+const assert = require('assert');
+const { AsyncLocalStorage } = require('async_hooks');
+const { setImmediate } = require('timers/promises');
+
+const als = new AsyncLocalStorage();
+
+function createPipeline(seen) {
+  return new ReadableStream({
+    start(controller) {
+      for (const chunk of 'abc') controller.enqueue(chunk);
+      controller.close();
+    },
+  }).pipeThrough(new TransformStream({
+    transform(chunk, controller) {
+      seen.push(`${chunk}:${als.getStore() ?? 'none'}`);
+      controller.enqueue(chunk);
+    },
+    flush() {
+      seen.push(`flush:${als.getStore() ?? 'none'}`);
+    },
+  }));
+}
+
+async function consume(readable) {
+  // eslint-disable-next-line no-unused-vars
+  for await (const _ of readable);
+}
+
+(async () => {
+  {
+    const seen = [];
+    const readable = als.run('writer', () => createPipeline(seen));
+    // Let the pipe park its first write on backpressure.
+    await setImmediate();
+    await consume(readable);
+    assert.deepStrictEqual(seen, [
+      'a:writer', 'b:writer', 'c:writer', 'flush:writer',
+    ]);
+  }
+
+  {
+    const seen = [];
+    const readable = createPipeline(seen);
+    await setImmediate();
+    await als.run('reader', () => consume(readable));
+    assert.deepStrictEqual(seen, [
+      'a:none', 'b:none', 'c:none', 'flush:none',
+    ]);
+  }
+})().then(common.mustCall());