feat: support TraceState in SamplingResult (#3530)

Co-authored-by: Daniel Dyla <dyladan@users.noreply.github.com> Co-authored-by: Gerhard Stöbich <deb2001-github@yahoo.de>
open-telemetry · Feb 9, 2023 · 65e83d4 · 65e83d4
1 parent c32a845
commit 65e83d4
Show file tree

Hide file tree

Showing 5 changed files with 45 additions and 2 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -13,6 +13,7 @@ For experimental package changes, see the [experimental CHANGELOG](experimental/
 
 * feat (api-logs): separate Events API into its own package [3550](https://github.com/open-telemetry/opentelemetry-js/pull/3550) @martinkuba
 * feat(sdk-metrics): apply binary search in histogram recording [#3539](https://github.com/open-telemetry/opentelemetry-js/pull/3539) @legendecas
+* feat: support TraceState in SamplingResult [#3530](https://github.com/open-telemetry/opentelemetry-js/pull/3530) @raphael-theriault-swi
 
 ### :bug: (Bug Fix)
 

diff --git a/api/src/trace/SamplingResult.ts b/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;
 }
diff --git a/packages/opentelemetry-sdk-trace-base/src/Sampler.ts b/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;
 }
 
 /**

diff --git a/packages/opentelemetry-sdk-trace-base/src/Tracer.ts b/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

diff --git a/packages/opentelemetry-sdk-trace-base/test/common/Tracer.test.ts b/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', () => {