feat(node): Add AsyncLocalStorage implementation of AsyncContextStrategy (#7800)

Author: timfish

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<T> {
+  getStore(): T | undefined;
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  run<R, TArgs extends any[]>(store: T, callback: (...args: TArgs) => R, ...args: TArgs): R;
+}
+
+type AsyncLocalStorageConstructor = { new <T>(): AsyncLocalStorage<T> };
+// 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<Hub>();
+
+  function getCurrentHub(): Hub | undefined {
+    return asyncStorage.getStore();
+  }
+
+  function createNewHub(parent: Hub | undefined): Hub {
+    const carrier: Carrier = {};
+    ensureHubOnCarrier(carrier, parent);
+    return getHubFromCarrier(carrier);
+  }
+
+  function runWithAsyncContext<T>(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 });
+}

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();
+        }
+      });
+    });
+  });
+});