chore: sampler gets a full context (#1631)

open-telemetry · Oct 30, 2020 · 5921901 · 5921901
1 parent b523dab
commit 5921901
Show file tree

Hide file tree

Showing 5 changed files with 27 additions and 29 deletions.
diff --git a/packages/opentelemetry-api/src/trace/Sampler.ts b/packages/opentelemetry-api/src/trace/Sampler.ts
@@ -14,11 +14,11 @@
  * limitations under the License.
  */
 
-import { SpanContext } from './span_context';
-import { SpanKind } from './span_kind';
+import { Context } from '@opentelemetry/context-base';
 import { Attributes } from './attributes';
 import { Link } from './link';
 import { SamplingResult } from './SamplingResult';
+import { SpanKind } from './span_kind';
 
 /**
  * This interface represent a sampler. Sampling is a mechanism to control the
@@ -29,8 +29,7 @@ export interface Sampler {
   /**
    * Checks whether span needs to be created and tracked.
    *
-   * @param parentContext Parent span context. Typically taken from the wire.
-   *     Can be null.
+   * @param context Parent Context which may contain a span.
    * @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.
@@ -42,7 +41,7 @@ export interface Sampler {
    * @returns a {@link SamplingResult}.
    */
   shouldSample(
-    parentContext: SpanContext | undefined,
+    context: Context,
     traceId: string,
     spanName: string,
     spanKind: SpanKind,

diff --git a/packages/opentelemetry-core/src/trace/sampler/ParentBasedSampler.ts b/packages/opentelemetry-core/src/trace/sampler/ParentBasedSampler.ts
@@ -16,16 +16,17 @@
 
 import {
   Attributes,
+  Context,
+  getParentSpanContext,
   Link,
   Sampler,
   SamplingResult,
-  SpanContext,
   SpanKind,
   TraceFlags,
 } from '@opentelemetry/api';
+import { globalErrorHandler } from '../../common/global-error-handler';
 import { AlwaysOffSampler } from './AlwaysOffSampler';
 import { AlwaysOnSampler } from './AlwaysOnSampler';
-import { globalErrorHandler } from '../../common/global-error-handler';
 
 /**
  * A composite sampler that either respects the parent span's sampling decision
@@ -59,16 +60,18 @@ export class ParentBasedSampler implements Sampler {
   }
 
   shouldSample(
-    parentContext: SpanContext | undefined,
+    context: Context,
     traceId: string,
     spanName: string,
     spanKind: SpanKind,
     attributes: Attributes,
     links: Link[]
   ): SamplingResult {
+    const parentContext = getParentSpanContext(context);
+
     if (!parentContext) {
       return this._root.shouldSample(
-        parentContext,
+        context,
         traceId,
         spanName,
         spanKind,
@@ -80,7 +83,7 @@ export class ParentBasedSampler implements Sampler {
     if (parentContext.isRemote) {
       if (parentContext.traceFlags & TraceFlags.SAMPLED) {
         return this._remoteParentSampled.shouldSample(
-          parentContext,
+          context,
           traceId,
           spanName,
           spanKind,
@@ -89,7 +92,7 @@ export class ParentBasedSampler implements Sampler {
         );
       }
       return this._remoteParentNotSampled.shouldSample(
-        parentContext,
+        context,
         traceId,
         spanName,
         spanKind,
@@ -100,7 +103,7 @@ export class ParentBasedSampler implements Sampler {
 
     if (parentContext.traceFlags & TraceFlags.SAMPLED) {
       return this._localParentSampled.shouldSample(
-        parentContext,
+        context,
         traceId,
         spanName,
         spanKind,
@@ -110,7 +113,7 @@ export class ParentBasedSampler implements Sampler {
     }
 
     return this._localParentNotSampled.shouldSample(
-      parentContext,
+      context,
       traceId,
       spanName,
       spanKind,

diff --git a/packages/opentelemetry-core/src/trace/sampler/TraceIdRatioBasedSampler.ts b/packages/opentelemetry-core/src/trace/sampler/TraceIdRatioBasedSampler.ts
@@ -14,23 +14,15 @@
  * limitations under the License.
  */
 
-import {
-  Sampler,
-  SamplingDecision,
-  SamplingResult,
-  SpanContext,
-} from '@opentelemetry/api';
+import { Sampler, SamplingDecision, SamplingResult } from '@opentelemetry/api';
 
 /** Sampler that samples a given fraction of traces based of trace id deterministically. */
 export class TraceIdRatioBasedSampler implements Sampler {
   constructor(private readonly _ratio: number = 0) {
     this._ratio = this._normalize(_ratio);
   }
 
-  shouldSample(
-    parentContext: SpanContext | undefined,
-    traceId: string
-  ): SamplingResult {
+  shouldSample(context: unknown, traceId: string): SamplingResult {
     let accumulation = 0;
     for (let idx = 0; idx < traceId.length; idx++) {
       accumulation += traceId.charCodeAt(idx);

diff --git a/packages/opentelemetry-core/test/trace/ParentBasedSampler.test.ts b/packages/opentelemetry-core/test/trace/ParentBasedSampler.test.ts
@@ -17,7 +17,11 @@ import * as assert from 'assert';
 import * as api from '@opentelemetry/api';
 import { AlwaysOnSampler } from '../../src/trace/sampler/AlwaysOnSampler';
 import { ParentBasedSampler } from '../../src/trace/sampler/ParentBasedSampler';
-import { TraceFlags, SpanKind } from '@opentelemetry/api';
+import {
+  TraceFlags,
+  SpanKind,
+  setExtractedSpanContext,
+} from '@opentelemetry/api';
 import { AlwaysOffSampler } from '../../src/trace/sampler/AlwaysOffSampler';
 import { TraceIdRatioBasedSampler } from '../../src';
 
@@ -58,7 +62,7 @@ describe('ParentBasedSampler', () => {
     };
     assert.deepStrictEqual(
       sampler.shouldSample(
-        spanContext,
+        setExtractedSpanContext(api.ROOT_CONTEXT, spanContext),
         traceId,
         spanName,
         SpanKind.CLIENT,
@@ -76,7 +80,7 @@ describe('ParentBasedSampler', () => {
 
     assert.deepStrictEqual(
       sampler.shouldSample(
-        undefined,
+        api.ROOT_CONTEXT,
         traceId,
         spanName,
         SpanKind.CLIENT,
@@ -99,7 +103,7 @@ describe('ParentBasedSampler', () => {
     };
     assert.deepStrictEqual(
       sampler.shouldSample(
-        spanContext,
+        setExtractedSpanContext(api.ROOT_CONTEXT, spanContext),
         traceId,
         spanName,
         SpanKind.CLIENT,
@@ -117,7 +121,7 @@ describe('ParentBasedSampler', () => {
 
     assert.deepStrictEqual(
       sampler.shouldSample(
-        undefined,
+        api.ROOT_CONTEXT,
         traceId,
         spanName,
         SpanKind.CLIENT,

diff --git a/packages/opentelemetry-tracing/src/Tracer.ts b/packages/opentelemetry-tracing/src/Tracer.ts
@@ -89,7 +89,7 @@ export class Tracer implements api.Tracer {
     const attributes = sanitizeAttributes(options.attributes);
     // make sampling decision
     const samplingResult = this._sampler.shouldSample(
-      parentContext,
+      context,
       traceId,
       name,
       spanKind,