feat(server-utils): Add tracingChannel-to-span binding

packages/cloudflare/src/async.ts

@@ -2,7 +2,12 @@
 // Note: Because we are using node:async_hooks, we need to set `node_compat` in the wrangler.toml
 import { AsyncLocalStorage } from 'node:async_hooks';
 import type { Scope } from '@sentry/core';
-import { getDefaultCurrentScope, getDefaultIsolationScope, setAsyncContextStrategy } from '@sentry/core';
+import {
+  _INTERNAL_setSpanForScope,
+  getDefaultCurrentScope,
+  getDefaultIsolationScope,
+  setAsyncContextStrategy,
+} from '@sentry/core';
 
 /**
  * Sets the async context strategy to use AsyncLocalStorage.
@@ -80,5 +85,14 @@ export function setAsyncLocalStorageAsyncContextStrategy(): void {
     withSetIsolationScope,
     getCurrentScope: () => getScopes().scope,
     getIsolationScope: () => getScopes().isolationScope,
+    getTracingChannelBinding: () => ({
+      asyncLocalStorage: asyncStorage,
+      getStoreWithActiveSpan: span => {
+        const scope = getScopes().scope.clone();
+        const isolationScope = getScopes().isolationScope;
+        _INTERNAL_setSpanForScope(scope, span);
+        return { scope, isolationScope };
+      },
+    }),
   });
 }

packages/core/src/asyncContext/index.ts

@@ -1,7 +1,7 @@
 import type { Carrier } from './../carrier';
 import { getMainCarrier, getSentryCarrier } from './../carrier';
 import { getStackAsyncContextStrategy } from './stackStrategy';
-import type { AsyncContextStrategy } from './types';
+import type { AsyncContextStrategy, TracingChannelBinding } from './types';
 
 /**
  * @private Private API with no semver guarantees!
@@ -29,3 +29,10 @@ export function getAsyncContextStrategy(carrier: Carrier): AsyncContextStrategy
   // Otherwise, use the default one (stack)
   return getStackAsyncContextStrategy();
 }
+
+/**
+ * Get the runtime binding needed to connect tracing channels to async context.
+ */
+export function getTracingChannelBinding(): TracingChannelBinding | undefined {
+  return getAsyncContextStrategy(getMainCarrier()).getTracingChannelBinding?.();
+}

packages/core/src/asyncContext/types.ts

@@ -1,4 +1,5 @@
 import type { Scope } from '../scope';
+import type { Span } from '../types/span';
 import type { getTraceData } from '../utils/traceData';
 import type {
   continueTrace,
@@ -11,6 +12,11 @@ import type {
 } from './../tracing/trace';
 import type { getActiveSpan } from './../utils/spanUtils';
 
+export interface TracingChannelBinding {
+  asyncLocalStorage: unknown;
+  getStoreWithActiveSpan: (span: Span) => unknown;
+}
+
 /**
  * @private Private API with no semver guarantees!
  *
@@ -80,4 +86,7 @@ export interface AsyncContextStrategy {
 
   /** Start a new trace, ensuring all spans in the callback share the same traceId. */
   startNewTrace?: typeof startNewTrace;
+
+  /** Get the runtime store required to bind tracing channels to an active span. */
+  getTracingChannelBinding?: () => TracingChannelBinding | undefined;
 }

packages/core/src/shared-exports.ts

@@ -4,7 +4,7 @@
 /* eslint-disable max-lines */
 
 export type { ClientClass as SentryCoreCurrentScopes } from './sdk';
-export type { AsyncContextStrategy } from './asyncContext/types';
+export type { AsyncContextStrategy, TracingChannelBinding } from './asyncContext/types';
 export type { Carrier } from './carrier';
 export type { OfflineStore, OfflineTransportOptions } from './transports/offline';
 export type { IntegrationIndex } from './integration';
@@ -47,7 +47,10 @@ export {
   hasExternalPropagationContext,
 } from './currentScopes';
 export { getDefaultCurrentScope, getDefaultIsolationScope } from './defaultScopes';
-export { setAsyncContextStrategy } from './asyncContext';
+export {
+  setAsyncContextStrategy,
+  getTracingChannelBinding as _INTERNAL_getTracingChannelBinding,
+} from './asyncContext';
 export { getGlobalSingleton, getMainCarrier } from './carrier';
 export { makeSession, closeSession, updateSession } from './session';
 export { Scope } from './scope';

packages/deno/src/async.ts

@@ -1,7 +1,12 @@
 // Need to use node: prefix for deno compatibility
 import { AsyncLocalStorage } from 'node:async_hooks';
 import type { Scope } from '@sentry/core';
-import { getDefaultCurrentScope, getDefaultIsolationScope, setAsyncContextStrategy } from '@sentry/core';
+import {
+  _INTERNAL_setSpanForScope,
+  getDefaultCurrentScope,
+  getDefaultIsolationScope,
+  setAsyncContextStrategy,
+} from '@sentry/core';
 
 let installed = false;
 
