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,11 @@ 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 TraceState from
+   * the passed-in {@link Context} if they do not intend to change it. Leaving
+   * the value undefined will also leave the TraceState unchanged.
+   */
+  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-core/test/trace/AlwaysOffSampler.test.ts

@@ -25,8 +25,9 @@ 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,
     });
   });
 });

packages/opentelemetry-core/test/trace/AlwaysOnSampler.test.ts

@@ -25,8 +25,9 @@ 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,
     });
   });
 });

packages/opentelemetry-core/test/trace/ParentBasedSampler.test.ts

@@ -67,6 +67,7 @@ describe('ParentBasedSampler', () => {
       ),
       {
         decision: api.SamplingDecision.NOT_RECORD,
+        traceState: undefined,
       }
     );
   });
@@ -85,6 +86,7 @@ describe('ParentBasedSampler', () => {
       ),
       {
         decision: api.SamplingDecision.RECORD_AND_SAMPLED,
+        traceState: undefined,
       }
     );
   });
@@ -103,6 +105,7 @@ describe('ParentBasedSampler', () => {
       ),
       {
         decision: api.SamplingDecision.RECORD_AND_SAMPLED,
+        traceState: undefined,
       }
     );
   });
@@ -126,6 +129,7 @@ describe('ParentBasedSampler', () => {
       ),
       {
         decision: api.SamplingDecision.RECORD_AND_SAMPLED,
+        traceState: undefined,
       }
     );
   });
@@ -144,6 +148,7 @@ describe('ParentBasedSampler', () => {
       ),
       {
         decision: api.SamplingDecision.NOT_RECORD,
+        traceState: undefined,
       }
     );
   });
@@ -162,6 +167,7 @@ describe('ParentBasedSampler', () => {
       ),
       {
         decision: api.SamplingDecision.NOT_RECORD,
+        traceState: undefined,
       }
     );
   });

packages/opentelemetry-core/test/trace/TraceIdRatioBasedSampler.test.ts

@@ -18,11 +18,12 @@ import * as assert from 'assert';
 import * as api from '@opentelemetry/api';
 import { TraceIdRatioBasedSampler } from '../../src/trace/sampler/TraceIdRatioBasedSampler';
 
-const spanContext = (traceId = '1') => ({
-  traceId,
-  spanId: '1.1',
-  traceFlags: api.TraceFlags.NONE,
-});
+const spanContext = (traceId = '1') =>
+  api.trace.setSpanContext(api.ROOT_CONTEXT, {
+    traceId,
+    spanId: '1.1',
+    traceFlags: api.TraceFlags.NONE,
+  });
 
 const traceId = (part: string) => ('0'.repeat(32) + part).slice(-32);
 
@@ -59,6 +60,7 @@ describe('TraceIdRatioBasedSampler', () => {
       sampler.shouldSample(spanContext(traceId('1')), traceId('1')),
       {
         decision: api.SamplingDecision.RECORD_AND_SAMPLED,
+        traceState: undefined,
       }
     );
   });
@@ -69,6 +71,7 @@ describe('TraceIdRatioBasedSampler', () => {
       sampler.shouldSample(spanContext(traceId('1')), traceId('1')),
       {
         decision: api.SamplingDecision.RECORD_AND_SAMPLED,
+        traceState: undefined,
       }
     );
   });
@@ -79,6 +82,7 @@ describe('TraceIdRatioBasedSampler', () => {
       sampler.shouldSample(spanContext(traceId('1')), traceId('1')),
       {
         decision: api.SamplingDecision.NOT_RECORD,
+        traceState: undefined,
       }
     );
   });
@@ -89,6 +93,7 @@ describe('TraceIdRatioBasedSampler', () => {
       sampler.shouldSample(spanContext(traceId('1')), traceId('1')),
       {
         decision: api.SamplingDecision.NOT_RECORD,
+        traceState: undefined,
       }
     );
   });
