From b16a232aeb73eb2456fa5959d6653ba1632a28ae Mon Sep 17 00:00:00 2001
From: Abhijeet Prasad <aprasad@sentry.io>
Date: Thu, 29 Jun 2023 12:56:28 -0400
Subject: [PATCH] ref(sveltekit): Extract propagation context

---
 packages/sveltekit/src/server/handle.ts |  3 ++-
 packages/sveltekit/src/server/load.ts   |  5 +++--
 packages/sveltekit/src/server/utils.ts  | 25 +++++++------------------
 3 files changed, 12 insertions(+), 21 deletions(-)

diff --git a/packages/sveltekit/src/server/handle.ts b/packages/sveltekit/src/server/handle.ts
index 6d41317cc044..5c90f42f4c9e 100644
--- a/packages/sveltekit/src/server/handle.ts
+++ b/packages/sveltekit/src/server/handle.ts
@@ -112,7 +112,8 @@ function instrumentHandle({ event, resolve }: Parameters<Handle>[0], options: Se
     return resolve(event);
   }
 
-  const { traceparentData, dynamicSamplingContext } = getTracePropagationData(event);
+  const { dynamicSamplingContext, traceparentData, propagationContext } = getTracePropagationData(event);
+  getCurrentHub().getScope().setPropagationContext(propagationContext);
 
   return trace(
     {
diff --git a/packages/sveltekit/src/server/load.ts b/packages/sveltekit/src/server/load.ts
index c776dc639d3a..0ad28e1cb4eb 100644
--- a/packages/sveltekit/src/server/load.ts
+++ b/packages/sveltekit/src/server/load.ts
@@ -1,5 +1,5 @@
 /* eslint-disable @sentry-internal/sdk/no-optional-chaining */
-import { trace } from '@sentry/core';
+import { getCurrentHub, trace } from '@sentry/core';
 import { captureException } from '@sentry/node';
 import type { TransactionContext } from '@sentry/types';
 import { addExceptionMechanism, addNonEnumerableProperty, objectify } from '@sentry/utils';
@@ -121,7 +121,8 @@ export function wrapServerLoadWithSentry<T extends (...args: any) => any>(origSe
 
       const routeId = event.route && event.route.id;
 
-      const { dynamicSamplingContext, traceparentData } = getTracePropagationData(event);
+      const { dynamicSamplingContext, traceparentData, propagationContext } = getTracePropagationData(event);
+      getCurrentHub().getScope().setPropagationContext(propagationContext);
 
       const traceLoadContext: TransactionContext = {
         op: 'function.sveltekit.server.load',
diff --git a/packages/sveltekit/src/server/utils.ts b/packages/sveltekit/src/server/utils.ts
index 8ef6acced314..1a9e1781643c 100644
--- a/packages/sveltekit/src/server/utils.ts
+++ b/packages/sveltekit/src/server/utils.ts
@@ -1,12 +1,5 @@
-import type { DynamicSamplingContext, StackFrame, TraceparentData } from '@sentry/types';
-import {
-  baggageHeaderToDynamicSamplingContext,
-  basename,
-  escapeStringForRegex,
-  extractTraceparentData,
-  GLOBAL_OBJ,
-  join,
-} from '@sentry/utils';
+import type { StackFrame } from '@sentry/types';
+import { basename, escapeStringForRegex, GLOBAL_OBJ, join, tracingContextFromHeaders } from '@sentry/utils';
 import type { RequestEvent } from '@sveltejs/kit';
 
 import { WRAPPED_MODULE_SUFFIX } from '../vite/autoInstrument';
@@ -15,17 +8,13 @@ import type { GlobalWithSentryValues } from '../vite/injectGlobalValues';
 /**
  * Takes a request event and extracts traceparent and DSC data
  * from the `sentry-trace` and `baggage` DSC headers.
+ *
+ * Sets propagation context as a side effect.
  */
-export function getTracePropagationData(event: RequestEvent): {
-  traceparentData?: TraceparentData;
-  dynamicSamplingContext?: Partial<DynamicSamplingContext>;
-} {
-  const sentryTraceHeader = event.request.headers.get('sentry-trace');
+export function getTracePropagationData(event: RequestEvent): ReturnType<typeof tracingContextFromHeaders> {
+  const sentryTraceHeader = event.request.headers.get('sentry-trace') || '';
   const baggageHeader = event.request.headers.get('baggage');
-  const traceparentData = sentryTraceHeader ? extractTraceparentData(sentryTraceHeader) : undefined;
-  const dynamicSamplingContext = baggageHeaderToDynamicSamplingContext(baggageHeader);
-
-  return { traceparentData, dynamicSamplingContext };
+  return tracingContextFromHeaders(sentryTraceHeader, baggageHeader);
 }
 
 /**