@@ -88,5 +93,14 @@ export function setAsyncLocalStorageAsyncContextStrategy(): void {
     withSetIsolationScope,
     getCurrentScope: () => getScopes().scope,
     getIsolationScope: () => getScopes().isolationScope,
+    getTracingChannelBinding: () => ({
+      asyncLocalStorage: asyncStorage,
+      getStoreWithActiveSpan: span => {
+        const scope = getScopes().scope.clone();
+        const isolationScope = getScopes().isolationScope;
+        _INTERNAL_setSpanForScope(scope, span);
+        return { scope, isolationScope };
+      },
+    }),
   });
 }

packages/node-core/src/light/asyncLocalStorageStrategy.ts

@@ -1,6 +1,7 @@
 import { AsyncLocalStorage } from 'node:async_hooks';
 import type { Scope } from '@sentry/core';
 import {
+  _INTERNAL_setSpanForScope,
   getDefaultCurrentScope,
   getDefaultIsolationScope,
   setAsyncContextStrategy,
@@ -80,5 +81,14 @@ export function setAsyncLocalStorageAsyncContextStrategy(): void {
     withSetIsolationScope,
     getCurrentScope: () => getScopes().scope,
     getIsolationScope: () => getScopes().isolationScope,
+    getTracingChannelBinding: () => ({
+      asyncLocalStorage: asyncStorage,
+      getStoreWithActiveSpan: span => {
+        const scope = getScopes().scope.clone();
+        const isolationScope = getScopes().isolationScope;
+        _INTERNAL_setSpanForScope(scope, span);
+        return { scope, isolationScope };
+      },
+    }),
   });
 }

packages/opentelemetry/src/asyncContextStrategy.ts

@@ -1,5 +1,5 @@
 import * as api from '@opentelemetry/api';
