diff --git a/packages/server-utils/src/integrations/tracing-channel/lru-memoizer.ts b/packages/server-utils/src/integrations/tracing-channel/lru-memoizer.ts index 1ca463b3b319..982e5200b4ef 100644 --- a/packages/server-utils/src/integrations/tracing-channel/lru-memoizer.ts +++ b/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,13 @@ 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() {}, + waitForTracingChannelBinding(() => { + bindTracingChannelToSpan( + diagnosticsChannel.tracingChannel(CHANNELS.LRU_MEMOIZER_LOAD), + // We only want the helper's caller-context restore for the callback lru-memoizer fires from a detached `setImmediate`. + () => undefined, + ); }); }, }; diff --git a/packages/server-utils/src/tracing-channel.ts b/packages/server-utils/src/tracing-channel.ts index 6fccf6820ec5..bedc59d2e652 100644 --- a/packages/server-utils/src/tracing-channel.ts +++ b/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 & { + /** + * 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( // 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) => { + // 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) => { + return data._sentryCallerStore as TData; + }); + return { channel, unbind: () => { - // Removes the store + // Removes the stores channel.start.unbindStore(asyncLocalStorage); + channel.asyncStart.unbindStore(asyncLocalStorage); }, }; } diff --git a/packages/server-utils/test/tracing-channel.test.ts b/packages/server-utils/test/tracing-channel.test.ts index fdaeee36476f..a6d26b3f5f3a 100644 --- a/packages/server-utils/test/tracing-channel.test.ts +++ b/packages/server-utils/test/tracing-channel.test.ts @@ -175,6 +175,81 @@ 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(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); + }); + + it('does not leak an unrelated active store into the callback when the caller had none', () => { + installTestAsyncContextStrategy(); + initTestClient(); + + const { channel } = bindTracingChannelToSpan(tracingChannel<{ operation: string }>('test:asyncStart:no-leak'), () => + startInactiveSpan({ name: 'channel-span' }), + ); + + // Caller issues the op with no active context, so the caller store is captured as `undefined`. + const ctx = { operation: 'read' }; + channel.start.runStores(ctx, () => undefined); + + let otherRequestSpanId: string | undefined; + let childParentSpanId: string | undefined; + + // The callback fires later, dispatched from *another* request's active context. + startSpan({ forceTransaction: true, name: 'other-request' }, other => { + otherRequestSpanId = other.spanContext().spanId; + channel.asyncStart.runStores(ctx, () => { + startSpan({ name: 'child-span' }, child => { + childParentSpanId = spanToJSON(child).parent_span_id; + }); + }); + }); + + // The caller had no context, so the callback must restore to none — not adopt the other request's. + expect(childParentSpanId).toBeUndefined(); + expect(childParentSpanId).not.toBe(otherRequestSpanId); + }); + describe('auto lifecycle ending strategy', () => { // Returns a channel whose span we can observe, plus spies for `span.end` and `captureException`. function setup(name: string): {