ensure propagation context is handled nicely

Author: mydea

packages/browser/src/tracing/browserTracingIntegration.ts

@@ -16,6 +16,7 @@ import {
   SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN,
   SEMANTIC_ATTRIBUTE_SENTRY_SOURCE,
   TRACING_DEFAULTS,
+  addNonEnumerableProperty,
   browserPerformanceTimeOrigin,
   generateTraceId,
   getActiveSpan,
@@ -246,7 +247,7 @@ export const browserTracingIntegration = ((_options: Partial<BrowserTracingOptio
   };
 
   /** Create routing idle transaction. */
-  function _createRouteSpan(client: Client, startSpanOptions: StartSpanOptions): Span {
+  function _createRouteSpan(client: Client, startSpanOptions: StartSpanOptions): void {
     const isPageloadTransaction = startSpanOptions.op === 'pageload';
 
     const finalStartSpanOptions: StartSpanOptions = beforeStartSpan
@@ -274,6 +275,15 @@ export const browserTracingIntegration = ((_options: Partial<BrowserTracingOptio
       beforeSpanEnd: span => {
         _collectWebVitals();
         addPerformanceEntries(span, { recordClsOnPageloadSpan: !enableStandaloneClsSpans });
+
+        // Ensure that DSC is updated with possibly final transaction etc.
+        const scope = getCurrentScope();
+        const oldPropagationContext = scope.getPropagationContext();
+
+        scope.setPropagationContext({
+          ...oldPropagationContext,
+          dsc: getDynamicSamplingContextFromSpan(span),
+        });
       },
     });
 
@@ -291,16 +301,31 @@ export const browserTracingIntegration = ((_options: Partial<BrowserTracingOptio
       emitFinish();
     }
 
-    return idleSpan;
+    // A trace should to stay the consistent over the entire time span of one route.
+    // Therefore, when the initial pageload or navigation root span ends, we update the
+    // scope's propagation context to keep span-specific attributes like the `sampled` decision and
+    // the dynamic sampling context valid, even after the root span has ended.
+    // This ensures that the trace data is consistent for the entire duration of the route.
+    const scope = getCurrentScope();
+    const oldPropagationContext = scope.getPropagationContext();
+
+    scope.setPropagationContext({
+      ...oldPropagationContext,
+      traceId: idleSpan.spanContext().traceId,
+      sampled: spanIsSampled(idleSpan),
+      dsc: getDynamicSamplingContextFromSpan(idleSpan),
+    });
+
+    setActiveIdleSpan(client, idleSpan);
   }
 
   return {
     name: BROWSER_TRACING_INTEGRATION_ID,
     afterAllSetup(client) {
-      let activeSpan: Span | undefined;
       let startingUrl: string | undefined = WINDOW.location && WINDOW.location.href;
 
       function maybeEndActiveSpan(): void {
+        const activeSpan = getActiveIdleSpan(client);
         if (activeSpan && !spanToJSON(activeSpan).timestamp) {
           DEBUG_BUILD && logger.log(`[Tracing] Finishing current active span with op: ${spanToJSON(activeSpan).op}`);
           // If there's an open active span, we need to finish it before creating an new one.
@@ -315,7 +340,7 @@ export const browserTracingIntegration = ((_options: Partial<BrowserTracingOptio
 
         maybeEndActiveSpan();
 
-        activeSpan = _createRouteSpan(client, {
+        _createRouteSpan(client, {
           op: 'navigation',
           ...startSpanOptions,
         });
@@ -333,33 +358,12 @@ export const browserTracingIntegration = ((_options: Partial<BrowserTracingOptio
         const propagationContext = propagationContextFromHeaders(sentryTrace, baggage);
         getCurrentScope().setPropagationContext(propagationContext);
 
-        activeSpan = _createRouteSpan(client, {
+        _createRouteSpan(client, {
           op: 'pageload',
           ...startSpanOptions,
         });
       });
 
-      // A trace should to stay the consistent over the entire time span of one route.
-      // Therefore, when the initial pageload or navigation root span ends, we update the
-      // scope's propagation context to keep span-specific attributes like the `sampled` decision and
-      // the dynamic sampling context valid, even after the root span has ended.
-      // This ensures that the trace data is consistent for the entire duration of the route.
-      client.on('spanEnd', span => {
-        const op = spanToJSON(span).op;
-        if (span !== getRootSpan(span) || (op !== 'navigation' && op !== 'pageload')) {
-          return;
-        }
-
-        const scope = getCurrentScope();
-        const oldPropagationContext = scope.getPropagationContext();
-
-        scope.setPropagationContext({
-          ...oldPropagationContext,
-          sampled: oldPropagationContext.sampled !== undefined ? oldPropagationContext.sampled : spanIsSampled(span),
-          dsc: oldPropagationContext.dsc || getDynamicSamplingContextFromSpan(span),
-        });
-      });
-
       if (WINDOW.location) {
         if (instrumentPageLoad) {
           startBrowserTracingPageLoadSpan(client, {
@@ -440,12 +444,9 @@ export function startBrowserTracingPageLoadSpan(
   traceOptions?: { sentryTrace?: string | undefined; baggage?: string | undefined },
 ): Span | undefined {
   client.emit('startPageLoadSpan', spanOptions, traceOptions);
-
   getCurrentScope().setTransactionName(spanOptions.name);
 
-  const span = getActiveSpan();
-  const op = span && spanToJSON(span).op;
-  return op === 'pageload' ? span : undefined;
+  return getActiveIdleSpan(client);
 }
 
 /**
@@ -454,15 +455,11 @@ export function startBrowserTracingPageLoadSpan(
  */
 export function startBrowserTracingNavigationSpan(client: Client, spanOptions: StartSpanOptions): Span | undefined {
   getIsolationScope().setPropagationContext({ traceId: generateTraceId() });
-  getCurrentScope().setPropagationContext({ traceId: generateTraceId() });
 
   client.emit('startNavigationSpan', spanOptions);
-
   getCurrentScope().setTransactionName(spanOptions.name);
 
-  const span = getActiveSpan();
-  const op = span && spanToJSON(span).op;
-  return op === 'navigation' ? span : undefined;
+  return getActiveIdleSpan(client);
 }
 
 /** Returns the value of a meta tag */
@@ -536,3 +533,13 @@ function registerInteractionListener(
     addEventListener('click', registerInteractionTransaction, { once: false, capture: true });
   }
 }
+
+// We store the active idle span on the client object, so we can access it from exported functions
+const ACTIVE_IDLE_SPAN_PROPERTY = '_sentry_idleSpan';
+function getActiveIdleSpan(client: Client): Span | undefined {
+  return (client as { [ACTIVE_IDLE_SPAN_PROPERTY]?: Span })[ACTIVE_IDLE_SPAN_PROPERTY];
+}
+
+function setActiveIdleSpan(client: Client, span: Span): void {
+  addNonEnumerableProperty(client, ACTIVE_IDLE_SPAN_PROPERTY, span);
+}