ensure trace state is picked up form spans & add tests

Author: mydea

packages/opentelemetry/src/propagator.ts

@@ -5,6 +5,7 @@ import { getClient, getDynamicSamplingContextFromClient } from '@sentry/core';
 import type { DynamicSamplingContext, PropagationContext } from '@sentry/types';
 import {
   SENTRY_BAGGAGE_KEY_PREFIX,
+  baggageHeaderToDynamicSamplingContext,
   dynamicSamplingContextToSentryBaggageHeader,
   generateSentryTraceHeader,
   propagationContextFromHeaders,
@@ -13,6 +14,27 @@ import {
 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;
+}
+
 /**
  * Injects and extracts `sentry-trace` and `baggage` headers from carriers.
  */
@@ -29,9 +51,8 @@ export class SentryPropagator extends W3CBaggagePropagator {
 
     const propagationContext = getPropagationContextFromContext(context);
     const { spanId, traceId, sampled } = getSentryTraceData(context, propagationContext);
-    const dynamicSamplingContext = propagationContext
-      ? getDynamicSamplingContext(propagationContext, traceId)
-      : undefined;
+
+    const dynamicSamplingContext = getDynamicSamplingContextFromContext(context);
 
     if (dynamicSamplingContext) {
       baggage = Object.entries(dynamicSamplingContext).reduce<Baggage>((b, [dscKey, dscValue]) => {

packages/opentelemetry/src/spanExporter.ts

@@ -5,7 +5,7 @@ import type { ReadableSpan, SpanExporter } from '@opentelemetry/sdk-trace-base';
 import { SemanticAttributes } from '@opentelemetry/semantic-conventions';
 import type { Transaction } from '@sentry/core';
 import { SEMANTIC_ATTRIBUTE_SENTRY_SOURCE } from '@sentry/core';
-import { SEMANTIC_ATTRIBUTE_SENTRY_SAMPLE_RATE, flush, getCurrentHub } from '@sentry/core';
+import { SEMANTIC_ATTRIBUTE_SENTRY_SAMPLE_RATE, getCurrentHub } from '@sentry/core';
 import type { Scope, Span as SentrySpan, SpanOrigin, TransactionSource } from '@sentry/types';
 import { addNonEnumerableProperty, dropUndefinedKeys, logger } from '@sentry/utils';
 import { startTransaction } from './custom/transaction';
@@ -74,14 +74,17 @@ export class SentrySpanExporter implements SpanExporter {
 
   /** @inheritDoc */
   public shutdown(): Promise<void> {
+    const forceFlush = this.forceFlush();
     this._stopped = true;
     this._finishedSpans = [];
-    return this.forceFlush();
+    return forceFlush;
   }
 
   /** @inheritDoc */
-  public async forceFlush(): Promise<void> {
-    await flush();
+  public forceFlush(): Promise<void> {
+    return new Promise(resolve => {
+      this.export(this._finishedSpans, () => resolve());
+    });
   }
 }
 

packages/opentelemetry/test/integration/transactions.test.ts

@@ -1,10 +1,12 @@
-import { TraceFlags, context, trace } from '@opentelemetry/api';
+import { SpanContext, TraceFlags, context, trace } from '@opentelemetry/api';
 import type { SpanProcessor } from '@opentelemetry/sdk-trace-base';
 import { SEMANTIC_ATTRIBUTE_SENTRY_SOURCE, addBreadcrumb, getClient, setTag, withIsolationScope } from '@sentry/core';
-import type { PropagationContext, TransactionEvent } from '@sentry/types';
+import type { Event, PropagationContext, TransactionEvent } from '@sentry/types';
 import { logger } from '@sentry/utils';
 
+import { TraceState } from '@opentelemetry/core';
 import { spanToJSON } from '@sentry/core';
+import { SENTRY_TRACE_STATE_DSC } from '../../src/constants';
 import { SentrySpanProcessor } from '../../src/spanProcessor';
 import { startInactiveSpan, startSpan } from '../../src/trace';
 import { setPropagationContextOnContext } from '../../src/utils/contextData';
@@ -14,6 +16,7 @@ import { cleanupOtel, getProvider, mockSdkInit } from '../helpers/mockSdkInit';
 describe('Integration | Transactions', () => {
   afterEach(() => {
     jest.restoreAllMocks();
+    jest.useRealTimers();
     cleanupOtel();
   });
 
@@ -515,4 +518,68 @@ describe('Integration | Transactions', () => {
       ]),
     );
   });
+
+  it('uses & inherits DSC on span trace state', async () => {
+    const transactionEvents: Event[] = [];
+    const beforeSendTransaction = jest.fn(event => {
+      transactionEvents.push(event);
+      return null;
+    });
+
+    const traceId = 'd4cda95b652f4a1592b449d5929fda1b';
+    const parentSpanId = '6e0c63257de34c92';
+
+    const dscString = `sentry-transaction=other-transaction,sentry-environment=other,sentry-release=8.0.0,sentry-public_key=public,sentry-trace_id=${traceId},sentry-sampled=true`;
+
+    const spanContext: SpanContext = {
+      traceId,
+      spanId: parentSpanId,
+      isRemote: true,
+      traceFlags: TraceFlags.SAMPLED,
+      traceState: new TraceState().set(SENTRY_TRACE_STATE_DSC, dscString),
+    };
+
+    const propagationContext: PropagationContext = {
+      traceId,
+      parentSpanId,
+      spanId: '6e0c63257de34c93',
+      sampled: true,
+    };
+
+    mockSdkInit({ enableTracing: true, beforeSendTransaction });
+
+    const client = getClient() as TestClientInterface;
+
+    // We simulate the correct context we'd normally get from the SentryPropagator
+    context.with(
+      trace.setSpanContext(setPropagationContextOnContext(context.active(), propagationContext), spanContext),
+      () => {
+        startSpan({ op: 'test op', name: 'test name', source: 'task', origin: 'auto.test' }, span => {
+          expect(span.spanContext().traceState?.get(SENTRY_TRACE_STATE_DSC)).toEqual(dscString);
+
+          const subSpan = startInactiveSpan({ name: 'inner span 1' });
+
+          expect(subSpan.spanContext().traceState?.get(SENTRY_TRACE_STATE_DSC)).toEqual(dscString);
+
+          subSpan.end();
+
+          startSpan({ name: 'inner span 2' }, subSpan => {
+            expect(subSpan.spanContext().traceState?.get(SENTRY_TRACE_STATE_DSC)).toEqual(dscString);
+          });
+        });
+      },
+    );
+
+    await client.flush();
+
+    expect(transactionEvents).toHaveLength(1);
+    expect(transactionEvents[0]?.sdkProcessingMetadata?.dynamicSamplingContext).toEqual({
+      environment: 'other',
+      public_key: 'public',
+      release: '8.0.0',
+      sampled: 'true',
+      trace_id: traceId,
+      transaction: 'other-transaction',
+    });
+  });
 });

packages/opentelemetry/test/propagator.test.ts

@@ -6,11 +6,11 @@ import {
   propagation,
   trace,
 } from '@opentelemetry/api';
-import { suppressTracing } from '@opentelemetry/core';
+import { TraceState, suppressTracing } from '@opentelemetry/core';
 import { addTracingExtensions, setCurrentClient } from '@sentry/core';
 import type { Client, PropagationContext } from '@sentry/types';
 
-import { SENTRY_BAGGAGE_HEADER, SENTRY_TRACE_HEADER } from '../src/constants';
+import { SENTRY_BAGGAGE_HEADER, SENTRY_TRACE_HEADER, SENTRY_TRACE_STATE_DSC } from '../src/constants';
 import { SentryPropagator } from '../src/propagator';
 import { getPropagationContextFromContext, setPropagationContextOnContext } from '../src/utils/contextData';
 
@@ -77,6 +77,33 @@ describe('SentryPropagator', () => {
           ],
           'd4cda95b652f4a1592b449d5929fda1b-6e0c63257de34c92-1',
         ],
+        [
+          'works with a DSC on the span trace state',
+          {
+            traceId: 'd4cda95b652f4a1592b449d5929fda1b',
+            spanId: '6e0c63257de34c92',
+            traceFlags: TraceFlags.SAMPLED,
+            traceState: new TraceState().set(
+              SENTRY_TRACE_STATE_DSC,
+              'sentry-transaction=other-transaction,sentry-environment=other,sentry-release=8.0.0,sentry-public_key=public,sentry-trace_id=d4cda95b652f4a1592b449d5929fda1b,sentry-sampled=true',
+            ),
+          },
+          {
+            traceId: 'd4cda95b652f4a1592b449d5929fda1b',
+            spanId: '6e0c63257de34c94',
+            parentSpanId: '6e0c63257de34c93',
+            sampled: true,
+          },
+          [
+            'sentry-environment=other',
+            'sentry-release=8.0.0',
+            'sentry-public_key=public',
+            'sentry-trace_id=d4cda95b652f4a1592b449d5929fda1b',
+            'sentry-transaction=other-transaction',
+            'sentry-sampled=true',
+          ],
+          'd4cda95b652f4a1592b449d5929fda1b-6e0c63257de34c92-1',
+        ],
         [
           'works with an unsampled propagation context',
           {