feat(core): Extend AsyncContextStrategy to allow reuse of existing …

…context (#7778)
getsentry · Apr 6, 2023 · 715876b · 715876b
1 parent f0f9b8a
commit 715876b
Show file tree

Hide file tree

Showing 4 changed files with 53 additions and 11 deletions.
diff --git a/packages/core/src/hub.ts b/packages/core/src/hub.ts
@@ -50,12 +50,27 @@ export const API_VERSION = 4;
  */
 const DEFAULT_BREADCRUMBS = 100;
 
+export interface RunWithAsyncContextOptions {
+  /** Whether to reuse an existing async context if one exists. Defaults to false. */
+  reuseExisting?: boolean;
+  /** Instances that should be referenced and retained in the new context */
+  args?: unknown[];
+}
+
 /**
+ * @private Private API with no semver guarantees!
+ *
  * Strategy used to track async context.
  */
 export interface AsyncContextStrategy {
+  /**
+   * Gets the current async context. Returns undefined if there is no current async context.
+   */
   getCurrentHub: () => Hub | undefined;
-  runWithAsyncContext<T>(callback: (hub: Hub) => T, ...args: unknown[]): T;
+  /**
+   * Runs the supplied callback in its own async context.
+   */
+  runWithAsyncContext<T>(callback: (hub: Hub) => T, options: RunWithAsyncContextOptions): T;
 }
 
 /**
@@ -583,13 +598,13 @@ export function setAsyncContextStrategy(strategy: AsyncContextStrategy | undefin
 /**
  * @private Private API with no semver guarantees!
  *
- * Runs the given callback function with the global async context strategy
+ * Runs the supplied callback in its own async context.
  */
-export function runWithAsyncContext<T>(callback: (hub: Hub) => T, ...args: unknown[]): T {
+export function runWithAsyncContext<T>(callback: (hub: Hub) => T, options: RunWithAsyncContextOptions = {}): T {
   const registry = getMainCarrier();
 
   if (registry.__SENTRY__ && registry.__SENTRY__.acs) {
-    return registry.__SENTRY__.acs.runWithAsyncContext(callback, ...args);
+    return registry.__SENTRY__.acs.runWithAsyncContext(callback, options);
   }
 
   // if there was no strategy, fallback to just calling the callback

diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts
@@ -1,5 +1,5 @@
 export type { ClientClass } from './sdk';
-export type { AsyncContextStrategy, Carrier, Layer } from './hub';
+export type { AsyncContextStrategy, Carrier, Layer, RunWithAsyncContextOptions } from './hub';
 export type { OfflineStore, OfflineTransportOptions } from './transports/offline';
 
 export * from './tracing';

diff --git a/packages/node/src/async/domain.ts b/packages/node/src/async/domain.ts
@@ -1,4 +1,4 @@
-import type { Carrier, Hub } from '@sentry/core';
+import type { Carrier, Hub, RunWithAsyncContextOptions } from '@sentry/core';
 import {
   ensureHubOnCarrier,
   getCurrentHub as getCurrentHubCore,
@@ -8,9 +8,13 @@ import {
 import * as domain from 'domain';
 import { EventEmitter } from 'events';
 
-function getCurrentHub(): Hub | undefined {
+function getActiveDomain<T>(): T | undefined {
   // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-explicit-any
-  const activeDomain = (domain as any).active as Carrier;
+  return (domain as any).active as T | undefined;
+}
+
+function getCurrentHub(): Hub | undefined {
+  const activeDomain = getActiveDomain<Carrier>();
 
   // If there's no active domain, just return undefined and the global hub will be used
   if (!activeDomain) {
@@ -22,10 +26,10 @@ function getCurrentHub(): Hub | undefined {
   return getHubFromCarrier(activeDomain);
 }
 
-function runWithAsyncContext<T, A>(callback: (hub: Hub) => T, ...args: A[]): T {
-  const local = domain.create();
+function runWithAsyncContext<T>(callback: (hub: Hub) => T, options: RunWithAsyncContextOptions): T {
+  const local = options?.reuseExisting ? getActiveDomain<domain.Domain>() || domain.create() : domain.create();
 
-  for (const emitter of args) {
+  for (const emitter of options.args || []) {
     if (emitter instanceof EventEmitter) {
       local.add(emitter);
     }

diff --git a/packages/node/test/async/domain.test.ts b/packages/node/test/async/domain.test.ts
@@ -36,6 +36,29 @@ describe('domains', () => {
     });
   });
 
+  test('domain within a domain not reused', () => {
+    setDomainAsyncContextStrategy();
+
+    runWithAsyncContext(hub1 => {
+      runWithAsyncContext(hub2 => {
+        expect(hub1).not.toBe(hub2);
+      });
+    });
+  });
+
+  test('domain within a domain reused when requested', () => {
+    setDomainAsyncContextStrategy();
+
+    runWithAsyncContext(hub1 => {
+      runWithAsyncContext(
+        hub2 => {
+          expect(hub1).toBe(hub2);
+        },
+        { reuseExisting: true },
+      );
+    });
+  });
+
   test('concurrent domain hubs', done => {
     setDomainAsyncContextStrategy();