@@ -100,6 +105,7 @@ describe('TraceIdRatioBasedSampler', () => {
       sampler.shouldSample(spanContext(traceId('1')), traceId('1')),
       {
         decision: api.SamplingDecision.NOT_RECORD,
+        traceState: undefined,
       }
     );
   });
@@ -111,6 +117,7 @@ describe('TraceIdRatioBasedSampler', () => {
       sampler.shouldSample(spanContext(traceId('1')), traceId('1')),
       {
         decision: api.SamplingDecision.NOT_RECORD,
+        traceState: undefined,
       }
     );
   });
@@ -122,6 +129,7 @@ describe('TraceIdRatioBasedSampler', () => {
       sampler.shouldSample(spanContext(traceId('1')), traceId('1')),
       {
         decision: api.SamplingDecision.NOT_RECORD,
+        traceState: undefined,
       }
     );
   });
@@ -132,6 +140,7 @@ describe('TraceIdRatioBasedSampler', () => {
       sampler.shouldSample(spanContext(traceId('1')), traceId('1')),
       {
         decision: api.SamplingDecision.RECORD_AND_SAMPLED,
+        traceState: undefined,
       }
     );
 
@@ -142,6 +151,7 @@ describe('TraceIdRatioBasedSampler', () => {
       ),
       {
         decision: api.SamplingDecision.NOT_RECORD,
+        traceState: undefined,
       }
     );
   });
@@ -150,12 +160,14 @@ describe('TraceIdRatioBasedSampler', () => {
     const sampler = new TraceIdRatioBasedSampler(1);
     assert.deepStrictEqual(sampler.shouldSample(spanContext(''), ''), {
       decision: api.SamplingDecision.NOT_RECORD,
+      traceState: undefined,
     });
 
     assert.deepStrictEqual(
       sampler.shouldSample(spanContext(traceId('g')), traceId('g')),
       {
         decision: api.SamplingDecision.NOT_RECORD,
+        traceState: undefined,
       }
     );
   });
@@ -167,30 +179,36 @@ describe('TraceIdRatioBasedSampler', () => {
     const id1 = traceId((Math.floor(0xffffffff * 0.1) - 1).toString(16));
     assert.deepStrictEqual(sampler10.shouldSample(spanContext(id1), id1), {
       decision: api.SamplingDecision.RECORD_AND_SAMPLED,
+      traceState: undefined,
     });
     assert.deepStrictEqual(sampler20.shouldSample(spanContext(id1), id1), {
       decision: api.SamplingDecision.RECORD_AND_SAMPLED,
+      traceState: undefined,
     });
 
     const id2 = traceId((Math.floor(0xffffffff * 0.2) - 1).toString(16));
     assert.deepStrictEqual(sampler10.shouldSample(spanContext(id2), id2), {
       decision: api.SamplingDecision.NOT_RECORD,
+      traceState: undefined,
     });
     assert.deepStrictEqual(sampler20.shouldSample(spanContext(id2), id2), {
       decision: api.SamplingDecision.RECORD_AND_SAMPLED,
+      traceState: undefined,
     });
 
     const id2delta = traceId(Math.floor(0xffffffff * 0.2).toString(16));
     assert.deepStrictEqual(
       sampler10.shouldSample(spanContext(id2delta), id2delta),
       {
         decision: api.SamplingDecision.NOT_RECORD,
+        traceState: undefined,
       }
     );
     assert.deepStrictEqual(
       sampler20.shouldSample(spanContext(id2delta), id2delta),
       {
         decision: api.SamplingDecision.NOT_RECORD,
+        traceState: undefined,
       }
     );
   });

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,13 @@ 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 TraceState from
+   * the passed-in {@link Context} if they do not intend to change it. Leaving
+   * the value undefined will also leave the TraceState unchanged.
+   */
+  traceState?: TraceState;
 }
 
 /**

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

@@ -117,6 +117,8 @@ export class Tracer implements api.Tracer {
       links
     );
 
+    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,
     };
   }