ref(nextjs): Set propagation context for tracing

AbhiPrasad · AbhiPrasad · commit 6bc752722a4e · 2023-06-29T12:26:47.000-04:00
diff --git a/packages/nextjs/src/client/performance.ts b/packages/nextjs/src/client/performance.ts
@@ -1,12 +1,14 @@
 import { getCurrentHub } from '@sentry/core';
 import { WINDOW } from '@sentry/react';
-import type { Primitive, TraceparentData, Transaction, TransactionContext, TransactionSource } from '@sentry/types';
-import {
-  baggageHeaderToDynamicSamplingContext,
-  extractTraceparentData,
-  logger,
-  stripUrlQueryAndFragment,
-} from '@sentry/utils';
+import type {
+  DynamicSamplingContext,
+  Primitive,
+  TraceparentData,
+  Transaction,
+  TransactionContext,
+  TransactionSource,
+} from '@sentry/types';
+import { logger, stripUrlQueryAndFragment, tracingContextFromHeaders } from '@sentry/utils';
 import type { NEXT_DATA as NextData } from 'next/dist/next-server/lib/utils';
 import { default as Router } from 'next/router';
 import type { ParsedUrlQuery } from 'querystring';
@@ -38,7 +40,7 @@ interface SentryEnhancedNextData extends NextData {
 interface NextDataTagInfo {
   route?: string;
   traceParentData?: TraceparentData;
-  baggage?: string;
+  dynamicSamplingContext?: Partial<DynamicSamplingContext>;
   params?: ParsedUrlQuery;
 }
 
@@ -83,13 +85,23 @@ function extractNextDataTagInformation(): NextDataTagInfo {
   nextDataTagInfo.params = query;
 
   if (props && props.pageProps) {
-    if (props.pageProps._sentryBaggage) {
-      nextDataTagInfo.baggage = props.pageProps._sentryBaggage;
+    const sentryTrace = props.pageProps._sentryTraceData;
+    const baggage = props.pageProps._sentryBaggage;
+
+    const { traceparentData, dynamicSamplingContext, propagationContext } = tracingContextFromHeaders(
+      sentryTrace,
+      baggage,
+    );
+
+    if (dynamicSamplingContext) {
+      nextDataTagInfo.dynamicSamplingContext = dynamicSamplingContext;
     }
 
-    if (props.pageProps._sentryTraceData) {
-      nextDataTagInfo.traceParentData = extractTraceparentData(props.pageProps._sentryTraceData);
+    if (traceparentData) {
+      nextDataTagInfo.traceParentData = traceparentData;
     }
+
+    getCurrentHub().getScope().setPropagationContext(propagationContext);
   }
 
   return nextDataTagInfo;
@@ -121,13 +133,11 @@ export function nextRouterInstrumentation(
   startTransactionOnPageLoad: boolean = true,
   startTransactionOnLocationChange: boolean = true,
 ): void {
-  const { route, traceParentData, baggage, params } = extractNextDataTagInformation();
+  const { route, traceParentData, dynamicSamplingContext, params } = extractNextDataTagInformation();
   prevLocationName = route || globalObject.location.pathname;
 
   if (startTransactionOnPageLoad) {
     const source = route ? 'route' : 'url';
-    const dynamicSamplingContext = baggageHeaderToDynamicSamplingContext(baggage);
-
     activeTransaction = startTransactionCb({
       name: prevLocationName,
       op: 'pageload',
diff --git a/packages/nextjs/src/edge/utils/edgeWrapperUtils.ts b/packages/nextjs/src/edge/utils/edgeWrapperUtils.ts
@@ -1,12 +1,6 @@
 import { captureException, getCurrentHub, hasTracingEnabled, startTransaction } from '@sentry/core';
 import type { Span } from '@sentry/types';
-import {
-  addExceptionMechanism,
-  baggageHeaderToDynamicSamplingContext,
-  extractTraceparentData,
-  logger,
-  objectify,
-} from '@sentry/utils';
+import { addExceptionMechanism, logger, objectify, tracingContextFromHeaders } from '@sentry/utils';
 
 import type { EdgeRouteHandler } from '../types';
 import { flush } from './flush';
@@ -32,17 +26,17 @@ export function withEdgeWrapping<H extends EdgeRouteHandler>(
           op: options.spanOp,
         });
       } else if (req instanceof Request) {
-        // If there is a trace header set, extract the data from it (parentSpanId, traceId, and sampling decision)
-        let traceparentData;
-
-        const sentryTraceHeader = req.headers.get('sentry-trace');
-        if (sentryTraceHeader) {
-          traceparentData = extractTraceparentData(sentryTraceHeader);
-          __DEBUG_BUILD__ && logger.log(`[Tracing] Continuing trace ${traceparentData?.traceId}.`);
+        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 dynamicSamplingContext = baggageHeaderToDynamicSamplingContext(req.headers.get('baggage'));
-
         span = startTransaction({
           name: options.spanDescription,
           op: options.spanOp,
diff --git a/packages/nextjs/src/server/utils/wrapperUtils.ts b/packages/nextjs/src/server/utils/wrapperUtils.ts
@@ -6,7 +6,7 @@ import {
   startTransaction,
 } from '@sentry/core';
 import type { Span, Transaction } from '@sentry/types';
-import { baggageHeaderToDynamicSamplingContext, extractTraceparentData } from '@sentry/utils';
+import { isString, tracingContextFromHeaders } from '@sentry/utils';
 import type { IncomingMessage, ServerResponse } from 'http';
 
 import { platformSupportsStreaming } from './platformSupportsStreaming';
@@ -38,6 +38,7 @@ function setTransactionOnRequest(transaction: Transaction, req: IncomingMessage)
  *
  * Note: This function turns the wrapped function into an asynchronous one.
  */
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
 export function withErrorInstrumentation<F extends (...args: any[]) => any>(
   origFunction: F,
 ): (...params: Parameters<F>) => Promise<ReturnType<F>> {
@@ -66,6 +67,7 @@ export function withErrorInstrumentation<F extends (...args: any[]) => any>(
  * @param options Options providing details for the created transaction and span.
  * @returns what the data fetching method call returned.
  */
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
 export function withTracedServerSideDataFetcher<F extends (...args: any[]) => Promise<any> | any>(
   origDataFetcher: F,
   req: IncomingMessage,
@@ -86,12 +88,14 @@ export function withTracedServerSideDataFetcher<F extends (...args: any[]) => Pr
       const previousSpan: Span | undefined = getTransactionFromRequest(req) ?? scope.getSpan();
       let dataFetcherSpan;
 
-      const sentryTraceHeader = req.headers['sentry-trace'];
-      const rawBaggageString = req.headers && req.headers.baggage;
-      const traceparentData =
-        typeof sentryTraceHeader === 'string' ? extractTraceparentData(sentryTraceHeader) : undefined;
-
-      const dynamicSamplingContext = baggageHeaderToDynamicSamplingContext(rawBaggageString);
+      const sentryTrace =
+        req.headers && isString(req.headers['sentry-trace']) ? req.headers['sentry-trace'] : undefined;
+      const baggage = req.headers?.baggage;
+      const { traceparentData, dynamicSamplingContext, propagationContext } = tracingContextFromHeaders(
+        sentryTrace,
+        baggage,
+      );
+      scope.setPropagationContext(propagationContext);
 
       if (platformSupportsStreaming()) {
         let spanToContinue: Span;
diff --git a/packages/nextjs/src/server/wrapApiHandlerWithSentry.ts b/packages/nextjs/src/server/wrapApiHandlerWithSentry.ts
@@ -3,12 +3,11 @@ import { captureException, startTransaction } from '@sentry/node';
 import type { Transaction } from '@sentry/types';
 import {
   addExceptionMechanism,
-  baggageHeaderToDynamicSamplingContext,
-  extractTraceparentData,
   isString,
   logger,
   objectify,
   stripUrlQueryAndFragment,
+  tracingContextFromHeaders,
 } from '@sentry/utils';
 
 import type { AugmentedNextApiRequest, AugmentedNextApiResponse, NextApiHandler } from './types';
@@ -73,15 +72,18 @@ export function withSentry(apiHandler: NextApiHandler, parameterizedRoute?: stri
             currentScope.setSDKProcessingMetadata({ request: req });
 
             if (hasTracingEnabled(options) && options?.instrumenter === 'sentry') {
-              // If there is a trace header set, extract the data from it (parentSpanId, traceId, and sampling decision)
-              let traceparentData;
-              if (req.headers && isString(req.headers['sentry-trace'])) {
-                traceparentData = extractTraceparentData(req.headers['sentry-trace']);
-                __DEBUG_BUILD__ && logger.log(`[Tracing] Continuing trace ${traceparentData?.traceId}.`);
-              }
+              const sentryTrace =
+                req.headers && isString(req.headers['sentry-trace']) ? req.headers['sentry-trace'] : undefined;
+              const baggage = req.headers?.baggage;
+              const { traceparentData, dynamicSamplingContext, propagationContext } = tracingContextFromHeaders(
+                sentryTrace,
+                baggage,
+              );
+              hub.getScope().setPropagationContext(propagationContext);
 
-              const baggageHeader = req.headers && req.headers.baggage;
-              const dynamicSamplingContext = baggageHeaderToDynamicSamplingContext(baggageHeader);
+              if (__DEBUG_BUILD__ && traceparentData) {
+                logger.log(`[Tracing] Continuing trace ${traceparentData.traceId}.`);
+              }
 
               // prefer the parameterized route, if we have it (which we will if we've auto-wrapped the route handler)
               let reqPath = parameterizedRoute;
diff --git a/packages/nextjs/src/server/wrapServerComponentWithSentry.ts b/packages/nextjs/src/server/wrapServerComponentWithSentry.ts
@@ -5,14 +5,15 @@ import {
   runWithAsyncContext,
   startTransaction,
 } from '@sentry/core';
-import { baggageHeaderToDynamicSamplingContext, extractTraceparentData } from '@sentry/utils';
+import { tracingContextFromHeaders } from '@sentry/utils';
 
 import { isNotFoundNavigationError, isRedirectNavigationError } from '../common/nextNavigationErrorUtils';
 import type { ServerComponentContext } from '../common/types';
 
 /**
  * Wraps an `app` directory server component with Sentry error instrumentation.
  */
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
 export function wrapServerComponentWithSentry<F extends (...args: any[]) => any>(
   appDirComponent: F,
   context: ServerComponentContext,
@@ -28,14 +29,15 @@ export function wrapServerComponentWithSentry<F extends (...args: any[]) => any>
     apply: (originalFunction, thisArg, args) => {
       return runWithAsyncContext(() => {
         const hub = getCurrentHub();
+        const currentScope = hub.getScope();
 
         let maybePromiseResult;
 
-        const traceparentData = context.sentryTraceHeader
-          ? extractTraceparentData(context.sentryTraceHeader)
-          : undefined;
-
-        const dynamicSamplingContext = baggageHeaderToDynamicSamplingContext(context.baggageHeader);
+        const { traceparentData, dynamicSamplingContext, propagationContext } = tracingContextFromHeaders(
+          context.sentryTraceHeader,
+          context.baggageHeader,
+        );
+        currentScope.setPropagationContext(propagationContext);
 
         const transaction = startTransaction({
           op: 'function.nextjs',
@@ -48,7 +50,6 @@ export function wrapServerComponentWithSentry<F extends (...args: any[]) => any>
           },
         });
 
-        const currentScope = hub.getScope();
         if (currentScope) {
           currentScope.setSpan(transaction);
         }
