fix(diagnostics): preserve AsyncLocalStorage across async lifecycle

maybeTraceMixed previously delegated start / end to traceSync and then
attached its own .then / .finally outside that scope. That worked for
event ordering but broke AsyncLocalStorage propagation: a subscriber
that called channel.start.bindStore(als, ...) could read that store in
its start and end handlers but not in asyncStart or asyncEnd, because
those fired in a .then attached outside the start.runStores scope.

Promise continuations inherit the AsyncLocalStorage context of the
frame that attaches them, so the fix is to attach .then inside the
start.runStores block. This mirrors Node's own TracingChannel.tracePromise
structure exactly, just with an additional sync branch that short-
circuits to [start, end] when fn returns synchronously.

Also updates the FakeChannel helper to match: its runStores now
publishes the context on entry (the behavior Node's Channel.runStores
has), so the fake traceSync / tracePromise implementations match real
Node's event counts without the old end.runStores workaround.

Adds an integration test that binds a store on graphql:execute's start
sub-channel and asserts every lifecycle handler sees it when a resolver
returns a promise.

Author: logaretm

integrationTests/diagnostics/test.js

@@ -3,6 +3,7 @@
 /* eslint-disable n/no-unsupported-features/node-builtins */
 
 import assert from 'node:assert/strict';
