Add otel contextManager & use it for hub management

Author: mydea

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,7 +4,7 @@ import { getCurrentHub } from '@sentry/core';
 import { SentryPropagator, SentrySpanProcessor } from '@sentry/opentelemetry-node';
 
 import type { NodeExperimentalClient } from './client';
-import { setOtelContextAsyncContextStrategy } from './otelContextAsyncContextStrategy';
+import { SentryContextManager } from './otelContextManager';
 
 /**
  * Initialize OpenTelemetry for Node.
@@ -23,13 +23,16 @@ 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,
   });
 
-  setOtelContextAsyncContextStrategy();
-
   // Cleanup function
   return () => {
     void provider.forceFlush();

packages/node-experimental/src/sdk/otelContextAsyncContextStrategy.ts renamed to packages/node-experimental/src/sdk/otelAsyncContextStrategy.ts

@@ -1,8 +1,8 @@
 import * as api from '@opentelemetry/api';
-import type { Carrier, Hub, RunWithAsyncContextOptions } from '@sentry/core';
-import { ensureHubOnCarrier, getHubFromCarrier, setAsyncContextStrategy } from '@sentry/core';
+import type { Hub, RunWithAsyncContextOptions } from '@sentry/core';
+import { setAsyncContextStrategy } from '@sentry/core';
 
-const hubKey = api.createContextKey('sentry_hub');
+import { OTEL_CONTEXT_HUB_KEY } from './otelContextManager';
 
 /**
  * Sets the async context strategy to use follow the OTEL context under the hood.
@@ -12,15 +12,10 @@ export function setOtelContextAsyncContextStrategy(): void {
     const ctx = api.context.active();
 
     // Returning undefined means the global hub will be used
-    return ctx.getValue(hubKey) as Hub | undefined;
-  }
-
-  function createNewHub(parent: Hub | undefined): Hub {
-    const carrier: Carrier = {};
-    ensureHubOnCarrier(carrier, parent);
-    return getHubFromCarrier(carrier);
+    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();
 
@@ -30,11 +25,10 @@ export function setOtelContextAsyncContextStrategy(): void {
       return callback();
     }
 
-    const newHub = createNewHub(existingHub);
-
     const ctx = api.context.active();
 
-    return api.context.with(ctx.setValue(hubKey, newHub), () => {
+    // We depend on the otelContextManager to handle the context/hub
+    return api.context.with(ctx, () => {
       return callback();
     });
   }

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

@@ -0,0 +1,32 @@
+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);
+}
+
+/** TODO docs */
+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"
@@ -4414,10 +4419,10 @@
     fflate "^0.4.4"
     mitt "^1.1.3"
 
-"@sentry/[email protected].0":
-  version "0.6.0"
-  resolved "https://registry.yarnpkg.com/@sentry/bundler-plugin-core/-/bundler-plugin-core-0.6.0.tgz#70ad3740b2f90cdca1aff5fdbcd7306566a2f51e"
-  integrity sha512-gDPBkFxiOkc525U9pxnGMI5B2DAG0+UCsNuiNgl9+AieDcPSYTwdzfGHytxDZrQgPMvIHEnTAp1VlNB+6UxUGQ==
+"@sentry/[email protected].1":
+  version "0.6.1"
+  resolved "https://registry.yarnpkg.com/@sentry/bundler-plugin-core/-/bundler-plugin-core-0.6.1.tgz#6c6a2ff3cdc98cd0ff1c30c59408cee9f067adf2"
+  integrity sha512-EecCJKp9ERM7J93DNDJTvkY78UiD/IfOjBdXWnaUVE0n619O7LfMVjwlXzxRJKl2x05dBE3lDraILLDGxCf6fg==
   dependencies:
     "@sentry/cli" "^2.17.0"
     "@sentry/node" "^7.19.0"
@@ -4463,12 +4468,12 @@
     proxy-from-env "^1.1.0"
     which "^2.0.2"
 
-"@sentry/vite-plugin@^0.6.0":
-  version "0.6.0"
-  resolved "https://registry.yarnpkg.com/@sentry/vite-plugin/-/vite-plugin-0.6.0.tgz#3902a5224d52b06d753a1deeb6b722bf6523840c"
-  integrity sha512-3J1ESvbI5okGJaSWm+gTAOOIa96u4ZwfI/C3n+0HSStz3e4vGiGUh59iNyb1/2m5HFgR5OLaHNfAvlyP8GM/ew==
+"@sentry/vite-plugin@^0.6.1":
+  version "0.6.1"
+  resolved "https://registry.yarnpkg.com/@sentry/vite-plugin/-/vite-plugin-0.6.1.tgz#31eb744e8d87b1528eed8d41433647727a62e7c0"
+  integrity sha512-qkvKaSOcNhNWcdxRXLSs+8cF3ey0XIRmEzTl8U7sTTcZwuOMHsJB+HsYij6aTGaqsKfP8w1ozVt9szBAiL4//w==
   dependencies:
-    "@sentry/bundler-plugin-core" "0.6.0"
+    "@sentry/bundler-plugin-core" "0.6.1"
 
 "@sentry/[email protected]":
   version "1.19.0"