feat(server-utils): Restore caller context for callback tracing channels

`bindTracingChannelToSpan` only bound the span store on `start`, which covers
the synchronous frame but not a callback the library dispatches from a detached
async context (e.g. a socket data handler or a `setImmediate` drain). There,
native async-context propagation no longer reaches the caller, so work issued
inside the callback lost its parent.

Stash the caller's store at `start` and re-bind it on `asyncStart`, so
callback-style channels run their continuation in the caller's context — the
same way a promise's `.then` does natively. It's inert for promise channels,
which `publish` `asyncStart` rather than `runStores` it.

Migrate the lru-memoizer subscriber onto the helper (`getSpan` returns
`undefined`, so no span is created — it only needs the context rebind),
dropping its hand-rolled callback re-wrapping.

Author: logaretm

packages/server-utils/src/integrations/tracing-channel/lru-memoizer.ts

@@ -1,14 +1,15 @@
 import * as diagnosticsChannel from 'node:diagnostics_channel';
 import type { IntegrationFn } from '@sentry/core';
-import { debug, defineIntegration, getCurrentScope, withScope } from '@sentry/core';
+import { debug, defineIntegration, waitForTracingChannelBinding } from '@sentry/core';
 import { DEBUG_BUILD } from '../../debug-build';
 import { CHANNELS } from '../../orchestrion/channels';
+import { bindTracingChannelToSpan } from '../../tracing-channel';
 
 // Same name as the OTel integration by design — when enabled, the OTel
 // 'LruMemoizer' integration is omitted from the default set.
 const INTEGRATION_NAME = 'LruMemoizer' as const;
 
