fix: pass same context to Sampler and SpanProcessor in root span case (…

…#2790)
open-telemetry · Feb 17, 2022 · eb19f3d · eb19f3d
1 parent 2da6754
commit eb19f3d
Show file tree

Hide file tree

Showing 2 changed files with 72 additions and 24 deletions.
diff --git a/packages/opentelemetry-sdk-trace-base/src/Tracer.ts b/packages/opentelemetry-sdk-trace-base/src/Tracer.ts
@@ -71,29 +71,32 @@ export class Tracer implements api.Tracer {
       return api.trace.wrapSpanContext(api.INVALID_SPAN_CONTEXT);
     }
 
-    const parentContext = getParent(options, context);
+    // remove span from context in case a root span is requested via options
+    if (options.root) {
+      context = api.trace.deleteSpan(context);
+    }
+
+    const parentSpanContext = api.trace.getSpanContext(context);
     const spanId = this._idGenerator.generateSpanId();
     let traceId;
     let traceState;
     let parentSpanId;
-    if (!parentContext || !api.trace.isSpanContextValid(parentContext)) {
+    if (!parentSpanContext || !api.trace.isSpanContextValid(parentSpanContext)) {
       // New root span.
       traceId = this._idGenerator.generateTraceId();
     } else {
       // New child span.
-      traceId = parentContext.traceId;
-      traceState = parentContext.traceState;
-      parentSpanId = parentContext.spanId;
+      traceId = parentSpanContext.traceId;
+      traceState = parentSpanContext.traceState;
+      parentSpanId = parentSpanContext.spanId;
     }
 
     const spanKind = options.kind ?? api.SpanKind.INTERNAL;
     const links = options.links ?? [];
     const attributes = sanitizeAttributes(options.attributes);
     // make sampling decision
     const samplingResult = this._sampler.shouldSample(
-      options.root
-        ? api.trace.setSpanContext(context, api.INVALID_SPAN_CONTEXT)
-        : context,
+      context,
       traceId,
       name,
       spanKind,
@@ -228,18 +231,3 @@ export class Tracer implements api.Tracer {
     return this._tracerProvider.getActiveSpanProcessor();
   }
 }
-
-/**
- * Get the parent to assign to a started span. If options.parent is null,
- * do not assign a parent.
- *
- * @param options span options
- * @param context context to check for parent
- */
-function getParent(
-  options: api.SpanOptions,
-  context: api.Context
-): api.SpanContext | undefined {
-  if (options.root) return undefined;
-  return api.trace.getSpanContext(context);
-}
diff --git a/packages/opentelemetry-sdk-trace-base/test/common/Tracer.test.ts b/packages/opentelemetry-sdk-trace-base/test/common/Tracer.test.ts
@@ -21,6 +21,7 @@ import {
   Sampler,
   SamplingDecision,
   SpanContext,
+  SpanKind,
   trace,
   TraceFlags
 } from '@opentelemetry/api';
@@ -32,7 +33,7 @@ import {
   suppressTracing
 } from '@opentelemetry/core';
 import * as assert from 'assert';
-import { BasicTracerProvider, Span, Tracer } from '../../src';
+import { BasicTracerProvider, Span, SpanProcessor, Tracer } from '../../src';
 import { TestStackContextManager } from './export/TestStackContextManager';
 import * as sinon from 'sinon';
 
@@ -53,6 +54,17 @@ describe('Tracer', () => {
     }
   }
 
+  class DummySpanProcessor implements SpanProcessor {
+      forceFlush () {
+        return Promise.resolve();
+      }
+      onStart() {}
+      onEnd() {}
+      shutdown() {
+        return Promise.resolve();
+      }
+  }
+
   beforeEach(() => {
     const contextManager = new TestStackContextManager().enable();
     context.setGlobalContextManager(contextManager);
@@ -189,6 +201,54 @@ describe('Tracer', () => {
     assert.strictEqual((span as Span).parentSpanId, undefined);
   });
 
+  it('should pass the same context to sampler and spanprocessor', () => {
+    const parent: SpanContext = {
+      traceId: '00112233445566778899001122334455',
+      spanId: '0011223344556677',
+      traceFlags: TraceFlags.SAMPLED,
+    };
+    const context = trace.setSpanContext(ROOT_CONTEXT, parent);
+
+    const sp: SpanProcessor = new DummySpanProcessor();
+    const onStartSpy = sinon.spy(sp, 'onStart');
+    const tp = new BasicTracerProvider();
+    tp.addSpanProcessor(sp);
+
+    const sampler: Sampler = new AlwaysOnSampler();
+    const shouldSampleSpy = sinon.spy(sampler, 'shouldSample');
+    const tracer = new Tracer({ name: 'default' }, { sampler }, tp);
+    const span = tracer.startSpan('a', {}, context) as Span;
+    assert.strictEqual(span.parentSpanId, parent.spanId);
+    sinon.assert.calledOnceWithExactly(shouldSampleSpy, context, parent.traceId, 'a', SpanKind.INTERNAL, {}, []);
+    sinon.assert.calledOnceWithExactly(onStartSpy, span, context);
+  });
+
+  it('should pass the same context to sampler and spanprocessor if options.root is true', () => {
+    const parent: SpanContext = {
+      traceId: '00112233445566778899001122334455',
+      spanId: '0011223344556677',
+      traceFlags: TraceFlags.SAMPLED,
+    };
+    const context = trace.setSpanContext(ROOT_CONTEXT, parent);
+
+    const sp: SpanProcessor = new DummySpanProcessor();
+    const onStartSpy = sinon.spy(sp, 'onStart');
+    const tp = new BasicTracerProvider();
+    tp.addSpanProcessor(sp);
+
+    const sampler: Sampler = new AlwaysOnSampler();
+    const shouldSampleSpy = sinon.spy(sampler, 'shouldSample');
+    const tracer = new Tracer({ name: 'default' }, { sampler }, tp);
+    const span = tracer.startSpan('a', { root: true }, context) as Span;
+    assert.strictEqual(span.parentSpanId, undefined);
+    sinon.assert.calledOnce(shouldSampleSpy);
+    sinon.assert.calledOnce(onStartSpy);
+    const samplerContext = shouldSampleSpy.firstCall.args[0];
+    const processorContext = onStartSpy.firstCall.args[1];
+    assert.strictEqual(samplerContext, processorContext);
+    assert.strictEqual(getSpan(samplerContext), undefined);
+  });
+
   it('should sample a trace when OTEL_TRACES_SAMPLER_ARG is unset', () => {
     envSource.OTEL_TRACES_SAMPLER = 'traceidratio';
     envSource.OTEL_TRACES_SAMPLER_ARG = '';