ref: Remove startTransaction call from withEdgeWrapping (#9984)

lforst · web-flow · commit ea74266f4967 · 2023-12-28T10:06:54.000+01:00
diff --git a/packages/nextjs/src/common/utils/edgeWrapperUtils.ts b/packages/nextjs/src/common/utils/edgeWrapperUtils.ts
@@ -1,15 +1,7 @@
-import { addTracingExtensions, captureException, getCurrentScope, startTransaction } from '@sentry/core';
-import type { Span } from '@sentry/types';
-import {
-  addExceptionMechanism,
-  logger,
-  objectify,
-  tracingContextFromHeaders,
-  winterCGRequestToRequestData,
-} from '@sentry/utils';
+import { addTracingExtensions, captureException, continueTrace, trace } from '@sentry/core';
+import { winterCGRequestToRequestData } from '@sentry/utils';
 
 import type { EdgeRouteHandler } from '../../edge/types';
-import { DEBUG_BUILD } from '../debug-build';
 import { flushQueue } from './responseEnd';
 
 /**
@@ -21,84 +13,56 @@ export function withEdgeWrapping<H extends EdgeRouteHandler>(
 ): (...params: Parameters<H>) => Promise<ReturnType<H>> {
   return async function (this: unknown, ...args) {
     addTracingExtensions();
+    const req: unknown = args[0];
 
-    const req = args[0];
-    const currentScope = getCurrentScope();
-    const prevSpan = currentScope.getSpan();
+    let sentryTrace;
+    let baggage;
 
-    let span: Span | undefined;
+    if (req instanceof Request) {
+      sentryTrace = req.headers.get('sentry-trace') || '';
+      baggage = req.headers.get('baggage');
+    }
 
-    if (prevSpan) {
-      span = prevSpan.startChild({
-        description: options.spanDescription,
-        op: options.spanOp,
-        origin: 'auto.function.nextjs',
-      });
-    } else if (req instanceof Request) {
-      const sentryTrace = req.headers.get('sentry-trace') || '';
-      const baggage = req.headers.get('baggage');
-      const { traceparentData, dynamicSamplingContext, propagationContext } = tracingContextFromHeaders(
-        sentryTrace,
-        baggage,
-      );
-      currentScope.setPropagationContext(propagationContext);
-      if (traceparentData) {
-        DEBUG_BUILD && logger.log(`[Tracing] Continuing trace ${traceparentData.traceId}.`);
-      }
+    const transactionContext = continueTrace({
+      sentryTrace,
+      baggage,
+    });
 
-      span = startTransaction({
+    return trace(
+      {
+        ...transactionContext,
         name: options.spanDescription,
         op: options.spanOp,
-        origin: 'auto.ui.nextjs.withEdgeWrapping',
-        ...traceparentData,
+        origin: 'auto.function.nextjs.withEdgeWrapping',
         metadata: {
-          request: winterCGRequestToRequestData(req),
-          dynamicSamplingContext: traceparentData && !dynamicSamplingContext ? {} : dynamicSamplingContext,
+          ...transactionContext.metadata,
+          request: req instanceof Request ? winterCGRequestToRequestData(req) : undefined,
           source: 'route',
         },
-      });
-    }
-
-    currentScope?.setSpan(span);
-
-    try {
-      const handlerResult: ReturnType<H> = await handler.apply(this, args);
+      },
+      async span => {
+        const handlerResult = await handler.apply(this, args);
 
-      if ((handlerResult as unknown) instanceof Response) {
-        span?.setHttpStatus(handlerResult.status);
-      } else {
-        span?.setStatus('ok');
-      }
+        if (handlerResult instanceof Response) {
+          span?.setHttpStatus(handlerResult.status);
+        } else {
+          span?.setStatus('ok');
+        }
 
-      return handlerResult;
-    } catch (e) {
-      // In case we have a primitive, wrap it in the equivalent wrapper class (string -> String, etc.) so that we can
-      // store a seen flag on it.
-      const objectifiedErr = objectify(e);
-
-      span?.setStatus('internal_error');
-
-      captureException(objectifiedErr, scope => {
-        scope.setSpan(span);
-        scope.addEventProcessor(event => {
-          addExceptionMechanism(event, {
+        return handlerResult;
+      },
+      (err, span) => {
+        span?.setStatus('internal_error');
+        captureException(err, {
+          mechanism: {
             type: 'instrument',
             handled: false,
             data: {
               function: options.mechanismFunctionName,
             },
-          });
-          return event;
+          },
         });
-
-        return scope;
-      });
-
-      throw objectifiedErr;
-    } finally {
-      span?.end();
-      currentScope?.setSpan(prevSpan);
-      await flushQueue();
-    }
+      },
+    ).finally(() => flushQueue());
   };
 }
diff --git a/packages/nextjs/test/edge/edgeWrapperUtils.test.ts b/packages/nextjs/test/edge/edgeWrapperUtils.test.ts
@@ -1,12 +1,7 @@
 import * as coreSdk from '@sentry/core';
-import { addTracingExtensions } from '@sentry/core';
 
 import { withEdgeWrapping } from '../../src/common/utils/edgeWrapperUtils';
 
-// The wrap* functions require the hub to have tracing extensions. This is normally called by the EdgeClient
-// constructor but the client isn't used in these tests.
-addTracingExtensions();
-
 // @ts-expect-error Request does not exist on type Global
 const origRequest = global.Request;
 // @ts-expect-error Response does not exist on type Global
@@ -33,8 +28,6 @@ afterAll(() => {
 
 beforeEach(() => {
   jest.clearAllMocks();
-  jest.resetAllMocks();
-  jest.spyOn(coreSdk, 'hasTracingEnabled').mockImplementation(() => true);
 });
 
 describe('withEdgeWrapping', () => {
@@ -71,27 +64,30 @@ describe('withEdgeWrapping', () => {
     expect(captureExceptionSpy).toHaveBeenCalledTimes(1);
   });
 
-  it('should return a function that starts a transaction when a request object is passed', async () => {
-    const startTransactionSpy = jest.spyOn(coreSdk, 'startTransaction');
+  it('should return a function that calls trace', async () => {
+    const traceSpy = jest.spyOn(coreSdk, 'trace');
 
-    const origFunctionReturnValue = new Response();
-    const origFunction = jest.fn(_req => origFunctionReturnValue);
+    const request = new Request('https://sentry.io/');
+    const origFunction = jest.fn(_req => new Response());
 
     const wrappedFunction = withEdgeWrapping(origFunction, {
       spanDescription: 'some label',
       mechanismFunctionName: 'some name',
       spanOp: 'some op',
     });
 
-    const request = new Request('https://sentry.io/');
     await wrappedFunction(request);
-    expect(startTransactionSpy).toHaveBeenCalledTimes(1);
-    expect(startTransactionSpy).toHaveBeenCalledWith(
+
+    expect(traceSpy).toHaveBeenCalledTimes(1);
+    expect(traceSpy).toHaveBeenCalledWith(
       expect.objectContaining({
-        metadata: expect.objectContaining({ source: 'route' }),
+        metadata: { request: { headers: {} }, source: 'route' },
         name: 'some label',
         op: 'some op',
+        origin: 'auto.function.nextjs.withEdgeWrapping',
       }),
+      expect.any(Function),
+      expect.any(Function),
     );
   });
 
diff --git a/packages/nextjs/test/edge/withSentryAPI.test.ts b/packages/nextjs/test/edge/withSentryAPI.test.ts
@@ -2,28 +2,30 @@ import * as coreSdk from '@sentry/core';
 
 import { wrapApiHandlerWithSentry } from '../../src/edge';
 
-// The wrap* functions require the hub to have tracing extensions. This is normally called by the EdgeClient
-// constructor but the client isn't used in these tests.
-coreSdk.addTracingExtensions();
-
 // @ts-expect-error Request does not exist on type Global
 const origRequest = global.Request;
 // @ts-expect-error Response does not exist on type Global
 const origResponse = global.Response;
 
 // @ts-expect-error Request does not exist on type Global
 global.Request = class Request {
-  headers = {
+  public url: string;
+
+  public headers = {
     get() {
       return null;
     },
   };
 
-  method = 'POST';
+  public method = 'POST';
+
+  public constructor(input: string) {
+    this.url = input;
+  }
 };
 
 // @ts-expect-error Response does not exist on type Global
-global.Response = class Request {};
+global.Response = class Response {};
 
 afterAll(() => {
   // @ts-expect-error Request does not exist on type Global
@@ -32,66 +34,51 @@ afterAll(() => {
   global.Response = origResponse;
 });
 
-beforeEach(() => {
-  jest.resetAllMocks();
-  jest.restoreAllMocks();
-  jest.spyOn(coreSdk, 'hasTracingEnabled').mockImplementation(() => true);
+const traceSpy = jest.spyOn(coreSdk, 'trace');
+
+afterEach(() => {
+  jest.clearAllMocks();
 });
 
 describe('wrapApiHandlerWithSentry', () => {
-  it('should return a function that starts a transaction with the correct name when there is no active transaction and a request is being passed', async () => {
-    const startTransactionSpy = jest.spyOn(coreSdk, 'startTransaction');
-
-    const origFunctionReturnValue = new Response();
-    const origFunction = jest.fn(_req => origFunctionReturnValue);
+  it('should return a function that calls trace', async () => {
+    const request = new Request('https://sentry.io/');
+    const origFunction = jest.fn(_req => new Response());
 
     const wrappedFunction = wrapApiHandlerWithSentry(origFunction, '/user/[userId]/post/[postId]');
 
-    const request = new Request('https://sentry.io/');
     await wrappedFunction(request);
-    expect(startTransactionSpy).toHaveBeenCalledTimes(1);
-    expect(startTransactionSpy).toHaveBeenCalledWith(
+
+    expect(traceSpy).toHaveBeenCalledTimes(1);
+    expect(traceSpy).toHaveBeenCalledWith(
       expect.objectContaining({
-        metadata: expect.objectContaining({ source: 'route' }),
+        metadata: { request: { headers: {}, method: 'POST', url: 'https://sentry.io/' }, source: 'route' },
         name: 'POST /user/[userId]/post/[postId]',
         op: 'http.server',
+        origin: 'auto.function.nextjs.withEdgeWrapping',
       }),
+      expect.any(Function),
+      expect.any(Function),
     );
   });
 
-  it('should return a function that should not start a transaction when there is no active span and no request is being passed', async () => {
-    const startTransactionSpy = jest.spyOn(coreSdk, 'startTransaction');
-
-    const origFunctionReturnValue = new Response();
-    const origFunction = jest.fn(() => origFunctionReturnValue);
+  it('should return a function that calls trace without throwing when no request is passed', async () => {
+    const origFunction = jest.fn(() => new Response());
 
     const wrappedFunction = wrapApiHandlerWithSentry(origFunction, '/user/[userId]/post/[postId]');
 
     await wrappedFunction();
-    expect(startTransactionSpy).not.toHaveBeenCalled();
-  });
-
-  it('should return a function that starts a span on the current transaction with the correct description when there is an active transaction and no request is being passed', async () => {
-    const testTransaction = coreSdk.startTransaction({ name: 'testTransaction' });
-    coreSdk.getCurrentHub().getScope().setSpan(testTransaction);
-
-    const startChildSpy = jest.spyOn(testTransaction, 'startChild');
 
-    const origFunctionReturnValue = new Response();
-    const origFunction = jest.fn(() => origFunctionReturnValue);
-
-    const wrappedFunction = wrapApiHandlerWithSentry(origFunction, '/user/[userId]/post/[postId]');
-
-    await wrappedFunction();
-    expect(startChildSpy).toHaveBeenCalledTimes(1);
-    expect(startChildSpy).toHaveBeenCalledWith(
+    expect(traceSpy).toHaveBeenCalledTimes(1);
+    expect(traceSpy).toHaveBeenCalledWith(
       expect.objectContaining({
-        description: 'handler (/user/[userId]/post/[postId])',
-        op: 'function',
+        metadata: { source: 'route' },
+        name: 'handler (/user/[userId]/post/[postId])',
+        op: 'http.server',
+        origin: 'auto.function.nextjs.withEdgeWrapping',
       }),
+      expect.any(Function),
+      expect.any(Function),
     );
-
-    testTransaction.end();
-    coreSdk.getCurrentHub().getScope().setSpan(undefined);
   });
 });
diff --git a/packages/utils/src/requestdata.ts b/packages/utils/src/requestdata.ts
@@ -374,6 +374,7 @@ function extractQueryParams(
  * Transforms a `Headers` object that implements the `Web Fetch API` (https://developer.mozilla.org/en-US/docs/Web/API/Headers) into a simple key-value dict.
  * The header keys will be lower case: e.g. A "Content-Type" header will be stored as "content-type".
  */
+// TODO(v8): Make this function return undefined when the extraction fails.
 export function winterCGHeadersToDict(winterCGHeaders: WebFetchHeaders): Record<string, string> {
   const headers: Record<string, string> = {};
   try {
