feat(nextjs): Add AsyncLocalStorage async context strategy to edge SDK (#8720)

Author: lforst

packages/e2e-tests/README.md

@@ -10,6 +10,9 @@ Prerequisites: Docker
 - Copy `.env.example` to `.env`
 - Fill in auth information in `.env` for an example Sentry project
   - The `E2E_TEST_AUTH_TOKEN` must have all the default permissions
+- Run `yarn build:tarball` in the root of the repository
+
+To finally run all of the tests:
 
 ```bash
 yarn test:e2e

packages/e2e-tests/test-applications/nextjs-app-dir/pages/api/async-context-edge-endpoint.ts

@@ -0,0 +1,27 @@
+import * as Sentry from '@sentry/nextjs';
+
+export const config = {
+  runtime: 'edge',
+};
+
+export default async function handler() {
+  // Without `runWithAsyncContext` and a working async context strategy the two spans created by `Sentry.trace()` would be nested.
+
+  const outerSpanPromise = Sentry.runWithAsyncContext(() => {
+    return Sentry.trace({ name: 'outer-span' }, () => {
+      return new Promise<void>(resolve => setTimeout(resolve, 300));
+    });
+  });
+
+  setTimeout(() => {
+    Sentry.runWithAsyncContext(() => {
+      return Sentry.trace({ name: 'inner-span' }, () => {
+        return new Promise<void>(resolve => setTimeout(resolve, 100));
+      });
+    });
+  }, 100);
+
+  await outerSpanPromise;
+
+  return new Response('ok', { status: 200 });
+}

packages/e2e-tests/test-applications/nextjs-app-dir/tests/async-context-edge.test.ts

@@ -0,0 +1,22 @@
+import { test, expect } from '@playwright/test';
+import { waitForTransaction, waitForError } from '../event-proxy-server';
+
+test('Should allow for async context isolation in the edge SDK', async ({ request }) => {
+  // test.skip(process.env.TEST_ENV === 'development', "Doesn't work in dev mode.");
+
+  const edgerouteTransactionPromise = waitForTransaction('nextjs-13-app-dir', async transactionEvent => {
+    return transactionEvent?.transaction === 'GET /api/async-context-edge-endpoint';
+  });
+
+  await request.get('/api/async-context-edge-endpoint');
+
+  const asyncContextEdgerouteTransaction = await edgerouteTransactionPromise;
+
+  const outerSpan = asyncContextEdgerouteTransaction.spans?.find(span => span.description === 'outer-span');
+  const innerSpan = asyncContextEdgerouteTransaction.spans?.find(span => span.description === 'inner-span');
+
+  // @ts-ignore parent_span_id exists
+  expect(outerSpan?.parent_span_id).toStrictEqual(asyncContextEdgerouteTransaction.contexts?.trace?.span_id);
+  // @ts-ignore parent_span_id exists
+  expect(innerSpan?.parent_span_id).toStrictEqual(asyncContextEdgerouteTransaction.contexts?.trace?.span_id);
+});

packages/nextjs/src/edge/asyncLocalStorageAsyncContextStrategy.ts

@@ -0,0 +1,55 @@
+import type { Carrier, Hub, RunWithAsyncContextOptions } from '@sentry/core';
+import { ensureHubOnCarrier, getHubFromCarrier, setAsyncContextStrategy } from '@sentry/core';
+import { GLOBAL_OBJ, logger } from '@sentry/utils';
+
+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;
+}
+
+// eslint-disable-next-line @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-explicit-any
+const MaybeGlobalAsyncLocalStorage = (GLOBAL_OBJ as any).AsyncLocalStorage;
+
+/**
+ * Sets the async context strategy to use AsyncLocalStorage which should be available in the edge runtime.
+ */
+export function setAsyncLocalStorageAsyncContextStrategy(): void {
+  if (!MaybeGlobalAsyncLocalStorage) {
+    __DEBUG_BUILD__ &&
+      logger.warn(
+        "Tried to register AsyncLocalStorage async context strategy in a runtime that doesn't support AsyncLocalStorage.",
+      );
+    return;
+  }
+
+  const asyncStorage: AsyncLocalStorage<Hub> = new MaybeGlobalAsyncLocalStorage();
+
+  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: () => 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();
+    }
+
+    const newHub = createNewHub(existingHub);
+
+    return asyncStorage.run(newHub, () => {
+      return callback();
+    });
+  }
+
+  setAsyncContextStrategy({ getCurrentHub, runWithAsyncContext });
+}

packages/nextjs/src/edge/index.ts

@@ -9,6 +9,7 @@ import {
 } from '@sentry/utils';
 
 import { getVercelEnv } from '../common/getVercelEnv';
+import { setAsyncLocalStorageAsyncContextStrategy } from './asyncLocalStorageAsyncContextStrategy';
 import { EdgeClient } from './edgeclient';
 import { makeEdgeTransport } from './transport';
 
@@ -20,6 +21,8 @@ export type EdgeOptions = Options;
 
 /** Inits the Sentry NextJS SDK on the Edge Runtime. */
 export function init(options: EdgeOptions = {}): void {
+  setAsyncLocalStorageAsyncContextStrategy();
+
   if (options.defaultIntegrations === undefined) {
     options.defaultIntegrations = defaultIntegrations;
   }