feat: spec compliant sampling result support (#1058)

Author: legendecas

packages/opentelemetry-api/src/index.ts

@@ -39,6 +39,7 @@ export * from './trace/NoopSpan';
 export * from './trace/NoopTracer';
 export * from './trace/NoopTracerProvider';
 export * from './trace/Sampler';
+export * from './trace/SamplingResult';
 export * from './trace/span_context';
 export * from './trace/span_kind';
 export * from './trace/span';

packages/opentelemetry-api/src/trace/Sampler.ts

@@ -15,6 +15,10 @@
  */
 
 import { SpanContext } from './span_context';
+import { SpanKind } from './span_kind';
+import { Attributes } from './attributes';
+import { Link } from './link';
+import { SamplingResult } from './SamplingResult';
 
 /**
  * This interface represent a sampler. Sampling is a mechanism to control the
@@ -25,12 +29,26 @@ export interface Sampler {
   /**
    * Checks whether span needs to be created and tracked.
    *
-   * TODO: Consider to add required arguments https://github.com/open-telemetry/opentelemetry-specification/blob/master/specification/sampling-api.md#shouldsample
-   * @param [parentContext] Parent span context. Typically taken from the wire.
+   * @param parentContext Parent span context. Typically taken from the wire.
    *     Can be null.
-   * @returns whether span should be sampled or not.
+   * @param traceId of the span to be created. It can be different from the
+   *     traceId in the {@link SpanContext}. Typically in situations when the
+   *     span to be created starts a new trace.
+   * @param spanName of the span to be created.
+   * @param spanKind of the span to be created.
+   * @param attributes Initial set of Attributes for the Span being constructed.
+   * @param links Collection of links that will be associated with the Span to
+   *     be created. Typically useful for batch operations.
+   * @returns a {@link SamplingResult}.
    */
-  shouldSample(parentContext?: SpanContext): boolean;
+  shouldSample(
+    parentContext: SpanContext | undefined,
+    traceId: string,
+    spanName: string,
+    spanKind: SpanKind,
+    attributes: Attributes,
+    links: Link[]
+  ): SamplingResult;
 
   /** Returns the sampler name or short description with the configuration. */
   toString(): string;

packages/opentelemetry-api/src/trace/SamplingResult.ts

@@ -0,0 +1,56 @@
+/*!
+ * Copyright 2020, OpenTelemetry Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { Attributes } from './attributes';
+
+/**
+ * A sampling decision that determines how a {@link Span} will be recorded
+ * and collected.
+ */
+export enum SamplingDecision {
+  /**
+   * `Span.isRecording() === false`, span will not be recorded and all events
+   * and attributes will be dropped.
+   */
+  NOT_RECORD,
+  /**
+   * `Span.isRecording() === true`, but `Sampled` flag in {@link TraceFlags}
+   * MUST NOT be set.
+   */
+  RECORD,
+  /**
+   * `Span.isRecording() === true` AND `Sampled` flag in {@link TraceFlags}
+   * MUST be set.
+   */
+  RECORD_AND_SAMPLED,
+}
+
+/**
+ * A sampling result contains a decision for a {@link Span} and additional
+ * attributes the sampler would like to added to the Span.
+ */
+export interface SamplingResult {
+  /**
+   * A sampling decision, refer to {@link SamplingDecision} for details.
+   */
+  decision: SamplingDecision;
+  /**
+   * The list of attributes returned by SamplingResult MUST be immutable.
+   * Caller may call {@link Sampler}.shouldSample any number of times and
+   * can safely cache the returned value.
+   */
+  attributes?: Readonly<Attributes>;
+}

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

@@ -14,7 +14,12 @@
  * limitations under the License.
  */
 
-import { Sampler, SpanContext, TraceFlags } from '@opentelemetry/api';
+import {
+  Sampler,
+  SpanContext,
+  TraceFlags,
+  SamplingDecision,
+} from '@opentelemetry/api';
 
 /** Sampler that samples a given fraction of traces. */
 export class ProbabilitySampler implements Sampler {
@@ -25,13 +30,19 @@ export class ProbabilitySampler implements Sampler {
   shouldSample(parentContext?: SpanContext) {
     // Respect the parent sampling decision if there is one
     if (parentContext && typeof parentContext.traceFlags !== 'undefined') {
-      return (
-        (TraceFlags.SAMPLED & parentContext.traceFlags) === TraceFlags.SAMPLED
-      );
+      return {
+        decision:
+          (TraceFlags.SAMPLED & parentContext.traceFlags) === TraceFlags.SAMPLED
+            ? SamplingDecision.RECORD_AND_SAMPLED
+            : SamplingDecision.NOT_RECORD,
+      };
     }
-    if (this._probability >= 1.0) return true;
-    else if (this._probability <= 0) return false;
-    return Math.random() < this._probability;
+    return {
+      decision:
+        Math.random() < this._probability
+          ? SamplingDecision.RECORD_AND_SAMPLED
+          : SamplingDecision.NOT_RECORD,
+    };
   }
 
   toString(): string {

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

@@ -15,6 +15,7 @@
  */
 
 import * as assert from 'assert';
+import * as api from '@opentelemetry/api';
 import {
   ProbabilitySampler,
   ALWAYS_SAMPLER,
@@ -24,60 +25,82 @@ import {
 describe('ProbabilitySampler', () => {
   it('should return a always sampler for 1', () => {
     const sampler = new ProbabilitySampler(1);
-    assert.strictEqual(sampler.shouldSample(), true);
+    assert.deepStrictEqual(sampler.shouldSample(), {
+      decision: api.SamplingDecision.RECORD_AND_SAMPLED,
+    });
   });
 
   it('should return a always sampler for >1', () => {
     const sampler = new ProbabilitySampler(100);
-    assert.strictEqual(sampler.shouldSample(), true);
+    assert.deepStrictEqual(sampler.shouldSample(), {
+      decision: api.SamplingDecision.RECORD_AND_SAMPLED,
+    });
     assert.strictEqual(sampler.toString(), 'ProbabilitySampler{1}');
   });
 
   it('should return a never sampler for 0', () => {
     const sampler = new ProbabilitySampler(0);
-    assert.strictEqual(sampler.shouldSample(), false);
+    assert.deepStrictEqual(sampler.shouldSample(), {
+      decision: api.SamplingDecision.NOT_RECORD,
+    });
   });
 
   it('should return a never sampler for <0', () => {
     const sampler = new ProbabilitySampler(-1);
-    assert.strictEqual(sampler.shouldSample(), false);
+    assert.deepStrictEqual(sampler.shouldSample(), {
+      decision: api.SamplingDecision.NOT_RECORD,
+    });
   });
 
   it('should sample according to the probability', () => {
     Math.random = () => 1 / 10;
     const sampler = new ProbabilitySampler(0.2);
-    assert.strictEqual(sampler.shouldSample(), true);
+    assert.deepStrictEqual(sampler.shouldSample(), {
+      decision: api.SamplingDecision.RECORD_AND_SAMPLED,
+    });
     assert.strictEqual(sampler.toString(), 'ProbabilitySampler{0.2}');
 
     Math.random = () => 5 / 10;
-    assert.strictEqual(sampler.shouldSample(), false);
+    assert.deepStrictEqual(sampler.shouldSample(), {
+      decision: api.SamplingDecision.NOT_RECORD,
+    });
   });
 
-  it('should return true for ALWAYS_SAMPLER', () => {
-    assert.strictEqual(ALWAYS_SAMPLER.shouldSample(), true);
+  it('should return api.SamplingDecision.RECORD_AND_SAMPLED for ALWAYS_SAMPLER', () => {
+    assert.deepStrictEqual(ALWAYS_SAMPLER.shouldSample(), {
+      decision: api.SamplingDecision.RECORD_AND_SAMPLED,
+    });
     assert.strictEqual(ALWAYS_SAMPLER.toString(), 'ProbabilitySampler{1}');
   });
 
-  it('should return false for NEVER_SAMPLER', () => {
-    assert.strictEqual(NEVER_SAMPLER.shouldSample(), false);
+  it('should return decision: api.SamplingDecision.NOT_RECORD for NEVER_SAMPLER', () => {
+    assert.deepStrictEqual(NEVER_SAMPLER.shouldSample(), {
+      decision: api.SamplingDecision.NOT_RECORD,
+    });
     assert.strictEqual(NEVER_SAMPLER.toString(), 'ProbabilitySampler{0}');
   });
 
   it('should handle NaN', () => {
     const sampler = new ProbabilitySampler(NaN);
-    assert.strictEqual(sampler.shouldSample(), false);
+    assert.deepStrictEqual(sampler.shouldSample(), {
+      decision: api.SamplingDecision.NOT_RECORD,
+    });
     assert.strictEqual(sampler.toString(), 'ProbabilitySampler{0}');
   });
 
   it('should handle -NaN', () => {
     const sampler = new ProbabilitySampler(-NaN);
-    assert.strictEqual(sampler.shouldSample(), false);
+    assert.deepStrictEqual(sampler.shouldSample(), {
+      decision: api.SamplingDecision.NOT_RECORD,
+    });
     assert.strictEqual(sampler.toString(), 'ProbabilitySampler{0}');
   });
 
   it('should handle undefined', () => {
     const sampler = new ProbabilitySampler(undefined);
-    assert.strictEqual(sampler.shouldSample(), false);
+    assert.deepStrictEqual(sampler.shouldSample(), {
+      decision: api.SamplingDecision.NOT_RECORD,
+    });
     assert.strictEqual(sampler.toString(), 'ProbabilitySampler{0}');
   });
 });

packages/opentelemetry-tracing/src/Tracer.ts

@@ -66,8 +66,6 @@ export class Tracer implements api.Tracer {
     context = api.context.active()
   ): api.Span {
     const parentContext = getParent(options, context);
-    // make sampling decision
-    const samplingDecision = this._sampler.shouldSample(parentContext);
     const spanId = randomSpanId();
     let traceId;
     let traceState;
@@ -79,28 +77,40 @@ export class Tracer implements api.Tracer {
       traceId = parentContext.traceId;
       traceState = parentContext.traceState;
     }
-    const traceFlags = samplingDecision
-      ? api.TraceFlags.SAMPLED
-      : api.TraceFlags.NONE;
+    const spanKind = options.kind ?? api.SpanKind.INTERNAL;
+    const links = options.links ?? [];
+    const attributes = { ...this._defaultAttributes, ...options.attributes };
+    // make sampling decision
+    const samplingResult = this._sampler.shouldSample(
+      parentContext,
+      traceId,
+      name,
+      spanKind,
+      attributes,
+      links
+    );
+
+    const traceFlags =
+      samplingResult.decision === api.SamplingDecision.RECORD_AND_SAMPLED
+        ? api.TraceFlags.SAMPLED
+        : api.TraceFlags.NONE;
     const spanContext = { traceId, spanId, traceFlags, traceState };
-    if (!samplingDecision) {
-      this.logger.debug('Sampling is off, starting no recording span');
+    if (samplingResult.decision === api.SamplingDecision.NOT_RECORD) {
+      this.logger.debug('Recording is off, starting no recording span');
       return new NoRecordingSpan(spanContext);
     }
 
     const span = new Span(
       this,
       name,
       spanContext,
-      options.kind || api.SpanKind.INTERNAL,
+      spanKind,
       parentContext ? parentContext.spanId : undefined,
-      options.links || [],
+      links,
       options.startTime
     );
     // Set default attributes
-    span.setAttributes(
-      Object.assign({}, this._defaultAttributes, options.attributes)
-    );
+    span.setAttributes(Object.assign(attributes, samplingResult.attributes));
     return span;
   }
 

packages/opentelemetry-tracing/test/Tracer.test.ts

@@ -0,0 +1,63 @@
+/*!
+ * Copyright 2020, OpenTelemetry Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import * as assert from 'assert';
+import { NoopSpan, Sampler, SamplingDecision } from '@opentelemetry/api';
+import { BasicTracerProvider, Tracer, Span } from '../src';
+import { NoopLogger, ALWAYS_SAMPLER, NEVER_SAMPLER } from '@opentelemetry/core';
+
+describe('Tracer', () => {
+  const tracerProvider = new BasicTracerProvider({
+    logger: new NoopLogger(),
+  });
+
+  class TestSampler implements Sampler {
+    shouldSample() {
+      return {
+        decision: SamplingDecision.RECORD_AND_SAMPLED,
+        attributes: {
+          testAttribute: 'foobar',
+        },
+      };
+    }
+  }
+
+  it('should create a Tracer instance', () => {
+    const tracer = new Tracer({}, tracerProvider);
+    assert.ok(tracer instanceof Tracer);
+  });
+
+  it('should respect NO_RECORD sampling result', () => {
+    const tracer = new Tracer({ sampler: NEVER_SAMPLER }, tracerProvider);
+    const span = tracer.startSpan('span1');
+    assert.ok(span instanceof NoopSpan);
+    span.end();
+  });
+
+  it('should respect RECORD_AND_SAMPLE sampling result', () => {
+    const tracer = new Tracer({ sampler: ALWAYS_SAMPLER }, tracerProvider);
+    const span = tracer.startSpan('span2');
+    assert.ok(!(span instanceof NoopSpan));
+    span.end();
+  });
+
+  it('should start a span with attributes in sampling result', () => {
+    const tracer = new Tracer({ sampler: new TestSampler() }, tracerProvider);
+    const span = tracer.startSpan('span3');
+    assert.strictEqual((span as Span).attributes.testAttribute, 'foobar');
+    span.end();
+  });
+});