diff --git a/packages/node/src/async/hooks.ts b/packages/node/src/async/hooks.ts new file mode 100644 index 000000000000..840236d5b2c1 --- /dev/null +++ b/packages/node/src/async/hooks.ts @@ -0,0 +1,48 @@ +import type { Carrier, Hub, RunWithAsyncContextOptions } from '@sentry/core'; +import { ensureHubOnCarrier, getHubFromCarrier, setAsyncContextStrategy } from '@sentry/core'; +import * as async_hooks from 'async_hooks'; + +interface AsyncLocalStorage { + getStore(): T | undefined; + // eslint-disable-next-line @typescript-eslint/no-explicit-any + run(store: T, callback: (...args: TArgs) => R, ...args: TArgs): R; +} + +type AsyncLocalStorageConstructor = { new (): AsyncLocalStorage }; +// AsyncLocalStorage only exists in async_hook after Node v12.17.0 or v13.10.0 +type NewerAsyncHooks = typeof async_hooks & { AsyncLocalStorage: AsyncLocalStorageConstructor }; + +/** + * Sets the async context strategy to use AsyncLocalStorage which requires Node v12.17.0 or v13.10.0. + */ +export function setHooksAsyncContextStrategy(): void { + const asyncStorage = new (async_hooks as NewerAsyncHooks).AsyncLocalStorage(); + + function getCurrentHub(): Hub | undefined { + return asyncStorage.getStore(); + } + + function createNewHub(parent: Hub | undefined): Hub { + const carrier: Carrier = {}; + ensureHubOnCarrier(carrier, parent); + return getHubFromCarrier(carrier); + } + + function runWithAsyncContext(callback: (hub: Hub) => 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(existingHub); + } + + const newHub = createNewHub(existingHub); + + return asyncStorage.run(newHub, () => { + return callback(newHub); + }); + } + + setAsyncContextStrategy({ getCurrentHub, runWithAsyncContext }); +} diff --git a/packages/node/test/async/hooks.test.ts b/packages/node/test/async/hooks.test.ts new file mode 100644 index 000000000000..dbd1904c34dc --- /dev/null +++ b/packages/node/test/async/hooks.test.ts @@ -0,0 +1,108 @@ +import { getCurrentHub, Hub, runWithAsyncContext, setAsyncContextStrategy } from '@sentry/core'; + +import { setHooksAsyncContextStrategy } from '../../src/async/hooks'; +import { conditionalTest } from '../utils'; + +conditionalTest({ min: 12 })('async_hooks', () => { + afterAll(() => { + // clear the strategy + setAsyncContextStrategy(undefined); + }); + + test('without context', () => { + const hub = getCurrentHub(); + expect(hub).toEqual(new Hub()); + }); + + test('without strategy hubs should be equal', () => { + runWithAsyncContext(hub1 => { + runWithAsyncContext(hub2 => { + expect(hub1).toBe(hub2); + }); + }); + }); + + test('hub scope inheritance', () => { + setHooksAsyncContextStrategy(); + + const globalHub = getCurrentHub(); + globalHub.setExtra('a', 'b'); + + runWithAsyncContext(hub1 => { + expect(hub1).toEqual(globalHub); + + hub1.setExtra('c', 'd'); + expect(hub1).not.toEqual(globalHub); + + runWithAsyncContext(hub2 => { + expect(hub2).toEqual(hub1); + expect(hub2).not.toEqual(globalHub); + + hub2.setExtra('e', 'f'); + expect(hub2).not.toEqual(hub1); + }); + }); + }); + + test('context single instance', () => { + setHooksAsyncContextStrategy(); + + runWithAsyncContext(hub => { + expect(hub).toBe(getCurrentHub()); + }); + }); + + test('context within a context not reused', () => { + setHooksAsyncContextStrategy(); + + runWithAsyncContext(hub1 => { + runWithAsyncContext(hub2 => { + expect(hub1).not.toBe(hub2); + }); + }); + }); + + test('context within a context reused when requested', () => { + setHooksAsyncContextStrategy(); + + runWithAsyncContext(hub1 => { + runWithAsyncContext( + hub2 => { + expect(hub1).toBe(hub2); + }, + { reuseExisting: true }, + ); + }); + }); + + test('concurrent hub contexts', done => { + setHooksAsyncContextStrategy(); + + let d1done = false; + let d2done = false; + + runWithAsyncContext(hub => { + hub.getStack().push({ client: 'process' } as any); + expect(hub.getStack()[1]).toEqual({ client: 'process' }); + // Just in case so we don't have to worry which one finishes first + // (although it always should be d2) + setTimeout(() => { + d1done = true; + if (d2done) { + done(); + } + }); + }); + + runWithAsyncContext(hub => { + hub.getStack().push({ client: 'local' } as any); + expect(hub.getStack()[1]).toEqual({ client: 'local' }); + setTimeout(() => { + d2done = true; + if (d1done) { + done(); + } + }); + }); + }); +});