|
| 1 | +import type { Context } from '@opentelemetry/api'; |
| 2 | +import { AsyncLocalStorageContextManager } from '@opentelemetry/context-async-hooks'; |
| 3 | +import { setHubOnContext } from '@sentry/opentelemetry'; |
| 4 | +import { getCurrentHub } from '../sdk/hub'; |
| 5 | + |
| 6 | +import { getCurrentScope, getIsolationScope } from './../sdk/api'; |
| 7 | +import { Scope } from './../sdk/scope'; |
| 8 | +import type { CurrentScopes } from './../sdk/types'; |
| 9 | +import { getScopesFromContext, setScopesOnContext } from './../utils/contextData'; |
| 10 | + |
| 11 | +/** |
| 12 | + * This is a custom ContextManager for OpenTelemetry, which extends the default AsyncLocalStorageContextManager. |
| 13 | + * It ensures that we create a new hub per context, so that the OTEL Context & the Sentry Hub are always in sync. |
| 14 | + * |
| 15 | + * Note that we currently only support AsyncHooks with this, |
| 16 | + * but since this should work for Node 14+ anyhow that should be good enough. |
| 17 | + */ |
| 18 | +export class SentryContextManager extends AsyncLocalStorageContextManager { |
| 19 | + /** |
| 20 | + * Overwrite with() of the original AsyncLocalStorageContextManager |
| 21 | + * to ensure we also create a new hub per context. |
| 22 | + */ |
| 23 | + public with<A extends unknown[], F extends (...args: A) => ReturnType<F>>( |
| 24 | + context: Context, |
| 25 | + fn: F, |
| 26 | + thisArg?: ThisParameterType<F>, |
| 27 | + ...args: A |
| 28 | + ): ReturnType<F> { |
| 29 | + const previousScopes = getScopesFromContext(context); |
| 30 | + |
| 31 | + const currentScope = previousScopes ? previousScopes.scope : getCurrentScope(); |
| 32 | + const isolationScope = previousScopes ? previousScopes.isolationScope : getIsolationScope(); |
| 33 | + |
| 34 | + const newCurrentScope = currentScope.clone(); |
| 35 | + const scopes: CurrentScopes = { scope: newCurrentScope, isolationScope }; |
| 36 | + |
| 37 | + // We also need to "mock" the hub on the context, as the original @sentry/opentelemetry uses that... |
| 38 | + const mockHub = { ...getCurrentHub(), getScope: () => scopes.scope }; |
| 39 | + |
| 40 | + const ctx1 = setHubOnContext(context, mockHub); |
| 41 | + const ctx2 = setScopesOnContext(ctx1, scopes); |
| 42 | + |
| 43 | + return super.with(ctx2, fn, thisArg, ...args); |
| 44 | + } |
| 45 | +} |
0 commit comments