fix(nextjs): Escape Next.js' OpenTelemetry instrumentation (#11625)

Author: lforst

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

@@ -4,15 +4,16 @@ import {
   SPAN_STATUS_OK,
   captureException,
   continueTrace,
-  getIsolationScope,
   handleCallbackErrors,
   setHttpStatus,
   startSpan,
+  withIsolationScope,
 } from '@sentry/core';
 import { winterCGRequestToRequestData } from '@sentry/utils';
 
 import type { EdgeRouteHandler } from '../../edge/types';
 import { flushQueue } from './responseEnd';
+import { commonObjectToIsolationScope, escapeNextjsTracing } from './tracingUtils';
 
 /**
  * Wraps a function on the edge runtime with error and performance monitoring.
@@ -22,61 +23,65 @@ export function withEdgeWrapping<H extends EdgeRouteHandler>(
   options: { spanDescription: string; spanOp: string; mechanismFunctionName: string },
 ): (...params: Parameters<H>) => Promise<ReturnType<H>> {
   return async function (this: unknown, ...args) {
-    const req: unknown = args[0];
+    return escapeNextjsTracing(() => {
+      const req: unknown = args[0];
+      return withIsolationScope(commonObjectToIsolationScope(req), isolationScope => {
+        let sentryTrace;
+        let baggage;
 
-    let sentryTrace;
-    let baggage;
+        if (req instanceof Request) {
+          sentryTrace = req.headers.get('sentry-trace') || '';
+          baggage = req.headers.get('baggage');
 
-    if (req instanceof Request) {
-      sentryTrace = req.headers.get('sentry-trace') || '';
-      baggage = req.headers.get('baggage');
-    }
+          isolationScope.setSDKProcessingMetadata({
+            request: winterCGRequestToRequestData(req),
+          });
+        }
 
-    return continueTrace(
-      {
-        sentryTrace,
-        baggage,
-      },
-      () => {
-        getIsolationScope().setSDKProcessingMetadata({
-          request: req instanceof Request ? winterCGRequestToRequestData(req) : undefined,
-        });
-        return startSpan(
+        return continueTrace(
           {
-            name: options.spanDescription,
-            op: options.spanOp,
-            forceTransaction: true,
-            attributes: {
-              [SEMANTIC_ATTRIBUTE_SENTRY_SOURCE]: 'route',
-              [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.function.nextjs.withEdgeWrapping',
-            },
+            sentryTrace,
+            baggage,
           },
-          async span => {
-            const handlerResult = await handleCallbackErrors(
-              () => handler.apply(this, args),
-              error => {
-                captureException(error, {
-                  mechanism: {
-                    type: 'instrument',
-                    handled: false,
-                    data: {
-                      function: options.mechanismFunctionName,
-                    },
-                  },
-                });
+          () => {
+            return startSpan(
+              {
+                name: options.spanDescription,
+                op: options.spanOp,
+                forceTransaction: true,
+                attributes: {
+                  [SEMANTIC_ATTRIBUTE_SENTRY_SOURCE]: 'route',
+                  [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.function.nextjs.withEdgeWrapping',
+                },
               },
-            );
+              async span => {
+                const handlerResult = await handleCallbackErrors(
+                  () => handler.apply(this, args),
+                  error => {
+                    captureException(error, {
+                      mechanism: {
+                        type: 'instrument',
+                        handled: false,
+                        data: {
+                          function: options.mechanismFunctionName,
+                        },
+                      },
+                    });
+                  },
+                );
 
-            if (handlerResult instanceof Response) {
-              setHttpStatus(span, handlerResult.status);
-            } else {
-              span.setStatus({ code: SPAN_STATUS_OK });
-            }
+                if (handlerResult instanceof Response) {
+                  setHttpStatus(span, handlerResult.status);
+                } else {
+                  span.setStatus({ code: SPAN_STATUS_OK });
+                }
 
-            return handlerResult;
+                return handlerResult;
+              },
+            ).finally(() => flushQueue());
           },
-        ).finally(() => flushQueue());
-      },
-    );
+        );
+      });
+    });
   };
 }

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

@@ -0,0 +1,98 @@
+import { Scope, getCurrentScope, withActiveSpan } from '@sentry/core';
+import type { PropagationContext } from '@sentry/types';
+import { GLOBAL_OBJ, logger, uuid4 } from '@sentry/utils';
+import { DEBUG_BUILD } from '../debug-build';
+
+const commonPropagationContextMap = new WeakMap<object, PropagationContext>();
+
+/**
+ * Takes a shared (garbage collectable) object between resources, e.g. a headers object shared between Next.js server components and returns a common propagation context.
+ *
+ * @param commonObject The shared object.
+ * @param propagationContext The propagation context that should be shared between all the resources if no propagation context was registered yet.
+ * @returns the shared propagation context.
+ */
+export function commonObjectToPropagationContext(
+  commonObject: unknown,
+  propagationContext: PropagationContext,
+): PropagationContext {
+  if (typeof commonObject === 'object' && commonObject) {
+    const memoPropagationContext = commonPropagationContextMap.get(commonObject);
+    if (memoPropagationContext) {
+      return memoPropagationContext;
+    } else {
+      commonPropagationContextMap.set(commonObject, propagationContext);
+      return propagationContext;
+    }
+  } else {
+    return propagationContext;
+  }
+}
+
+const commonIsolationScopeMap = new WeakMap<object, Scope>();
+
+/**
+ * Takes a shared (garbage collectable) object between resources, e.g. a headers object shared between Next.js server components and returns a common propagation context.
+ *
+ * @param commonObject The shared object.
+ * @param isolationScope The isolationScope that should be shared between all the resources if no isolation scope was created yet.
+ * @returns the shared isolation scope.
+ */
+export function commonObjectToIsolationScope(commonObject: unknown): Scope {
+  if (typeof commonObject === 'object' && commonObject) {
+    const memoIsolationScope = commonIsolationScopeMap.get(commonObject);
+    if (memoIsolationScope) {
+      return memoIsolationScope;
+    } else {
+      const newIsolationScope = new Scope();
+      commonIsolationScopeMap.set(commonObject, newIsolationScope);
+      return newIsolationScope;
+    }
+  } else {
+    return new Scope();
+  }
+}
+
+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;
+}
+
+let nextjsEscapedAsyncStorage: AsyncLocalStorage<true>;
+
+/**
+ * Will mark the execution context of the callback as "escaped" from Next.js internal tracing by unsetting the active
+ * span and propagation context. When an execution passes through this function multiple times, it is a noop after the
+ * first time.
+ */
+export function escapeNextjsTracing<T>(cb: () => T): T {
+  // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-explicit-any
+  const MaybeGlobalAsyncLocalStorage = (GLOBAL_OBJ as any).AsyncLocalStorage;
+
+  if (!MaybeGlobalAsyncLocalStorage) {
+    DEBUG_BUILD &&
+      logger.warn(
+        "Tried to register AsyncLocalStorage async context strategy in a runtime that doesn't support AsyncLocalStorage.",
+      );
+    return cb();
+  }
+
+  if (!nextjsEscapedAsyncStorage) {
+    nextjsEscapedAsyncStorage = new MaybeGlobalAsyncLocalStorage();
+  }
+
+  if (nextjsEscapedAsyncStorage.getStore()) {
+    return cb();
+  } else {
+    return withActiveSpan(null, () => {
+      getCurrentScope().setPropagationContext({
+        traceId: uuid4(),
+        spanId: uuid4().substring(16),
+      });
+      return nextjsEscapedAsyncStorage.run(true, () => {
+        return cb();
+      });
+    });
+  }
+}

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

@@ -10,13 +10,14 @@ 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';
+import { commonObjectToIsolationScope } from './tracingUtils';
 
 declare module 'http' {
   interface IncomingMessage {
@@ -89,7 +90,8 @@ 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 withIsolationScopeOrReuseFromRootSpan(async isolationScope => {
+    const isolationScope = commonObjectToIsolationScope(req);
+    return withIsolationScope(isolationScope, () => {
       isolationScope.setSDKProcessingMetadata({
         request: req,
       });
@@ -100,7 +102,6 @@ export function withTracedServerSideDataFetcher<F extends (...args: any[]) => Pr
 
       return continueTrace({ sentryTrace, baggage }, () => {
         const requestSpan = getOrStartRequestSpan(req, res, options.requestedRouteName);
-
         return withActiveSpan(requestSpan, () => {
           return startSpanManual(
             {