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: #66906

Signed-off-by: Matteo Collina <hello@matteocollina.com>
PR-URL: #66907
Backport-PR-URL: #66915
Fixes: #66906
Refs: #66906
Refs: #66907
Reviewed-By: James M Snell <jasnell@gmail.com>

Author: mcollina

lib/internal/webstreams/transformstream.js

@@ -25,6 +25,9 @@ const {
   DOMException,
 } = internalBinding('messaging');
 
+const AsyncContextFrame = require('internal/async_context_frame');
+const { enabledHooksExist } = require('internal/async_hooks');
+
 const {
   customInspectSymbol: kInspect,
   kEmptyObject,
@@ -267,6 +270,8 @@ class TransformStreamState {
   pullPending = false;
   pendingWriteParked = false;
   pendingWriteChunk = undefined;
+  pendingWriteContextFrame = undefined;
+  pendingWriteResolve = undefined;
   writeContinuation = undefined;
 }
 ObjectSetPrototypeOf(TransformStreamState.prototype, null);
@@ -474,7 +479,13 @@ function transformStreamSetBackpressure(stream, backpressure) {
         state.readable[kState].controller[kState].pullFulfilled);
     }
   } else if (state.pendingWriteParked) {
-    PromisePrototypeThen(kResolvedPromise, state.writeContinuation);
+    const resolve = state.pendingWriteResolve;
+    if (resolve !== undefined) {
+      state.pendingWriteResolve = undefined;
+      resolve();
+    } else {
+      PromisePrototypeThen(kResolvedPromise, state.writeContinuation);
+    }
   }
 }
 
@@ -626,31 +637,48 @@ function transformStreamDefaultSinkWriteAlgorithm(stream, chunk) {
     // 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);
     };
+    if (enabledHooksExist()) {
+      // Async hooks (and the async_hooks-based AsyncLocalStorage) track
+      // context through promise reactions, so register the continuation
+      // now, in the writer's context, and resolve it on unpark.
+      const { promise, resolve } = PromiseWithResolvers();
+      state.pendingWriteResolve = resolve;
+      PromisePrototypeThen(promise, state.writeContinuation);
+    }
     return kParkedAlgorithmResult;
   }
   return transformStreamDefaultControllerPerformTransform(controller, chunk);

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

@@ -0,0 +1,60 @@
+// Flags: --no-async-context-frame
+'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.
+// This variant covers the async_hooks-based AsyncLocalStorage.
+// 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());

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());