ref(nextjs): Simplify wrapServerComponentWithSentry (#9844)

Author: Luca Forstner

packages/core/src/tracing/trace.ts

@@ -24,7 +24,9 @@ export function trace<T>(
   context: TransactionContext,
   callback: (span?: Span) => T,
   // eslint-disable-next-line @typescript-eslint/no-empty-function
-  onError: (error: unknown) => void = () => {},
+  onError: (error: unknown, span?: Span) => void = () => {},
+  // eslint-disable-next-line @typescript-eslint/no-empty-function
+  afterFinish: () => void = () => {},
 ): T {
   const ctx = normalizeContext(context);
 
@@ -46,24 +48,28 @@ export function trace<T>(
     maybePromiseResult = callback(activeSpan);
   } catch (e) {
     activeSpan && activeSpan.setStatus('internal_error');
-    onError(e);
+    onError(e, activeSpan);
     finishAndSetSpan();
+    afterFinish();
     throw e;
   }
 
   if (isThenable(maybePromiseResult)) {
     Promise.resolve(maybePromiseResult).then(
       () => {
         finishAndSetSpan();
+        afterFinish();
       },
       e => {
         activeSpan && activeSpan.setStatus('internal_error');
-        onError(e);
+        onError(e, activeSpan);
         finishAndSetSpan();
+        afterFinish();
       },
     );
   } else {
     finishAndSetSpan();
+    afterFinish();
   }
 
   return maybePromiseResult;

packages/nextjs/src/common/wrapServerComponentWithSentry.ts

@@ -1,11 +1,5 @@
-import {
-  addTracingExtensions,
-  captureException,
-  getCurrentScope,
-  runWithAsyncContext,
-  startTransaction,
-} from '@sentry/core';
-import { tracingContextFromHeaders, winterCGHeadersToDict } from '@sentry/utils';
+import { addTracingExtensions, captureException, continueTrace, runWithAsyncContext, trace } from '@sentry/core';
+import { winterCGHeadersToDict } from '@sentry/utils';
 
 import { isNotFoundNavigationError, isRedirectNavigationError } from '../common/nextNavigationErrorUtils';
 import type { ServerComponentContext } from '../common/types';
@@ -28,88 +22,57 @@ export function wrapServerComponentWithSentry<F extends (...args: any[]) => any>
   return new Proxy(appDirComponent, {
     apply: (originalFunction, thisArg, args) => {
       return runWithAsyncContext(() => {
-        const currentScope = getCurrentScope();
-        let maybePromiseResult;
-
         const completeHeadersDict: Record<string, string> = context.headers
           ? winterCGHeadersToDict(context.headers)
           : {};
 
-        const { traceparentData, dynamicSamplingContext, propagationContext } = tracingContextFromHeaders(
+        const transactionContext = continueTrace({
           // eslint-disable-next-line deprecation/deprecation
-          context.sentryTraceHeader ?? completeHeadersDict['sentry-trace'],
+          sentryTrace: context.sentryTraceHeader ?? completeHeadersDict['sentry-trace'],
           // eslint-disable-next-line deprecation/deprecation
-          context.baggageHeader ?? completeHeadersDict['baggage'],
-        );
-        currentScope.setPropagationContext(propagationContext);
-
-        const transaction = startTransaction({
-          op: 'function.nextjs',
-          name: `${componentType} Server Component (${componentRoute})`,
-          status: 'ok',
-          origin: 'auto.function.nextjs',
-          ...traceparentData,
-          metadata: {
-            request: {
-              headers: completeHeadersDict,
-            },
-            source: 'component',
-            dynamicSamplingContext: traceparentData && !dynamicSamplingContext ? {} : dynamicSamplingContext,
-          },
+          baggage: context.baggageHeader ?? completeHeadersDict['baggage'],
         });
 
-        currentScope.setSpan(transaction);
-
-        const handleErrorCase = (e: unknown): void => {
-          if (isNotFoundNavigationError(e)) {
-            // We don't want to report "not-found"s
-            transaction.setStatus('not_found');
-          } else if (isRedirectNavigationError(e)) {
-            // We don't want to report redirects
-          } else {
-            transaction.setStatus('internal_error');
-
-            captureException(e, {
-              mechanism: {
-                handled: false,
+        const res = trace(
+          {
+            ...transactionContext,
+            op: 'function.nextjs',
+            name: `${componentType} Server Component (${componentRoute})`,
+            status: 'ok',
+            origin: 'auto.function.nextjs',
+            metadata: {
+              ...transactionContext.metadata,
+              request: {
+                headers: completeHeadersDict,
               },
-            });
-          }
-
-          transaction.finish();
-        };
-
-        try {
-          maybePromiseResult = originalFunction.apply(thisArg, args);
-        } catch (e) {
-          handleErrorCase(e);
-          void flushQueue();
-          throw e;
-        }
+              source: 'component',
+            },
+          },
+          () => originalFunction.apply(thisArg, args),
+          (e, span) => {
+            if (isNotFoundNavigationError(e)) {
+              // We don't want to report "not-found"s
+              span?.setStatus('not_found');
+            } else if (isRedirectNavigationError(e)) {
+              // We don't want to report redirects
+              // Since `trace` will automatically set the span status to "internal_error" we need to set it back to "ok"
+              span?.setStatus('ok');
+            } else {
+              span?.setStatus('internal_error');
 
-        if (typeof maybePromiseResult === 'object' && maybePromiseResult !== null && 'then' in maybePromiseResult) {
-          // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access
-          Promise.resolve(maybePromiseResult)
-            .then(
-              () => {
-                transaction.finish();
-              },
-              e => {
-                handleErrorCase(e);
-              },
-            )
-            .finally(() => {
-              void flushQueue();
-            });
+              captureException(e, {
+                mechanism: {
+                  handled: false,
+                },
+              });
+            }
+          },
+          () => {
+            void flushQueue();
+          },
+        );
 
-          // It is very important that we return the original promise here, because Next.js attaches various properties
-          // to that promise and will throw if they are not on the returned value.
-          return maybePromiseResult;
-        } else {
-          transaction.finish();
-          void flushQueue();
-          return maybePromiseResult;
-        }
+        return res;
       });
     },
   });