fix(opentelemetry): Ensure DSC propagation works correctly (#10904)

mydea · web-flow · commit 09f2a69d740a · 2024-03-04T15:33:17.000Z
This updates the propagation handling in OTEL to the following: 1. If there is an active (local) span, we always pick up propagation context/DSC data from it. 2. Else, we try to pick it from the current scope 3. If we can't find this (e.g. a detached otel context is used), we try to pick it from the remote span context 4. Finally, if that also fails, we have no DSC Closes #10899
diff --git a/packages/node-experimental/test/integration/transactions.test.ts b/packages/node-experimental/test/integration/transactions.test.ts
@@ -1,8 +1,8 @@
 import { TraceFlags, context, trace } from '@opentelemetry/api';
 import type { SpanProcessor } from '@opentelemetry/sdk-trace-base';
 import { SEMANTIC_ATTRIBUTE_SENTRY_SOURCE, spanToJSON } from '@sentry/core';
-import { SentrySpanProcessor, setPropagationContextOnContext } from '@sentry/opentelemetry';
-import type { PropagationContext, TransactionEvent } from '@sentry/types';
+import { SentrySpanProcessor } from '@sentry/opentelemetry';
+import type { TransactionEvent } from '@sentry/types';
 import { logger } from '@sentry/utils';
 
 import * as Sentry from '../../src';
@@ -488,37 +488,27 @@ describe('Integration | Transactions', () => {
       traceFlags: TraceFlags.SAMPLED,
     };
 
-    const propagationContext: PropagationContext = {
-      traceId,
-      parentSpanId,
-      spanId: '6e0c63257de34c93',
-      sampled: true,
-    };
-
     mockSdkInit({ enableTracing: true, beforeSendTransaction });
 
     const client = Sentry.getClient()!;
 
     // We simulate the correct context we'd normally get from the SentryPropagator
-    context.with(
-      trace.setSpanContext(setPropagationContextOnContext(context.active(), propagationContext), spanContext),
-      () => {
-        Sentry.startSpan(
-          {
-            op: 'test op',
-            name: 'test name',
-            origin: 'auto.test',
-            attributes: { [Sentry.SEMANTIC_ATTRIBUTE_SENTRY_SOURCE]: 'task' },
-          },
-          () => {
-            const subSpan = Sentry.startInactiveSpan({ name: 'inner span 1' });
-            subSpan.end();
+    context.with(trace.setSpanContext(context.active(), spanContext), () => {
+      Sentry.startSpan(
+        {
+          op: 'test op',
+          name: 'test name',
+          origin: 'auto.test',
+          attributes: { [Sentry.SEMANTIC_ATTRIBUTE_SENTRY_SOURCE]: 'task' },
+        },
+        () => {
+          const subSpan = Sentry.startInactiveSpan({ name: 'inner span 1' });
+          subSpan.end();
 
-            Sentry.startSpan({ name: 'inner span 2' }, () => {});
-          },
-        );
-      },
-    );
+          Sentry.startSpan({ name: 'inner span 2' }, () => {});
+        },
+      );
+    });
 
     await client.flush();
 
diff --git a/packages/opentelemetry/src/constants.ts b/packages/opentelemetry/src/constants.ts
@@ -3,9 +3,7 @@ import { createContextKey } from '@opentelemetry/api';
 export const SENTRY_TRACE_HEADER = 'sentry-trace';
 export const SENTRY_BAGGAGE_HEADER = 'baggage';
 export const SENTRY_TRACE_STATE_DSC = 'sentry.trace';
-
-/** Context Key to hold a PropagationContext. */
-export const SENTRY_PROPAGATION_CONTEXT_CONTEXT_KEY = createContextKey('SENTRY_PROPAGATION_CONTEXT_CONTEXT_KEY');
+export const SENTRY_TRACE_STATE_PARENT_SPAN_ID = 'sentry.parent_span_id';
 
 /** Context Key to hold a Hub. */
 export const SENTRY_HUB_CONTEXT_KEY = createContextKey('sentry_hub');
diff --git a/packages/opentelemetry/src/index.ts b/packages/opentelemetry/src/index.ts
@@ -14,11 +14,7 @@ export {
   getSpanScopes,
 } from './utils/spanData';
 
-export {
-  getPropagationContextFromContext,
-  setPropagationContextOnContext,
-  getScopesFromContext,
-} from './utils/contextData';
+export { getScopesFromContext } from './utils/contextData';
 
 export {
   spanHasAttributes,
diff --git a/packages/opentelemetry/src/propagator.ts b/packages/opentelemetry/src/propagator.ts
@@ -1,7 +1,7 @@
 import type { Baggage, Context, SpanContext, TextMapGetter, TextMapSetter } from '@opentelemetry/api';
 import { TraceFlags, propagation, trace } from '@opentelemetry/api';
 import { TraceState, W3CBaggagePropagator, isTracingSuppressed } from '@opentelemetry/core';
-import { getClient, getDynamicSamplingContextFromClient } from '@sentry/core';
+import { getClient, getCurrentScope, getDynamicSamplingContextFromClient, getIsolationScope } from '@sentry/core';
 import type { DynamicSamplingContext, PropagationContext } from '@sentry/types';
 import {
   SENTRY_BAGGAGE_KEY_PREFIX,
@@ -11,28 +11,30 @@ import {
   propagationContextFromHeaders,
 } from '@sentry/utils';
 
-import { SENTRY_BAGGAGE_HEADER, SENTRY_TRACE_HEADER, SENTRY_TRACE_STATE_DSC } from './constants';
-import { getPropagationContextFromContext, setPropagationContextOnContext } from './utils/contextData';
-
-function getDynamicSamplingContextFromContext(context: Context): Partial<DynamicSamplingContext> | undefined {
-  // If possible, we want to take the DSC from the active span
-  // That should take precedence over the DSC from the propagation context
-  const activeSpan = trace.getSpan(context);
-  const traceStateDsc = activeSpan?.spanContext().traceState?.get(SENTRY_TRACE_STATE_DSC);
-  const dscOnSpan = traceStateDsc ? baggageHeaderToDynamicSamplingContext(traceStateDsc) : undefined;
-
-  if (dscOnSpan) {
-    return dscOnSpan;
-  }
-
-  const propagationContext = getPropagationContextFromContext(context);
-
-  if (propagationContext) {
-    const { traceId } = getSentryTraceData(context, propagationContext);
-    return getDynamicSamplingContext(propagationContext, traceId);
-  }
-
-  return undefined;
+import {
+  SENTRY_BAGGAGE_HEADER,
+  SENTRY_TRACE_HEADER,
+  SENTRY_TRACE_STATE_DSC,
+  SENTRY_TRACE_STATE_PARENT_SPAN_ID,
+} from './constants';
+import { getScopesFromContext, setScopesOnContext } from './utils/contextData';
+
+/** Get the Sentry propagation context from a span context. */
+export function getPropagationContextFromSpanContext(spanContext: SpanContext): PropagationContext {
+  const { traceId, spanId, traceFlags, traceState } = spanContext;
+
+  const dscString = traceState ? traceState.get(SENTRY_TRACE_STATE_DSC) : undefined;
+  const dsc = dscString ? baggageHeaderToDynamicSamplingContext(dscString) : undefined;
+  const parentSpanId = traceState ? traceState.get(SENTRY_TRACE_STATE_PARENT_SPAN_ID) : undefined;
+  const sampled = traceFlags === TraceFlags.SAMPLED;
+
+  return {
+    traceId,
+    spanId,
+    sampled,
+    parentSpanId,
+    dsc,
+  };
 }
 
 /**
@@ -49,10 +51,7 @@ export class SentryPropagator extends W3CBaggagePropagator {
 
     let baggage = propagation.getBaggage(context) || propagation.createBaggage({});
 
-    const propagationContext = getPropagationContextFromContext(context);
-    const { spanId, traceId, sampled } = getSentryTraceData(context, propagationContext);
-
-    const dynamicSamplingContext = getDynamicSamplingContextFromContext(context);
+    const { dynamicSamplingContext, traceId, spanId, sampled } = getInjectionData(context);
 
     if (dynamicSamplingContext) {
       baggage = Object.entries(dynamicSamplingContext).reduce<Baggage>((b, [dscKey, dscValue]) => {
@@ -83,15 +82,11 @@ export class SentryPropagator extends W3CBaggagePropagator {
 
     const propagationContext = propagationContextFromHeaders(sentryTraceHeader, maybeBaggageHeader);
 
-    // Add propagation context to context
-    const contextWithPropagationContext = setPropagationContextOnContext(context, propagationContext);
-
     // We store the DSC as OTEL trace state on the span context
-    const dscString = propagationContext.dsc
-      ? dynamicSamplingContextToSentryBaggageHeader(propagationContext.dsc)
-      : undefined;
-
-    const traceState = dscString ? new TraceState().set(SENTRY_TRACE_STATE_DSC, dscString) : undefined;
+    const traceState = makeTraceState({
+      parentSpanId: propagationContext.parentSpanId,
+      dsc: propagationContext.dsc,
+    });
 
     const spanContext: SpanContext = {
       traceId: propagationContext.traceId,
@@ -101,8 +96,18 @@ export class SentryPropagator extends W3CBaggagePropagator {
       traceState,
     };
 
-    // Add remote parent span context
-    return trace.setSpanContext(contextWithPropagationContext, spanContext);
+    // Add remote parent span context,
+    const ctxWithSpanContext = trace.setSpanContext(context, spanContext);
+
+    // Also update the scope on the context (to be sure this is picked up everywhere)
+    const scopes = getScopesFromContext(ctxWithSpanContext);
+    const newScopes = {
+      scope: scopes ? scopes.scope.clone() : getCurrentScope().clone(),
+      isolationScope: scopes ? scopes.isolationScope : getIsolationScope(),
+    };
+    newScopes.scope.setPropagationContext(propagationContext);
+
+    return setScopesOnContext(ctxWithSpanContext, newScopes);
   }
 
   /**
@@ -113,13 +118,91 @@ export class SentryPropagator extends W3CBaggagePropagator {
   }
 }
 
-/** Get the DSC. */
+/** Exported for tests. */
+export function makeTraceState({
+  parentSpanId,
+  dsc,
+}: { parentSpanId?: string; dsc?: Partial<DynamicSamplingContext> }): TraceState | undefined {
+  if (!parentSpanId && !dsc) {
+    return undefined;
+  }
+
+  // We store the DSC as OTEL trace state on the span context
+  const dscString = dsc ? dynamicSamplingContextToSentryBaggageHeader(dsc) : undefined;
+
+  const traceStateBase = parentSpanId
+    ? new TraceState().set(SENTRY_TRACE_STATE_PARENT_SPAN_ID, parentSpanId)
+    : new TraceState();
+
+  return dscString ? traceStateBase.set(SENTRY_TRACE_STATE_DSC, dscString) : traceStateBase;
+}
+
+function getInjectionData(context: Context): {
+  dynamicSamplingContext: Partial<DynamicSamplingContext> | undefined;
+  traceId: string | undefined;
+  spanId: string | undefined;
+  sampled: boolean | undefined;
+} {
+  const span = trace.getSpan(context);
+  const spanIsRemote = span?.spanContext().isRemote;
+
+  // If we have a local span, we can just pick everything from it
+  if (span && !spanIsRemote) {
+    const spanContext = span.spanContext();
+    const propagationContext = getPropagationContextFromSpanContext(spanContext);
+    const dynamicSamplingContext = getDynamicSamplingContext(propagationContext, spanContext.traceId);
+    return {
+      dynamicSamplingContext,
+      traceId: spanContext.traceId,
+      spanId: spanContext.spanId,
+      sampled: spanContext.traceFlags === TraceFlags.SAMPLED,
+    };
+  }
+
+  // Else we try to use the propagation context from the scope
+  const scope = getScopesFromContext(context)?.scope;
+  if (scope) {
+    const propagationContext = scope.getPropagationContext();
+    const dynamicSamplingContext = getDynamicSamplingContext(propagationContext, propagationContext.traceId);
+    return {
+      dynamicSamplingContext,
+      traceId: propagationContext.traceId,
+      spanId: propagationContext.spanId,
+      sampled: propagationContext.sampled,
+    };
+  }
+
+  // Else, we look at the remote span context
+  const spanContext = trace.getSpanContext(context);
+  if (spanContext) {
+    const propagationContext = getPropagationContextFromSpanContext(spanContext);
+    const dynamicSamplingContext = getDynamicSamplingContext(propagationContext, spanContext.traceId);
+
+    return {
+      dynamicSamplingContext,
+      traceId: spanContext.traceId,
+      spanId: spanContext.spanId,
+      sampled: spanContext.traceFlags === TraceFlags.SAMPLED,
+    };
+  }
+
+  // If we have neither, there is nothing much we can do, but that should not happen usually
+  // Unless there is a detached OTEL context being passed around
+  return {
+    dynamicSamplingContext: undefined,
+    traceId: undefined,
+    spanId: undefined,
+    sampled: undefined,
+  };
+}
+
+/** Get the DSC from a context, or fall back to use the one from the client. */
 function getDynamicSamplingContext(
   propagationContext: PropagationContext,
   traceId: string | undefined,
 ): Partial<DynamicSamplingContext> | undefined {
   // If we have a DSC on the propagation context, we just use it
-  if (propagationContext.dsc) {
+  if (propagationContext?.dsc) {
     return propagationContext.dsc;
   }
 
@@ -132,30 +215,3 @@ function getDynamicSamplingContext(
 
   return undefined;
 }
-
-/** Get the trace data for propagation. */
-function getSentryTraceData(
-  context: Context,
-  propagationContext: PropagationContext | undefined,
-): {
-  spanId: string | undefined;
-  traceId: string | undefined;
-  sampled: boolean | undefined;
-} {
-  const span = trace.getSpan(context);
-  const spanContext = span && span.spanContext();
-
-  const traceId = spanContext ? spanContext.traceId : propagationContext?.traceId;
-
-  // We have a few scenarios here:
-  // If we have an active span, and it is _not_ remote, we just use the span's ID
-  // If we have an active span that is remote, we do not want to use the spanId, as we don't want to attach it to the parent span
-  // If `isRemote === true`, the span is bascially virtual
-  // If we don't have a local active span, we use the generated spanId from the propagationContext
-  const spanId = spanContext && !spanContext.isRemote ? spanContext.spanId : propagationContext?.spanId;
-
-  // eslint-disable-next-line no-bitwise
-  const sampled = spanContext ? Boolean(spanContext.traceFlags & TraceFlags.SAMPLED) : propagationContext?.sampled;
-
-  return { traceId, spanId, sampled };
-}
diff --git a/packages/opentelemetry/src/sampler.ts b/packages/opentelemetry/src/sampler.ts
@@ -8,8 +8,8 @@ import type { Client, ClientOptions, SamplingContext } from '@sentry/types';
 import { isNaN, logger } from '@sentry/utils';
 
 import { DEBUG_BUILD } from './debug-build';
+import { getPropagationContextFromSpanContext } from './propagator';
 import { InternalSentrySemanticAttributes } from './semanticAttributes';
-import { getPropagationContextFromContext } from './utils/contextData';
 
 /**
  * A custom OTEL sampler that uses Sentry sampling rates to make it's decision
@@ -44,7 +44,7 @@ export class SentrySampler implements Sampler {
     // Note for testing: `isSpanContextValid()` checks the format of the traceId/spanId, so we need to pass valid ones
     if (parentContext && isSpanContextValid(parentContext) && parentContext.traceId === traceId) {
       if (parentContext.isRemote) {
-        parentSampled = getParentRemoteSampled(parentContext, context);
+        parentSampled = getParentRemoteSampled(parentContext);
         DEBUG_BUILD &&
           logger.log(`[Tracing] Inheriting remote parent's sampled decision for ${spanName}: ${parentSampled}`);
       } else {
@@ -178,10 +178,10 @@ function isValidSampleRate(rate: unknown): boolean {
   return true;
 }
 
-function getParentRemoteSampled(spanContext: SpanContext, context: Context): boolean | undefined {
+function getParentRemoteSampled(spanContext: SpanContext): boolean | undefined {
   const traceId = spanContext.traceId;
-  const traceparentData = getPropagationContextFromContext(context);
+  const traceparentData = getPropagationContextFromSpanContext(spanContext);
 
-  // Only inherit sample rate if `traceId` is the same
+  // Only inherit sampled if `traceId` is the same
   return traceparentData && traceId === traceparentData.traceId ? traceparentData.sampled : undefined;
 }
diff --git a/packages/opentelemetry/src/utils/contextData.ts b/packages/opentelemetry/src/utils/contextData.ts
@@ -1,31 +1,11 @@
 import type { Context } from '@opentelemetry/api';
-import type { Hub, PropagationContext, Scope } from '@sentry/types';
+import type { Hub, Scope } from '@sentry/types';
 
-import {
-  SENTRY_HUB_CONTEXT_KEY,
-  SENTRY_PROPAGATION_CONTEXT_CONTEXT_KEY,
-  SENTRY_SCOPES_CONTEXT_KEY,
-} from '../constants';
+import { SENTRY_HUB_CONTEXT_KEY, SENTRY_SCOPES_CONTEXT_KEY } from '../constants';
 import type { CurrentScopes } from '../types';
 
 const SCOPE_CONTEXT_MAP = new WeakMap<Scope, Context>();
 
-/**
- * Try to get the Propagation Context from the given OTEL context.
- * This requires the SentryPropagator to be registered.
- */
-export function getPropagationContextFromContext(context: Context): PropagationContext | undefined {
-  return context.getValue(SENTRY_PROPAGATION_CONTEXT_CONTEXT_KEY) as PropagationContext | undefined;
-}
-
-/**
- * Set a Propagation Context on an OTEL context..
- * This will return a forked context with the Propagation Context set.
- */
-export function setPropagationContextOnContext(context: Context, propagationContext: PropagationContext): Context {
-  return context.setValue(SENTRY_PROPAGATION_CONTEXT_CONTEXT_KEY, propagationContext);
-}
-
 /**
  * Try to get the Hub from the given OTEL context.
  * This requires a Context Manager that was wrapped with getWrappedContextManager.
diff --git a/packages/opentelemetry/test/integration/transactions.test.ts b/packages/opentelemetry/test/integration/transactions.test.ts
diff --git a/packages/opentelemetry/test/propagator.test.ts b/packages/opentelemetry/test/propagator.test.ts
diff --git a/packages/opentelemetry/test/trace.test.ts b/packages/opentelemetry/test/trace.test.ts
