feat(node-experimental): Sync OTEL context with Sentry AsyncContext

packages/node-experimental/package.json

@@ -37,6 +37,7 @@
     "@opentelemetry/instrumentation-pg": "~0.36.0",
     "@opentelemetry/sdk-trace-node": "~1.15.0",
     "@opentelemetry/semantic-conventions": "~1.15.0",
+    "@opentelemetry/context-async-hooks": "~1.15.0",
     "@prisma/instrumentation": "~5.0.0",
     "@sentry/core": "7.66.0",
     "@sentry/node": "7.66.0",

packages/node-experimental/src/sdk/init.ts

@@ -6,6 +6,7 @@ import { Http } from '../integrations/http';
 import type { NodeExperimentalOptions } from '../types';
 import { NodeExperimentalClient } from './client';
 import { initOtel } from './initOtel';
+import { setOtelContextAsyncContextStrategy } from './otelAsyncContextStrategy';
 
 const ignoredDefaultIntegrations = ['Http', 'Undici'];
 
@@ -35,4 +36,5 @@ export function init(options: NodeExperimentalOptions | undefined = {}): void {
 
   // Always init Otel, even if tracing is disabled, because we need it for trace propagation & the HTTP integration
   initOtel();
+  setOtelContextAsyncContextStrategy();
 }

packages/node-experimental/src/sdk/initOtel.ts

@@ -4,6 +4,7 @@ import { getCurrentHub } from '@sentry/core';
 import { SentryPropagator, SentrySpanProcessor } from '@sentry/opentelemetry-node';
 
 import type { NodeExperimentalClient } from './client';
+import { SentryContextManager } from './otelContextManager';
 
 /**
  * Initialize OpenTelemetry for Node.
@@ -22,9 +23,14 @@ export function initOtel(): () => void {
   });
   provider.addSpanProcessor(new SentrySpanProcessor());
 
+  // We use a custom context manager to keep context in sync with sentry scope
+  const contextManager = new SentryContextManager();
+  contextManager.enable();
+
   // Initialize the provider
   provider.register({
     propagator: new SentryPropagator(),
+    contextManager,
   });
 
   // Cleanup function

packages/node-experimental/src/sdk/otelAsyncContextStrategy.ts

@@ -0,0 +1,38 @@
+import * as api from '@opentelemetry/api';
+import type { Hub, RunWithAsyncContextOptions } from '@sentry/core';
+import { setAsyncContextStrategy } from '@sentry/core';
+
+import { OTEL_CONTEXT_HUB_KEY } from './otelContextManager';
+
+/**
+ * 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)
+ */
+export function setOtelContextAsyncContextStrategy(): void {
+  function getCurrentHub(): Hub | undefined {
+    const ctx = api.context.active();
+
+    // Returning undefined means the global hub will be used
+    return ctx.getValue(OTEL_CONTEXT_HUB_KEY) as Hub | undefined;
+  }
+
+  /* This is more or less a NOOP - we rely on the OTEL context manager for this */
+  function runWithAsyncContext<T>(callback: () => T, options: RunWithAsyncContextOptions): T {
+    const existingHub = getCurrentHub();
+
+    if (existingHub && options?.reuseExisting) {
+      // We're already in an async context, so we don't need to create a new one
+      // just call the callback with the current hub
+      return callback();
+    }
+
+    const ctx = api.context.active();
+
+    // We depend on the otelContextManager to handle the context/hub
+    return api.context.with(ctx, () => {
+      return callback();
+    });
+  }
+
+  setAsyncContextStrategy({ getCurrentHub, runWithAsyncContext });
+}

packages/node-experimental/src/sdk/otelContextManager.ts

@@ -0,0 +1,38 @@
+import type { Context } from '@opentelemetry/api';
+import * as api from '@opentelemetry/api';
+import { AsyncLocalStorageContextManager } from '@opentelemetry/context-async-hooks';
+import type { Carrier, Hub } from '@sentry/core';
+import { ensureHubOnCarrier, getCurrentHub, getHubFromCarrier } from '@sentry/core';
+
+export const OTEL_CONTEXT_HUB_KEY = api.createContextKey('sentry_hub');
+
+function createNewHub(parent: Hub | undefined): Hub {
+  const carrier: Carrier = {};
+  ensureHubOnCarrier(carrier, parent);
+  return getHubFromCarrier(carrier);
+}
+
+/**
+ * This is a custom ContextManager for OpenTelemetry, which extends the default AsyncLocalStorageContextManager.
+ * It ensures that we create a new hub per context, so that the OTEL Context & the Sentry Hub are always in sync.
+ *
+ * Note that we currently only support AsyncHooks with this,
+ * but since this should work for Node 14+ anyhow that should be good enough.
+ */
+export class SentryContextManager extends AsyncLocalStorageContextManager {
+  /**
+   * Overwrite with() of the original AsyncLocalStorageContextManager
+   * to ensure we also create a new hub per context.
+   */
+  public with<A extends unknown[], F extends (...args: A) => ReturnType<F>>(
+    context: Context,
+    fn: F,
+    thisArg?: ThisParameterType<F>,
+    ...args: A
+  ): ReturnType<F> {
+    const existingHub = getCurrentHub();
+    const newHub = createNewHub(existingHub);
+
+    return super.with(context.setValue(OTEL_CONTEXT_HUB_KEY, newHub), fn, thisArg, ...args);
+  }
+}

yarn.lock

@@ -3733,6 +3733,11 @@
   dependencies:
     tslib "^2.3.1"
 
+"@opentelemetry/context-async-hooks@~1.15.0":
+  version "1.15.2"
+  resolved "https://registry.yarnpkg.com/@opentelemetry/context-async-hooks/-/context-async-hooks-1.15.2.tgz#116bd5fef231137198d5bf551e8c0521fbdfe928"
+  integrity sha512-VAMHG67srGFQDG/N2ns5AyUT9vUcoKpZ/NpJ5fDQIPfJd7t3ju+aHwvDsMcrYBWuCh03U3Ky6o16+872CZchBg==
+
 "@opentelemetry/context-base@^0.12.0":
   version "0.12.0"
   resolved "https://registry.yarnpkg.com/@opentelemetry/context-base/-/context-base-0.12.0.tgz#4906ae27359d3311e3dea1b63770a16f60848550"