feat(next): Handle existing root spans for isolation scope (#11479)

This updates handling of next.js instrumentation to re-use an isolation
scope from a root span.

This should ensure we have consistent isolation scopes, no matter if
next.js auto creates spans or not.

Author: mydea

packages/nextjs/src/common/utils/withIsolationScopeOrReuseFromRootSpan.ts

@@ -0,0 +1,39 @@
+import {
+  getActiveSpan,
+  getCapturedScopesOnSpan,
+  getDefaultIsolationScope,
+  getRootSpan,
+  spanToJSON,
+  withIsolationScope,
+} from '@sentry/core';
+import type { Scope } from '@sentry/types';
+
+/**
+ * Wrap a callback with a new isolation scope.
+ * However, if we have an active root span that was generated by next, we want to reuse the isolation scope from that span.
+ */
+export function withIsolationScopeOrReuseFromRootSpan<T>(cb: (isolationScope: Scope) => T): T {
+  const activeSpan = getActiveSpan();
+
+  if (!activeSpan) {
+    return withIsolationScope(cb);
+  }
+
+  const rootSpan = getRootSpan(activeSpan);
+
+  // Verify this is a next span
+  if (!spanToJSON(rootSpan).data?.['next.route']) {
+    return withIsolationScope(cb);
+  }
+
+  const scopes = getCapturedScopesOnSpan(rootSpan);
+
+  const isolationScope = scopes.isolationScope;
+
+  // If this is the default isolation scope, we still want to fork one
+  if (isolationScope === getDefaultIsolationScope()) {
+    return withIsolationScope(cb);
+  }
+
+  return withIsolationScope(isolationScope, cb);
+}

packages/nextjs/src/common/utils/wrapperUtils.ts

@@ -10,13 +10,13 @@ import {
   startSpan,
   startSpanManual,
   withActiveSpan,
-  withIsolationScope,
 } from '@sentry/core';
 import type { Span } from '@sentry/types';
 import { isString } from '@sentry/utils';
 
 import { platformSupportsStreaming } from './platformSupportsStreaming';
 import { autoEndSpanOnResponseEnd, flushQueue } from './responseEnd';
+import { withIsolationScopeOrReuseFromRootSpan } from './withIsolationScopeOrReuseFromRootSpan';
 
 declare module 'http' {
   interface IncomingMessage {
@@ -89,7 +89,7 @@ export function withTracedServerSideDataFetcher<F extends (...args: any[]) => Pr
   },
 ): (...params: Parameters<F>) => Promise<ReturnType<F>> {
   return async function (this: unknown, ...args: Parameters<F>): Promise<ReturnType<F>> {
-    return withIsolationScope(async isolationScope => {
+    return withIsolationScopeOrReuseFromRootSpan(async isolationScope => {
       isolationScope.setSDKProcessingMetadata({
         request: req,
       });

packages/nextjs/src/common/withServerActionInstrumentation.ts

@@ -6,14 +6,14 @@ import {
   getClient,
   handleCallbackErrors,
   startSpan,
-  withIsolationScope,
 } from '@sentry/core';
 import { logger } from '@sentry/utils';
 
 import { DEBUG_BUILD } from './debug-build';
 import { isNotFoundNavigationError, isRedirectNavigationError } from './nextNavigationErrorUtils';
 import { platformSupportsStreaming } from './utils/platformSupportsStreaming';
 import { flushQueue } from './utils/responseEnd';
+import { withIsolationScopeOrReuseFromRootSpan } from './utils/withIsolationScopeOrReuseFromRootSpan';
 
 interface Options {
   formData?: FormData;
@@ -58,7 +58,7 @@ async function withServerActionInstrumentationImplementation<A extends (...args:
   callback: A,
 ): Promise<ReturnType<A>> {
   addTracingExtensions();
-  return withIsolationScope(isolationScope => {
+  return withIsolationScopeOrReuseFromRootSpan(isolationScope => {
     const sendDefaultPii = getClient()?.getOptions().sendDefaultPii;
 
     let sentryTraceHeader;

packages/nextjs/src/common/wrapApiHandlerWithSentry.ts

@@ -5,14 +5,14 @@ import {
   continueTrace,
   setHttpStatus,
   startSpanManual,
-  withIsolationScope,
 } from '@sentry/core';
 import { consoleSandbox, isString, logger, objectify, stripUrlQueryAndFragment } from '@sentry/utils';
 
 import { SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN } from '@sentry/core';
 import type { AugmentedNextApiRequest, AugmentedNextApiResponse, NextApiHandler } from './types';
 import { platformSupportsStreaming } from './utils/platformSupportsStreaming';
 import { flushQueue } from './utils/responseEnd';
+import { withIsolationScopeOrReuseFromRootSpan } from './utils/withIsolationScopeOrReuseFromRootSpan';
 
 /**
  * Wrap the given API route handler for tracing and error capturing. Thin wrapper around `withSentry`, which only
@@ -54,7 +54,7 @@ export function wrapApiHandlerWithSentry(apiHandler: NextApiHandler, parameteriz
 
       addTracingExtensions();
 
-      return withIsolationScope(isolationScope => {
+      return withIsolationScopeOrReuseFromRootSpan(isolationScope => {
         return continueTrace(
           {
             // TODO(v8): Make it so that continue trace will allow null as sentryTrace value and remove this fallback here

packages/nextjs/src/common/wrapApiHandlerWithSentryVercelCrons.ts

@@ -1,7 +1,8 @@
-import { addTracingExtensions, captureCheckIn, withIsolationScope } from '@sentry/core';
+import { addTracingExtensions, captureCheckIn } from '@sentry/core';
 import type { NextApiRequest } from 'next';
 
 import type { VercelCronsConfig } from './types';
+import { withIsolationScopeOrReuseFromRootSpan } from './utils/withIsolationScopeOrReuseFromRootSpan';
 
 type EdgeRequest = {
   nextUrl: URL;
@@ -19,7 +20,7 @@ export function wrapApiHandlerWithSentryVercelCrons<F extends (...args: any[]) =
   return new Proxy(handler, {
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     apply: (originalFunction, thisArg, args: any[]) => {
-      return withIsolationScope(() => {
+      return withIsolationScopeOrReuseFromRootSpan(() => {
         if (!args || !args[0]) {
           return originalFunction.apply(thisArg, args);
         }

packages/nextjs/src/common/wrapGenerationFunctionWithSentry.ts

@@ -8,7 +8,6 @@ import {
   getCurrentScope,
   handleCallbackErrors,
   startSpanManual,
-  withIsolationScope,
 } from '@sentry/core';
 import type { WebFetchHeaders } from '@sentry/types';
 import { propagationContextFromHeaders, winterCGHeadersToDict } from '@sentry/utils';
@@ -17,6 +16,7 @@ import { SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN } from '@sentry/core';
 import type { GenerationFunctionContext } from '../common/types';
 import { isNotFoundNavigationError, isRedirectNavigationError } from './nextNavigationErrorUtils';
 import { commonObjectToPropagationContext } from './utils/commonObjectTracing';
+import { withIsolationScopeOrReuseFromRootSpan } from './utils/withIsolationScopeOrReuseFromRootSpan';
 
 /**
  * Wraps a generation function (e.g. generateMetadata) with Sentry error and performance instrumentation.
@@ -47,7 +47,7 @@ export function wrapGenerationFunctionWithSentry<F extends (...args: any[]) => a
         data = { params, searchParams };
       }
 
-      return withIsolationScope(isolationScope => {
+      return withIsolationScopeOrReuseFromRootSpan(isolationScope => {
         isolationScope.setSDKProcessingMetadata({
           request: {
             headers: headers ? winterCGHeadersToDict(headers) : undefined,

packages/nextjs/src/common/wrapPageComponentWithSentry.ts

@@ -1,5 +1,6 @@
-import { addTracingExtensions, captureException, getCurrentScope, withIsolationScope } from '@sentry/core';
+import { addTracingExtensions, captureException, getCurrentScope } from '@sentry/core';
 import { extractTraceparentData } from '@sentry/utils';
+import { withIsolationScopeOrReuseFromRootSpan } from './utils/withIsolationScopeOrReuseFromRootSpan';
 
 interface FunctionComponent {
   (...args: unknown[]): unknown;
@@ -25,7 +26,7 @@ export function wrapPageComponentWithSentry(pageComponent: FunctionComponent | C
   if (isReactClassComponent(pageComponent)) {
     return class SentryWrappedPageComponent extends pageComponent {
       public render(...args: unknown[]): unknown {
-        return withIsolationScope(() => {
+        return withIsolationScopeOrReuseFromRootSpan(() => {
           const scope = getCurrentScope();
           // We extract the sentry trace data that is put in the component props by datafetcher wrappers
           const sentryTraceData =
@@ -60,7 +61,7 @@ export function wrapPageComponentWithSentry(pageComponent: FunctionComponent | C
   } else if (typeof pageComponent === 'function') {
     return new Proxy(pageComponent, {
       apply(target, thisArg, argArray: [{ _sentryTraceData?: string } | undefined]) {
-        return withIsolationScope(() => {
+        return withIsolationScopeOrReuseFromRootSpan(() => {
           const scope = getCurrentScope();
           // We extract the sentry trace data that is put in the component props by datafetcher wrappers
           const sentryTraceData = argArray?.[0]?._sentryTraceData;

packages/nextjs/src/common/wrapRouteHandlerWithSentry.ts

@@ -8,14 +8,14 @@ import {
   handleCallbackErrors,
   setHttpStatus,
   startSpan,
-  withIsolationScope,
 } from '@sentry/core';
 import { winterCGHeadersToDict } from '@sentry/utils';
 
 import { isNotFoundNavigationError, isRedirectNavigationError } from './nextNavigationErrorUtils';
 import type { RouteHandlerContext } from './types';
 import { platformSupportsStreaming } from './utils/platformSupportsStreaming';
 import { flushQueue } from './utils/responseEnd';
+import { withIsolationScopeOrReuseFromRootSpan } from './utils/withIsolationScopeOrReuseFromRootSpan';
 
 /**
  * Wraps a Next.js route handler with performance and error instrumentation.
@@ -29,7 +29,7 @@ export function wrapRouteHandlerWithSentry<F extends (...args: any[]) => any>(
   const { method, parameterizedRoute, headers } = context;
   return new Proxy(routeHandler, {
     apply: (originalFunction, thisArg, args) => {
-      return withIsolationScope(async isolationScope => {
+      return withIsolationScopeOrReuseFromRootSpan(async isolationScope => {
         isolationScope.setSDKProcessingMetadata({
           request: {
             headers: headers ? winterCGHeadersToDict(headers) : undefined,

packages/nextjs/src/common/wrapServerComponentWithSentry.ts

@@ -7,7 +7,6 @@ import {
   getCurrentScope,
   handleCallbackErrors,
   startSpanManual,
-  withIsolationScope,
 } from '@sentry/core';
 import { propagationContextFromHeaders, winterCGHeadersToDict } from '@sentry/utils';
 
@@ -16,6 +15,7 @@ import { isNotFoundNavigationError, isRedirectNavigationError } from '../common/
 import type { ServerComponentContext } from '../common/types';
 import { commonObjectToPropagationContext } from './utils/commonObjectTracing';
 import { flushQueue } from './utils/responseEnd';
+import { withIsolationScopeOrReuseFromRootSpan } from './utils/withIsolationScopeOrReuseFromRootSpan';
 
 /**
  * Wraps an `app` directory server component with Sentry error instrumentation.
@@ -34,7 +34,7 @@ export function wrapServerComponentWithSentry<F extends (...args: any[]) => any>
   return new Proxy(appDirComponent, {
     apply: (originalFunction, thisArg, args) => {
       // TODO: If we ever allow withIsolationScope to take a scope, we should pass a scope here that is shared between all of the server components, similar to what `commonObjectToPropagationContext` does.
-      return withIsolationScope(isolationScope => {
+      return withIsolationScopeOrReuseFromRootSpan(isolationScope => {
         const completeHeadersDict: Record<string, string> = context.headers
           ? winterCGHeadersToDict(context.headers)
           : {};

packages/nextjs/test/edge/withIsolationScopeOrReuseFromRootSpan.test.ts

@@ -0,0 +1,96 @@
+import {
+  Scope,
+  getCurrentScope,
+  getGlobalScope,
+  getIsolationScope,
+  setCapturedScopesOnSpan,
+  startSpan,
+} from '@sentry/core';
+import { GLOBAL_OBJ } from '@sentry/utils';
+import { init } from '@sentry/vercel-edge';
+import { AsyncLocalStorage } from 'async_hooks';
+
+import { withIsolationScopeOrReuseFromRootSpan } from '../../src/common/utils/withIsolationScopeOrReuseFromRootSpan';
+
+describe('withIsolationScopeOrReuseFromRootSpan', () => {
+  beforeEach(() => {
+    getIsolationScope().clear();
+    getCurrentScope().clear();
+    getGlobalScope().clear();
+    (GLOBAL_OBJ as any).AsyncLocalStorage = AsyncLocalStorage;
+
+    init({
+      enableTracing: true,
+    });
+  });
+
+  it('works without any span', () => {
+    const initialIsolationScope = getIsolationScope();
+    initialIsolationScope.setTag('aa', 'aa');
+
+    withIsolationScopeOrReuseFromRootSpan(isolationScope => {
+      isolationScope.setTag('bb', 'bb');
+      expect(isolationScope).not.toBe(initialIsolationScope);
+      expect(isolationScope.getScopeData().tags).toEqual({ aa: 'aa', bb: 'bb' });
+    });
+  });
+
+  it('works with a non-next.js span', () => {
+    const initialIsolationScope = getIsolationScope();
+    initialIsolationScope.setTag('aa', 'aa');
+
+    const customScope = new Scope();
+
+    startSpan({ name: 'other' }, span => {
+      setCapturedScopesOnSpan(span, getCurrentScope(), customScope);
+
+      withIsolationScopeOrReuseFromRootSpan(isolationScope => {
+        isolationScope.setTag('bb', 'bb');
+        expect(isolationScope).not.toBe(initialIsolationScope);
+        expect(isolationScope.getScopeData().tags).toEqual({ aa: 'aa', bb: 'bb' });
+      });
+    });
+  });
+
+  it('works with a next.js span', () => {
+    const initialIsolationScope = getIsolationScope();
+    initialIsolationScope.setTag('aa', 'aa');
+
+    const customScope = new Scope();
+
+    startSpan(
+      {
+        name: 'other',
+        attributes: { 'next.route': 'aha' },
+      },
+      span => {
+        setCapturedScopesOnSpan(span, getCurrentScope(), customScope);
+
+        withIsolationScopeOrReuseFromRootSpan(isolationScope => {
+          isolationScope.setTag('bb', 'bb');
+          expect(isolationScope).toBe(customScope);
+          expect(isolationScope.getScopeData().tags).toEqual({ bb: 'bb' });
+        });
+      },
+    );
+  });
+
+  it('works with a next.js span that has default isolation scope', () => {
+    const initialIsolationScope = getIsolationScope();
+    initialIsolationScope.setTag('aa', 'aa');
+
+    startSpan(
+      {
+        name: 'other',
+        attributes: { 'next.route': 'aha' },
+      },
+      () => {
+        withIsolationScopeOrReuseFromRootSpan(isolationScope => {
+          isolationScope.setTag('bb', 'bb');
+          expect(isolationScope).not.toBe(initialIsolationScope);
+          expect(isolationScope.getScopeData().tags).toEqual({ aa: 'aa', bb: 'bb' });
+        });
+      },
+    );
+  });
+});