-interface LruMemoizerChannelContext {
+interface LruMemoizerLoadContext {
   arguments: unknown[];
 }
 
@@ -22,35 +23,19 @@ const _lruMemoizerChannelIntegration = (() => {
       }
 
       DEBUG_BUILD && debug.log(`[orchestrion:lru-memoizer] subscribing to channel "${CHANNELS.LRU_MEMOIZER_LOAD}"`);
-      const lruMemoizerCh = diagnosticsChannel.tracingChannel(CHANNELS.LRU_MEMOIZER_LOAD);
 
-      lruMemoizerCh.subscribe({
-        start(rawCtx) {
-          const ctx = rawCtx as LruMemoizerChannelContext;
-          if (ctx.arguments.length === 0) {
-            return;
-          }
-
-          // Capture the scope while we're still synchronously inside the memoized call.
-          // lru-memoizer queues the callback and fires it later via setImmediate, where the
-          // active scope no longer reflects the caller's context.
-          const scope = getCurrentScope();
-          const cbIdx = ctx.arguments.length - 1;
-          const orchestrionWrappedCb = ctx.arguments[cbIdx];
-
-          if (typeof orchestrionWrappedCb !== 'function') {
-            return;
-          }
-
-          const wrapped = orchestrionWrappedCb as (...a: unknown[]) => unknown;
-          ctx.arguments[cbIdx] = function (this: unknown, ...args: unknown[]): unknown {
-            return withScope(scope, () => wrapped.apply(this, args));
-          };
-        },
-        end() {},
-        asyncStart() {},
-        asyncEnd() {},
-        error() {},
+      // lru-memoizer creates no span: it queues the load callback and fires it later via
+      // `setImmediate`, from a detached context where the caller's scope is no longer active.
+      // Returning `undefined` from `getSpan` opts out of span creation entirely — the helper's
+      // `asyncStart` rebind still restores the caller's context for that callback, which is all this
+      // instrumentation needs (keeping memoized work parented to the caller). `bindTracingChannelToSpan`
+      // uses `bindStore`, which needs the async-context binding `initOpenTelemetry()` registers after
+      // integration `setupOnce`, so defer until it's available (matches the other channel subscribers).
+      waitForTracingChannelBinding(() => {
+        bindTracingChannelToSpan(
+          diagnosticsChannel.tracingChannel<LruMemoizerLoadContext>(CHANNELS.LRU_MEMOIZER_LOAD),
+          () => undefined,
+        );
       });
     },
   };

packages/server-utils/src/tracing-channel.ts

@@ -6,7 +6,15 @@ import { DEBUG_BUILD } from './debug-build';
 import { ERROR_TYPE } from '@sentry/conventions/attributes';
 
 export type TracingChannelPayloadWithSpan<TData extends object> = TData & {
+  /**
+   * The current active span for the traced call.
+   */
   _sentrySpan?: Span;
+
+  /**
+   * The context's active store value, used to restore the context for asyncStart continuations for callback-based tracing.
+   */
+  _sentryCallerStore?: unknown;
 };
 
 /*
@@ -158,21 +166,33 @@ function bindSpanToChannelStore<TData extends object>(
   // 3. Read: inside the op, Sentry's scope machinery calls getScopes() → asyncStorage.getStore() on that same ALS, so getCurrentScope/getIsolationScope/getActiveSpan resolve to the scope carrying our span.
   // 4. Nest: any child span started in the traced op parents to that active span.
   channel.start.bindStore(asyncLocalStorage, (data: TracingChannelPayloadWithSpan<TData>) => {
+    // Stash the caller's store before we swap in the span store, so `asyncStart` can restore it for
+    // callback-style channels (see `_sentryCallerStore`).
+    data._sentryCallerStore = asyncLocalStorage.getStore();
+
     const span = getSpan(data);
     if (!span) {
       // Leave the active context untouched so nested operations keep parenting to the enclosing span.
-      return asyncLocalStorage.getStore() as TData;
+      return data._sentryCallerStore as TData;
     }
     data._sentrySpan = span;
 
     return binding.getStoreWithActiveSpan(span) as TData;
   });
 
+  // Restore the caller's context for the async continuation. Only callback-style channels `runStores`
+  // `asyncStart` (so the callback runs inside this store); promise channels `publish` it, leaving this
+  // inert — their continuation already inherits the caller's context natively.
+  channel.asyncStart.bindStore(asyncLocalStorage, (data: TracingChannelPayloadWithSpan<TData>) => {
+    return (data._sentryCallerStore ?? asyncLocalStorage.getStore()) as TData;
+  });
+
   return {
     channel,
     unbind: () => {
-      // Removes the store
+      // Removes the stores
       channel.start.unbindStore(asyncLocalStorage);
+      channel.asyncStart.unbindStore(asyncLocalStorage);
     },
   };
 }

packages/server-utils/test/tracing-channel.test.ts

@@ -175,6 +175,51 @@ describe('bindTracingChannelToSpan', () => {
     expect(childParentSpanId).toBe(parent.spanContext().spanId);
   });
 
+  it('restores the caller context in a callback dispatched from a detached context (asyncStart rebind)', async () => {
+    installTestAsyncContextStrategy();
+    initTestClient();
+
+    let channelSpanId: string | undefined;
+    const { channel } = bindTracingChannelToSpan(
+      tracingChannel<{ operation: string }>('test:asyncStart:caller-context'),
+      () => {
+        const span = startInactiveSpan({ name: 'channel-span' });
+        channelSpanId = span.spanContext().spanId;
+        return span;
+      },
+    );
+
+    let enclosingSpanId: string | undefined;
+    let childParentSpanId: string | undefined;
+
+    await new Promise<void>(done => {
+      startSpan({ forceTransaction: true, name: 'enclosing-span' }, enclosing => {
+        enclosingSpanId = enclosing.spanContext().spanId;
+        channel.traceCallback(
+          (cb: (err: Error | null, result?: string) => void) => {
+            // Fire the callback after the enclosing scope has exited, so it runs in a detached
+            // async context — the asyncStart rebind is the only thing that can restore the caller's.
+            setTimeout(() => cb(null, 'ok'), 1);
+          },
+          0,
+          { operation: 'read' },
+          undefined,
+          () => {
+            startSpan({ name: 'child-span' }, child => {
+              childParentSpanId = spanToJSON(child).parent_span_id;
+            });
+            done();
+          },
+        );
+      });
+    });
+
+    // A span started inside the callback parents to the caller (the enclosing span), not to the
+    // channel span — matching how a promise's `.then` continuation behaves.
+    expect(childParentSpanId).toBe(enclosingSpanId);
+    expect(childParentSpanId).not.toBe(channelSpanId);
+  });
+
   describe('auto lifecycle ending strategy', () => {
     // Returns a channel whose span we can observe, plus spies for `span.end` and `captureException`.
     function setup(name: string): {