+import { AsyncLocalStorage } from 'node:async_hooks';
 import dc from 'node:diagnostics_channel';
 
 import {
@@ -247,12 +248,48 @@ function runNoSubscriberCase() {
   assert.equal(doc.kind, 'Document');
 }
 
+async function runAlsPropagationCase() {
+  // A subscriber that binds a store on the `start` sub-channel should be able
+  // to read it in every lifecycle handler (start, end, asyncStart, asyncEnd).
+  // This is what APMs use to parent child spans to the current operation
+  // without threading state through the ctx object.
+  const als = new AsyncLocalStorage();
+  const channel = dc.tracingChannel('graphql:execute');
+  channel.start.bindStore(als, (ctx) => ({ operationName: ctx.operationName }));
+
+  const seen = {};
+  const handler = {
+    start: () => (seen.start = als.getStore()),
+    end: () => (seen.end = als.getStore()),
+    asyncStart: () => (seen.asyncStart = als.getStore()),
+    asyncEnd: () => (seen.asyncEnd = als.getStore()),
+  };
+  channel.subscribe(handler);
+
+  try {
+    const schema = buildSchema(`type Query { slow: String }`);
+    const document = parse('query Slow { slow }');
+    const rootValue = { slow: () => Promise.resolve('done') };
+
+    await execute({ schema, document, rootValue });
+
+    assert.deepEqual(seen.start, { operationName: 'Slow' });
+    assert.deepEqual(seen.end, { operationName: 'Slow' });
+    assert.deepEqual(seen.asyncStart, { operationName: 'Slow' });
+    assert.deepEqual(seen.asyncEnd, { operationName: 'Slow' });
+  } finally {
+    channel.unsubscribe(handler);
+    channel.start.unbindStore(als);
+  }
+}
+
 async function main() {
   runParseCases();
   runValidateCase();
   runExecuteCase();
   await runSubscribeCase();
   runResolveCase();
+  await runAlsPropagationCase();
   runNoSubscriberCase();
   console.log('diagnostics integration test passed');
 }

src/__testUtils__/fakeDiagnosticsChannel.ts

@@ -29,11 +29,15 @@ export class FakeChannel implements MinimalChannel {
   }
 
   runStores<T, ContextType extends object>(
-    _ctx: ContextType,
+    ctx: ContextType,
     fn: (this: ContextType, ...args: Array<unknown>) => T,
     thisArg?: unknown,
     ...args: Array<unknown>
   ): T {
+    // Node's Channel.runStores publishes the context on the channel before
+    // invoking fn. Mirror that here so traceSync / tracePromise fake exactly
+    // matches real Node's start / end event counts.
+    this.publish(ctx);
     return fn.apply(thisArg as ContextType, args);
   }
 
@@ -81,23 +85,24 @@ export class FakeTracingChannel implements MinimalTracingChannel {
     thisArg?: unknown,
     ...args: Array<unknown>
   ): T {
-    this.start.publish(ctx);
-    let result: T;
-    try {
-      result = this.end.runStores(ctx, fn, thisArg, ...args);
-    } catch (err) {
-      (ctx as { error: unknown }).error = err;
-      this.error.publish(ctx);
+    return this.start.runStores(ctx, () => {
+      let result: T;
+      try {
+        result = fn.apply(thisArg as object, args);
+      } catch (err) {
+        (ctx as { error: unknown }).error = err;
+        this.error.publish(ctx);
+        this.end.publish(ctx);
+        throw err;
+      }
+      // Node's real traceSync sets `ctx.result` before publishing `end`, so
+      // subscribers can inspect `isPromise(ctx.result)` inside their `end`
+      // handler to decide whether the operation is complete or async events
+      // will follow. Match that semantic here.
+      (ctx as { result: unknown }).result = result;
       this.end.publish(ctx);
-      throw err;
-    }
-    // Node's real traceSync sets `ctx.result` before publishing `end`, so
-    // subscribers can inspect `isPromise(ctx.result)` inside their `end`
-    // handler to decide whether the operation is complete or async events
-    // will follow. Match that semantic here.
-    (ctx as { result: unknown }).result = result;
-    this.end.publish(ctx);
-    return result;
+      return result;
+    });
   }
 
   tracePromise<T>(
@@ -106,33 +111,39 @@ export class FakeTracingChannel implements MinimalTracingChannel {
     thisArg?: unknown,
     ...args: Array<unknown>
   ): Promise<T> {
-    this.start.publish(ctx);
-    let promise: Promise<T>;
-    try {
-      promise = this.end.runStores(ctx, fn, thisArg, ...args);
-    } catch (err) {
-      (ctx as { error: unknown }).error = err;
-      this.error.publish(ctx);
+    return this.start.runStores(ctx, () => {
+      let promise: Promise<T>;
+      try {
+        promise = fn.apply(thisArg as object, args);
+      } catch (err) {
+        (ctx as { error: unknown }).error = err;
+        this.error.publish(ctx);
+        this.end.publish(ctx);
+        throw err;
+      }
       this.end.publish(ctx);
-      throw err;
-    }
-    this.end.publish(ctx);
-    this.asyncStart.publish(ctx);
-    return promise
-      .then(
+      return promise.then(
         (result) => {
           (ctx as { result: unknown }).result = result;
-          return result;
+          this.asyncStart.publish(ctx);
+          try {
+            return result;
+          } finally {
+            this.asyncEnd.publish(ctx);
+          }
         },
         (err: unknown) => {
           (ctx as { error: unknown }).error = err;
           this.error.publish(ctx);
-          throw err;
+          this.asyncStart.publish(ctx);
+          try {
+            throw err;
+          } finally {
+            this.asyncEnd.publish(ctx);
+          }
         },
-      )
-      .finally(() => {
-        this.asyncEnd.publish(ctx);
-      });
+      );
+    });
   }
 }
 

src/diagnostics.ts

@@ -185,7 +185,21 @@ export function maybeTracePromise<T>(
 }
 
 /**
- * Publish a mixed sync-or-promise operation through the named graphql tracing channel.
+ * Publish a mixed sync-or-promise operation through the named graphql tracing
+ * channel.
+ *
+ * Mirrors Node's own `TracingChannel.tracePromise` for the async branch while
+ * handling sync returns without the cost of a promise wrap. The entire
+ * lifecycle runs inside `start.runStores`, which is what lets subscribers
+ * that call `channel.start.bindStore(als, ...)` read that store in every
+ * sub-channel handler: promise continuations attached inside a `runStores`
+ * block inherit the AsyncLocalStorage context via async_hooks, so
+ * `asyncStart` and `asyncEnd` fire with the same store active as `start`
+ * and `end`.
+ *
+ * Subscribers can inspect `isPromise(ctx.result)` inside their `end` handler
+ * to know whether `asyncEnd` will follow or the operation is complete. This
+ * matches Node's convention.
  *
  * @internal
  */
@@ -203,29 +217,44 @@ export function maybeTraceMixed<T>(
     result?: unknown;
   };
 
-  // traceSync fires start/end (and error, if fn throws synchronously)
-  const result = channel.traceSync(fn, ctx);
-  if (!isPromise(result)) {
-    return result;
-  }
+  return channel.start.runStores(ctx, () => {
+    let result: T | Promise<T>;
+    try {
+      result = fn();
+    } catch (err) {
+      ctx.error = err;
+      channel.error.publish(ctx);
+      channel.end.publish(ctx);
+      throw err;
+    }
 
-  // Fires off `asyncStart` and `asyncEnd` lifecycle events.
-  channel.asyncStart.publish(ctx);
-  return result
-    .then(
+    if (!isPromise(result)) {
+      ctx.result = result;
+      channel.end.publish(ctx);
+      return result;
+    }
+
+    channel.end.publish(ctx);
+    return result.then(
       (value) => {
         ctx.result = value;
-
-        return value;
+        channel.asyncStart.publish(ctx);
+        try {
+          return value;
+        } finally {
+          channel.asyncEnd.publish(ctx);
+        }
       },
       (err: unknown) => {
         ctx.error = err;
         channel.error.publish(ctx);
-
-        throw err;
+        channel.asyncStart.publish(ctx);
+        try {
+          throw err;
+        } finally {
+          channel.asyncEnd.publish(ctx);
+        }
       },
-    )
-    .finally(() => {
-      channel.asyncEnd.publish(ctx);
-    });
+    );
+  });
 }