ref(serverless): Extract propagation context (#8429)

getsentry · Jul 4, 2023 · 02e5035 · 02e5035
1 parent a3178f7
commit 02e5035
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 30 deletions.
diff --git a/packages/serverless/src/awslambda.ts b/packages/serverless/src/awslambda.ts
@@ -3,7 +3,7 @@ import type { Scope } from '@sentry/node';
 import * as Sentry from '@sentry/node';
 import { captureException, captureMessage, flush, getCurrentHub, withScope } from '@sentry/node';
 import type { Integration } from '@sentry/types';
-import { baggageHeaderToDynamicSamplingContext, extractTraceparentData, isString, logger } from '@sentry/utils';
+import { isString, logger, tracingContextFromHeaders } from '@sentry/utils';
 // NOTE: I have no idea how to fix this right now, and don't want to waste more time, as it builds just fine — Kamil
 // eslint-disable-next-line import/no-unresolved
 import type { Context, Handler } from 'aws-lambda';
@@ -274,17 +274,20 @@ export function wrapHandler<TEvent, TResult>(
       }, timeoutWarningDelay) as unknown as NodeJS.Timeout;
     }
 
-    // Applying `sentry-trace` to context
-    let traceparentData;
-    const eventWithHeaders = event as { headers?: { [key: string]: string } };
-    if (eventWithHeaders.headers && isString(eventWithHeaders.headers['sentry-trace'])) {
-      traceparentData = extractTraceparentData(eventWithHeaders.headers['sentry-trace']);
-    }
+    const hub = getCurrentHub();
 
-    const baggageHeader = eventWithHeaders.headers && eventWithHeaders.headers.baggage;
-    const dynamicSamplingContext = baggageHeaderToDynamicSamplingContext(baggageHeader);
+    const eventWithHeaders = event as { headers?: { [key: string]: string } };
 
-    const hub = getCurrentHub();
+    const sentryTrace =
+      eventWithHeaders.headers && isString(eventWithHeaders.headers['sentry-trace'])
+        ? eventWithHeaders.headers['sentry-trace']
+        : undefined;
+    const baggage = eventWithHeaders.headers?.baggage;
+    const { traceparentData, dynamicSamplingContext, propagationContext } = tracingContextFromHeaders(
+      sentryTrace,
+      baggage,
+    );
+    hub.getScope().setPropagationContext(propagationContext);
 
     const transaction = hub.startTransaction({
       name: context.functionName,

diff --git a/packages/serverless/src/gcpfunction/http.ts b/packages/serverless/src/gcpfunction/http.ts
@@ -1,13 +1,6 @@
 import type { AddRequestDataToEventOptions } from '@sentry/node';
 import { captureException, flush, getCurrentHub } from '@sentry/node';
-import {
-  baggageHeaderToDynamicSamplingContext,
-  extractTraceparentData,
-  isString,
-  isThenable,
-  logger,
-  stripUrlQueryAndFragment,
-} from '@sentry/utils';
+import { isString, isThenable, logger, stripUrlQueryAndFragment, tracingContextFromHeaders } from '@sentry/utils';
 
 import { domainify, proxyFunction } from './../utils';
 import type { HttpFunction, WrapperOptions } from './general';
@@ -68,21 +61,18 @@ function _wrapHttpFunction(fn: HttpFunction, wrapOptions: Partial<HttpFunctionWr
     ...wrapOptions,
   };
   return (req, res) => {
+    const hub = getCurrentHub();
+
     const reqMethod = (req.method || '').toUpperCase();
     const reqUrl = stripUrlQueryAndFragment(req.originalUrl || req.url || '');
 
-    // Applying `sentry-trace` to context
-    let traceparentData;
-    const reqWithHeaders = req as { headers?: { [key: string]: string } };
-    if (reqWithHeaders.headers && isString(reqWithHeaders.headers['sentry-trace'])) {
-      traceparentData = extractTraceparentData(reqWithHeaders.headers['sentry-trace']);
-    }
-
-    const baggageHeader = reqWithHeaders.headers?.baggage;
-
-    const dynamicSamplingContext = baggageHeaderToDynamicSamplingContext(baggageHeader);
-
-    const hub = getCurrentHub();
+    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 transaction = hub.startTransaction({
       name: `${reqMethod} ${reqUrl}`,

diff --git a/packages/serverless/test/__mocks__/@sentry/node.ts b/packages/serverless/test/__mocks__/@sentry/node.ts
@@ -22,6 +22,7 @@ export const fakeScope = {
   setSpan: jest.fn(),
   getTransaction: jest.fn(() => fakeTransaction),
   setSDKProcessingMetadata: jest.fn(),
+  setPropagationContext: jest.fn(),
 };
 export const fakeSpan = {
   finish: jest.fn(),