feat: support TraceState in SamplingResult

CHANGELOG.md

@@ -11,6 +11,8 @@ For experimental package changes, see the [experimental CHANGELOG](experimental/
 
 ### :rocket: (Enhancement)
 
+* feat: support TraceState in SamplingResult [#3530](https://github.com/open-telemetry/opentelemetry-js/pull/3530) @raphael-theriault-swi
+
 ### :bug: (Bug Fix)
 
 ### :books: (Refine Doc)

api/src/trace/SamplingResult.ts

@@ -15,6 +15,7 @@
  */
 
 import { SpanAttributes } from './attributes';
+import { TraceState } from './trace_state';
 
 /**
  * @deprecated use the one declared in @opentelemetry/sdk-trace-base instead.
@@ -55,4 +56,10 @@ export interface SamplingResult {
    * can safely cache the returned value.
    */
   attributes?: Readonly<SpanAttributes>;
+  /**
+   * A {@link TraceState} that will be associated with the {@link Span}
+   * through the new {@link SpanContext}. Samplers SHOULD return the
+   * passed-in {@link TraceState} if they do not intend to change it.
+   */
+  traceState?: TraceState;
 }

packages/opentelemetry-core/src/trace/sampler/AlwaysOffSampler.ts

@@ -14,16 +14,23 @@
  * limitations under the License.
  */
 
-import { Sampler, SamplingDecision, SamplingResult } from '@opentelemetry/api';
+import {
+  Context,
+  Sampler,
+  SamplingDecision,
+  SamplingResult,
+  trace,
+} from '@opentelemetry/api';
 
 /**
  * @deprecated Use the one defined in @opentelemetry/sdk-trace-base instead.
  * Sampler that samples no traces.
  */
 export class AlwaysOffSampler implements Sampler {
-  shouldSample(): SamplingResult {
+  shouldSample(context: Context): SamplingResult {
     return {
       decision: SamplingDecision.NOT_RECORD,
+      traceState: trace.getSpanContext(context)?.traceState,
     };
   }
 

packages/opentelemetry-core/src/trace/sampler/AlwaysOnSampler.ts

@@ -14,16 +14,23 @@
  * limitations under the License.
  */
 
-import { Sampler, SamplingDecision, SamplingResult } from '@opentelemetry/api';
+import {
+  Context,
+  Sampler,
+  SamplingDecision,
+  SamplingResult,
+  trace,
+} from '@opentelemetry/api';
 
 /**
  * @deprecated Use the one defined in @opentelemetry/sdk-trace-base instead.
  * Sampler that samples all traces.
  */
 export class AlwaysOnSampler implements Sampler {
-  shouldSample(): SamplingResult {
+  shouldSample(context: Context): SamplingResult {
     return {
       decision: SamplingDecision.RECORD_AND_SAMPLED,
+      traceState: trace.getSpanContext(context)?.traceState,
     };
   }
 

packages/opentelemetry-core/src/trace/sampler/TraceIdRatioBasedSampler.ts

@@ -19,6 +19,8 @@ import {
   SamplingDecision,
   SamplingResult,
   isValidTraceId,
+  Context,
+  trace,
 } from '@opentelemetry/api';
 
 /**
@@ -33,12 +35,13 @@ export class TraceIdRatioBasedSampler implements Sampler {
     this._upperBound = Math.floor(this._ratio * 0xffffffff);
   }
 
-  shouldSample(context: unknown, traceId: string): SamplingResult {
+  shouldSample(context: Context, traceId: string): SamplingResult {
     return {
       decision:
         isValidTraceId(traceId) && this._accumulate(traceId) < this._upperBound
           ? SamplingDecision.RECORD_AND_SAMPLED
           : SamplingDecision.NOT_RECORD,
+      traceState: trace.getSpanContext(context)?.traceState,
     };
   }
 

packages/opentelemetry-sdk-trace-base/src/Sampler.ts

@@ -14,7 +14,13 @@
  * limitations under the License.
  */
 
-import { Context, Link, SpanAttributes, SpanKind } from '@opentelemetry/api';
+import {
+  Context,
+  Link,
+  SpanAttributes,
+  SpanKind,
+  TraceState,
+} from '@opentelemetry/api';
 
 /**
  * A sampling decision that determines how a {@link Span} will be recorded
@@ -53,6 +59,12 @@ export interface SamplingResult {
    * can safely cache the returned value.
    */
   attributes?: Readonly<SpanAttributes>;
+  /**
+   * A {@link TraceState} that will be associated with the {@link Span}
+   * through the new {@link SpanContext}. Samplers SHOULD return the
+   * passed-in {@link TraceState} if they do not intend to change it.
+   */
+  traceState?: TraceState;
 }
 
 /**

packages/opentelemetry-sdk-trace-base/src/Tracer.ts

@@ -117,6 +117,10 @@ export class Tracer implements api.Tracer {
       links
     );
 
+    // according to the spec an empty Tracestate means it should be cleared but
+    // to maintain backwards compatibility we ignore empty values
+    traceState = samplingResult.traceState ?? traceState;
+
     const traceFlags =
       samplingResult.decision === api.SamplingDecision.RECORD_AND_SAMPLED
         ? api.TraceFlags.SAMPLED

packages/opentelemetry-sdk-trace-base/src/sampler/AlwaysOffSampler.ts

@@ -14,13 +14,15 @@
  * limitations under the License.
  */
 
+import { Context, trace } from '@opentelemetry/api';
 import { Sampler, SamplingDecision, SamplingResult } from '../Sampler';
 
 /** Sampler that samples no traces. */
 export class AlwaysOffSampler implements Sampler {
-  shouldSample(): SamplingResult {
+  shouldSample(context: Context): SamplingResult {
     return {
       decision: SamplingDecision.NOT_RECORD,
+      traceState: trace.getSpanContext(context)?.traceState,
     };
   }
 

packages/opentelemetry-sdk-trace-base/src/sampler/AlwaysOnSampler.ts

@@ -14,13 +14,15 @@
  * limitations under the License.
  */
 
+import { Context, trace } from '@opentelemetry/api';
 import { Sampler, SamplingDecision, SamplingResult } from '../Sampler';
 
 /** Sampler that samples all traces. */
 export class AlwaysOnSampler implements Sampler {
-  shouldSample(): SamplingResult {
+  shouldSample(context: Context): SamplingResult {
     return {
       decision: SamplingDecision.RECORD_AND_SAMPLED,
+      traceState: trace.getSpanContext(context)?.traceState,
     };
   }
 

packages/opentelemetry-sdk-trace-base/src/sampler/TraceIdRatioBasedSampler.ts

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-import { isValidTraceId } from '@opentelemetry/api';
+import { Context, isValidTraceId, trace } from '@opentelemetry/api';
 import { Sampler, SamplingDecision, SamplingResult } from '../Sampler';
 
 /** Sampler that samples a given fraction of traces based of trace id deterministically. */
@@ -26,12 +26,13 @@ export class TraceIdRatioBasedSampler implements Sampler {
     this._upperBound = Math.floor(this._ratio * 0xffffffff);
   }
 
-  shouldSample(context: unknown, traceId: string): SamplingResult {
+  shouldSample(context: Context, traceId: string): SamplingResult {
     return {
       decision:
         isValidTraceId(traceId) && this._accumulate(traceId) < this._upperBound
           ? SamplingDecision.RECORD_AND_SAMPLED
           : SamplingDecision.NOT_RECORD,
+      traceState: trace.getSpanContext(context)?.traceState,
     };
   }
 

packages/opentelemetry-sdk-trace-base/test/common/Tracer.test.ts

@@ -18,13 +18,15 @@ import {
   Context,
   context,
   createContextKey,
+  createTraceState,
   INVALID_TRACEID,
   Link,
   ROOT_CONTEXT,
   SpanContext,
   SpanKind,
   trace,
   TraceFlags,
+  TraceState,
 } from '@opentelemetry/api';
 import { getSpan } from '@opentelemetry/api/build/src/trace/context-utils';
 import {
@@ -57,6 +59,8 @@ describe('Tracer', () => {
   }
 
   class TestSampler implements Sampler {
+    constructor(private readonly traceState?: TraceState) {}
+
     shouldSample(
       _context: Context,
       _traceId: string,
@@ -80,6 +84,7 @@ describe('Tracer', () => {
           // invalid attributes should be sanitized.
           ...invalidAttributes,
         } as unknown as SpanAttributes,
+        traceState: this.traceState,
       };
     }
   }
@@ -160,6 +165,17 @@ describe('Tracer', () => {
     span.end();
   });
 
+  it('should start a span with traceState in sampling result', () => {
+    const traceState = createTraceState();
+    const tracer = new Tracer(
+      { name: 'default', version: '0.0.1' },
+      { sampler: new TestSampler(traceState) },
+      tracerProvider
+    );
+    const span = tracer.startSpan('stateSpan');
+    assert.strictEqual(span.spanContext().traceState, traceState);
+  });
+
   it('should have an instrumentationLibrary', () => {
     const tracer = new Tracer(
       { name: 'default', version: '0.0.1' },
@@ -192,11 +208,13 @@ describe('Tracer', () => {
     });
   });
 
-  it('should use traceId and spanId from parent', () => {
+  it('should use traceId, spanId and traceState from parent', () => {
+    const traceState = createTraceState();
     const parent: SpanContext = {
       traceId: '00112233445566778899001122334455',
       spanId: '0011223344556677',
       traceFlags: TraceFlags.SAMPLED,
+      traceState,
     };
     const tracer = new Tracer(
       { name: 'default', version: '0.0.1' },
@@ -210,6 +228,7 @@ describe('Tracer', () => {
     );
     assert.strictEqual((span as Span).parentSpanId, parent.spanId);
     assert.strictEqual(span.spanContext().traceId, parent.traceId);
+    assert.strictEqual(span.spanContext().traceState, traceState);
   });
 
   it('should not use spanId from invalid parent', () => {

packages/opentelemetry-sdk-trace-base/test/common/sampler/AlwaysOffSampler.test.ts

@@ -25,8 +25,23 @@ describe('AlwaysOffSampler', () => {
 
   it('should return decision: api.SamplingDecision.NOT_RECORD for AlwaysOffSampler', () => {
     const sampler = new AlwaysOffSampler();
-    assert.deepStrictEqual(sampler.shouldSample(), {
+    assert.deepStrictEqual(sampler.shouldSample(api.ROOT_CONTEXT), {
       decision: api.SamplingDecision.NOT_RECORD,
+      traceState: undefined,
     });
   });
+
+  it('should forward the traceState', () => {
+    const sampler = new AlwaysOffSampler();
+    const traceState = api.createTraceState();
+    assert.strictEqual(
+      sampler.shouldSample(
+        api.trace.setSpanContext(api.ROOT_CONTEXT, {
+          ...api.INVALID_SPAN_CONTEXT,
+          traceState,
+        })
+      ).traceState,
+      traceState
+    );
+  });
 });

packages/opentelemetry-sdk-trace-base/test/common/sampler/AlwaysOnSampler.test.ts

@@ -25,8 +25,23 @@ describe('AlwaysOnSampler', () => {
 
   it('should return api.SamplingDecision.RECORD_AND_SAMPLED for AlwaysOnSampler', () => {
     const sampler = new AlwaysOnSampler();
-    assert.deepStrictEqual(sampler.shouldSample(), {
+    assert.deepStrictEqual(sampler.shouldSample(api.ROOT_CONTEXT), {
       decision: api.SamplingDecision.RECORD_AND_SAMPLED,
+      traceState: undefined,
     });
   });
+
+  it('should forward the traceState', () => {
+    const sampler = new AlwaysOnSampler();
+    const traceState = api.createTraceState();
+    assert.deepStrictEqual(
+      sampler.shouldSample(
+        api.trace.setSpanContext(api.ROOT_CONTEXT, {
+          ...api.INVALID_SPAN_CONTEXT,
+          traceState,
+        })
+      ).traceState,
+      traceState
+    );
+  });
 });

packages/opentelemetry-sdk-trace-base/test/common/sampler/ParentBasedSampler.test.ts

@@ -69,6 +69,7 @@ describe('ParentBasedSampler', () => {
       ),
       {
         decision: api.SamplingDecision.NOT_RECORD,
+        traceState: undefined,
       }
     );
   });
@@ -87,6 +88,7 @@ describe('ParentBasedSampler', () => {
       ),
       {
         decision: api.SamplingDecision.RECORD_AND_SAMPLED,
+        traceState: undefined,
       }
     );
   });
@@ -105,6 +107,7 @@ describe('ParentBasedSampler', () => {
       ),
       {
         decision: api.SamplingDecision.RECORD_AND_SAMPLED,
+        traceState: undefined,
       }
     );
   });
@@ -128,6 +131,7 @@ describe('ParentBasedSampler', () => {
       ),
       {
         decision: api.SamplingDecision.RECORD_AND_SAMPLED,
+        traceState: undefined,
       }
     );
   });
@@ -146,6 +150,7 @@ describe('ParentBasedSampler', () => {
       ),
       {
         decision: api.SamplingDecision.NOT_RECORD,
+        traceState: undefined,
       }
     );
   });
@@ -164,6 +169,7 @@ describe('ParentBasedSampler', () => {
       ),
       {
         decision: api.SamplingDecision.NOT_RECORD,
+        traceState: undefined,
       }
     );
   });