-import type { Scope, withActiveSpan as defaultWithActiveSpan } from '@sentry/core';
+import type { Scope, Span, withActiveSpan as defaultWithActiveSpan } from '@sentry/core';
 import { getDefaultCurrentScope, getDefaultIsolationScope, setAsyncContextStrategy } from '@sentry/core';
 import {
   SENTRY_FORK_ISOLATION_SCOPE_CONTEXT_KEY,
@@ -13,6 +13,14 @@ import { getActiveSpan } from './utils/getActiveSpan';
 import { getTraceData } from './utils/getTraceData';
 import { suppressTracing } from './utils/suppressTracing';
 
+interface ContextApi {
+  _getContextManager(): {
+    getAsyncLocalStorageLookup(): {
+      asyncLocalStorage: unknown;
+    };
+  };
+}
+
 /**
  * Sets the async context strategy to use follow the OTEL context under the hood.
  * We handle forking a hub inside of our custom OTEL Context Manager (./otelContextManager.ts)
@@ -108,5 +116,18 @@ export function setOpenTelemetryContextAsyncContextStrategy(): void {
     // The types here don't fully align, because our own `Span` type is narrower
     // than the OTEL one - but this is OK for here, as we now we'll only have OTEL spans passed around
     withActiveSpan: withActiveSpan as typeof defaultWithActiveSpan,
+    getTracingChannelBinding: () => {
+      try {
+        const contextManager = (api.context as unknown as ContextApi)._getContextManager();
+        const lookup = contextManager.getAsyncLocalStorageLookup();
+
+        return {
+          asyncLocalStorage: lookup.asyncLocalStorage,
+          getStoreWithActiveSpan: (span: Span) => api.trace.setSpan(api.context.active(), span as api.Span),
+        };
+      } catch {
+        return undefined;
+      }
+    },
   });
 }

packages/server-utils/src/index.ts

@@ -23,3 +23,9 @@ export type {
   RedisTracingChannelFactory,
   RedisTracingChannelSubscribers,
 } from './redis/redis-dc-subscriber';
+export { bindTracingChannelToSpan } from './tracing-channel';
+export type {
+  SentryTracingChannel,
+  TracingChannelBindingHandle,
+  TracingChannelPayloadWithSpan,
+} from './tracing-channel';

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

@@ -0,0 +1,153 @@
+import type { TracingChannel, TracingChannelSubscribers } from 'node:diagnostics_channel';
+import type { AsyncLocalStorage } from 'node:async_hooks';
+import type { Span } from '@sentry/core';
+import { _INTERNAL_getTracingChannelBinding, debug, captureException, SPAN_STATUS_ERROR } from '@sentry/core';
+import { DEBUG_BUILD } from './debug-build';
+
+export type TracingChannelPayloadWithSpan<TData extends object> = TData & {
+  _sentrySpan?: Span;
+};
+
+/*
+ * A type patch so that we don't have to handle all subscription types.
+ */
+export interface SentryTracingChannel<TData extends object = object> extends Omit<
+  TracingChannel<TData, TracingChannelPayloadWithSpan<TData>>,
+  'subscribe' | 'unsubscribe'
+> {
+  subscribe(subscribers: Partial<TracingChannelSubscribers<TracingChannelPayloadWithSpan<TData>>>): void;
+  unsubscribe(subscribers: Partial<TracingChannelSubscribers<TracingChannelPayloadWithSpan<TData>>>): void;
+}
+
+interface TracingChannelManualBindingOptions {
+  /**
+   * Whether the span is ended automatically (`auto`, default) or left to the caller (`manual`).
+   */
+  lifecycle: 'manual';
+}
+
+interface TracingChannelAutoBindingOptions<TData extends object = object> {
+  /**
+   * Whether the span is ended automatically (`auto`, default) or left to the caller (`manual`).
+   */
+  lifecycle?: 'auto' | undefined;
+
+  /**
+   * Invoked with the span and the channel context object once the traced operation completes
+   * Use it to enrich the span from the result/error (branch on `'error' in data` / `'result' in data`) or to run cleanup.
+   */
+  beforeSpanEnd?: (span: Span, data: TracingChannelPayloadWithSpan<TData>) => void;
+
+  /**
+   * Whether a thrown error is captured as a Sentry event. The span is always marked with error
+   * status regardless. Defaults to `true`.
+   * Set `false` for instrumentation that only annotates the span and lets the error be captured at the boundary that owns it (e.g. db spans).
+   */
+  captureError?: boolean;
+}
+
+export type TracingChannelBindingOptions<TData extends object = object> =
+  | TracingChannelAutoBindingOptions<TData>
+  | TracingChannelManualBindingOptions;
+
+/** Returned by {@link bindTracingChannelToSpan}: the bound channel plus a teardown handle. */
+export interface TracingChannelBindingHandle<TData extends object = object> {
+  /** The tracing channel with the span bound into async context (and, in `auto` mode, its lifecycle subscribed). */
+  channel: SentryTracingChannel<TData>;
+  /**
+   * Tears down the binding: unsubscribes the auto lifecycle handlers and unbinds the start store.
+   * Idempotent, and a no-op when no async context binding was available.
+   */
+  unbind: () => void;
+}
+
+const NOOP = (): void => {};
+
+export function bindTracingChannelToSpan<TData extends object>(
+  channel: TracingChannel<TData, TData>,
+  getSpan: (data: TracingChannelPayloadWithSpan<TData>) => Span,
+  opts?: TracingChannelBindingOptions<TData>,
+): TracingChannelBindingHandle<TData> {
+  const sentryChannel = channel as SentryTracingChannel<TData>;
+  const binding = _INTERNAL_getTracingChannelBinding();
+
+  if (!binding) {
+    DEBUG_BUILD && debug.log('[TracingChannel] Could not access async context binding.');
+    return { channel: sentryChannel, unbind: NOOP };
+  }
+
+  const asyncLocalStorage = binding.asyncLocalStorage as AsyncLocalStorage<TData>;
+
+  channel.start.bindStore(asyncLocalStorage, (data: TracingChannelPayloadWithSpan<TData>) => {
+    const span = getSpan(data);
+    data._sentrySpan = span;
+
+    return binding.getStoreWithActiveSpan(span) as TData;
+  });
+
+  const unbindStore = (): void => {
+    channel.start.unbindStore(asyncLocalStorage);
+  };
+
+  if (opts && 'lifecycle' in opts && opts.lifecycle === 'manual') {
+    return { channel: sentryChannel, unbind: unbindStore };
+  }
+
+  const beforeSpanEnd = opts?.beforeSpanEnd;
+
+  const subscribers: Partial<TracingChannelSubscribers<TracingChannelPayloadWithSpan<TData>>> = {
+    start: NOOP,
+    asyncStart: NOOP,
+    end(data) {
+      // The operation settled synchronously (returned or threw)
+      // Presence checks because caller can return `undefined` result or throw a falsy value.
+      if ('error' in data || 'result' in data) {
+        endBoundSpan(data, beforeSpanEnd);
+      }
+    },
+    error(data) {
+      if (opts?.captureError !== false) {
+        captureException(data.error, {
+          mechanism: {
+            type: 'auto.diagnostic_channels.bind_span',
+            handled: false,
+          },
+        });
+      }
+      data._sentrySpan?.setStatus({ code: SPAN_STATUS_ERROR, message: getErrorMessage(data.error) });
+    },
+    asyncEnd(data) {
+      endBoundSpan(data, beforeSpanEnd);
+    },
+  };
+
+  sentryChannel.subscribe(subscribers);
+
+  return {
+    channel: sentryChannel,
+    unbind: () => {
+      sentryChannel.unsubscribe(subscribers);
+      unbindStore();
+    },
+  };
+}
+
+function endBoundSpan<TData extends object>(
+  data: TracingChannelPayloadWithSpan<TData>,
+  beforeSpanEnd: TracingChannelAutoBindingOptions<TData>['beforeSpanEnd'],
+): void {
+  const span = data._sentrySpan;
+  if (!span) {
+    return;
+  }
+  beforeSpanEnd?.(span, data);
+  span.end();
+}
+
+/** Best-effort short message for a span status: an error-like's `message`, otherwise its string form. */
+function getErrorMessage(error: unknown): string {
+  if (error && typeof error === 'object' && 'message' in error && typeof error.message === 'string') {
+    return error.message;
+  }
+  return String(error);
+}