style: refactor b3 extract and tests

open-telemetry · Jul 31, 2020 · dcca659 · dcca659
1 parent bbc08dc
commit dcca659
Show file tree

Hide file tree

Showing 5 changed files with 364 additions and 281 deletions.
diff --git a/packages/opentelemetry-api/src/trace/span_context.ts b/packages/opentelemetry-api/src/trace/span_context.ts
@@ -49,7 +49,7 @@ export interface SpanContext {
    *
    * SAMPLED = 0x1 and NONE = 0x0;
    */
-  traceFlags: TraceFlags;
+  traceFlags?: TraceFlags;
   /**
    * Tracing-system-specific info to propagate.
    *

diff --git a/packages/opentelemetry-core/src/context/propagation/B3Propagator.ts b/packages/opentelemetry-core/src/context/propagation/B3Propagator.ts
@@ -37,7 +37,8 @@ export const DEBUG_FLAG_KEY = Context.createKey(
 const VALID_TRACEID_REGEX = /^([0-9a-f]{16}){1,2}$/i;
 const VALID_SPANID_REGEX = /^[0-9a-f]{16}$/i;
 const INVALID_ID_REGEX = /^0+$/i;
-const VALID_SAMPLED_VALUES = [true, 'true', '1'];
+const VALID_SAMPLED_VALUES = [true, 'true', '1', 1];
+const VALID_UNSAMPLED_VALUES = [0, '0', 'false', false];
 
 function isValidTraceId(traceId: string): boolean {
   return VALID_TRACEID_REGEX.test(traceId) && !INVALID_ID_REGEX.test(traceId);
@@ -47,6 +48,18 @@ function isValidSpanId(spanId: string): boolean {
   return VALID_SPANID_REGEX.test(spanId) && !INVALID_ID_REGEX.test(spanId);
 }
 
+function isValidParentSpanID(spanId: string | undefined): boolean {
+  return spanId === undefined || isValidSpanId(spanId);
+}
+
+function isValidSampledValue(sampled: number | undefined): boolean {
+  return (
+    sampled === undefined ||
+    VALID_SAMPLED_VALUES.includes(sampled) ||
+    VALID_UNSAMPLED_VALUES.includes(sampled)
+  );
+}
+
 function parseHeader(header: unknown) {
   return Array.isArray(header) ? header[0] : header;
 }
@@ -56,24 +69,78 @@ function parseHeader(header: unknown) {
  * Based on: https://github.com/openzipkin/b3-propagation
  */
 export class B3Propagator implements HttpTextPropagator {
+  _getHeaderValue(carrier: unknown, getter: GetterFunction, key: string) {
+    const header = getter(carrier, key);
+    return parseHeader(header);
+  }
+
+  _getTraceId(carrier: unknown, getter: GetterFunction): string {
+    const traceId = this._getHeaderValue(carrier, getter, X_B3_TRACE_ID);
+    if (typeof traceId === 'string') {
+      return traceId.padStart(32, '0');
+    }
+    return '';
+  }
+
+  _getSpanId(carrier: unknown, getter: GetterFunction): string {
+    const spanId = this._getHeaderValue(carrier, getter, X_B3_SPAN_ID);
+    if (typeof spanId === 'string') {
+      return spanId;
+    }
+    return '';
+  }
+
+  _getParentSpanId(
+    carrier: unknown,
+    getter: GetterFunction
+  ): string | undefined {
+    const spanId = this._getHeaderValue(carrier, getter, X_B3_PARENT_SPAN_ID);
+    if (typeof spanId === 'string') {
+      return spanId;
+    }
+    return;
+  }
+
+  _getDebug(carrier: unknown, getter: GetterFunction): string | undefined {
+    const debug = this._getHeaderValue(carrier, getter, X_B3_FLAGS);
+    return debug === '1' ? '1' : undefined;
+  }
+
+  _getTraceFlags(carrier: unknown, getter: GetterFunction): number | undefined {
+    const traceFlags = this._getHeaderValue(carrier, getter, X_B3_SAMPLED);
+    const debug = this._getDebug(carrier, getter);
+    if (debug === '1') {
+      return TraceFlags.SAMPLED;
+    } else if (traceFlags !== undefined) {
+      if (VALID_SAMPLED_VALUES.includes(traceFlags)) {
+        return TraceFlags.SAMPLED;
+      } else if (VALID_UNSAMPLED_VALUES.includes(traceFlags)) {
+        return TraceFlags.NONE;
+      } else {
+        // Invalid traceflag
+        return 2;
+      }
+    }
+    return;
+  }
+
   inject(context: Context, carrier: unknown, setter: SetterFunction) {
     const spanContext = getParentSpanContext(context);
     if (!spanContext) return;
-    const parentSpanId = context.getValue(PARENT_SPAN_ID_KEY);
+    const parentSpanId = context.getValue(PARENT_SPAN_ID_KEY) as
+      | undefined
+      | string;
     if (
       isValidTraceId(spanContext.traceId) &&
-      isValidSpanId(spanContext.spanId)
+      isValidSpanId(spanContext.spanId) &&
+      isValidParentSpanID(parentSpanId)
     ) {
-      if (parentSpanId) {
-        if (isValidSpanId(parentSpanId as string)) {
-          setter(carrier, X_B3_PARENT_SPAN_ID, parentSpanId);
-        } else {
-          return;
-        }
-      }
       const debug = context.getValue(DEBUG_FLAG_KEY);
       setter(carrier, X_B3_TRACE_ID, spanContext.traceId);
       setter(carrier, X_B3_SPAN_ID, spanContext.spanId);
+      if (parentSpanId) {
+        setter(carrier, X_B3_PARENT_SPAN_ID, parentSpanId);
+      }
       // According to the B3 spec, if the debug flag is set,
       // the sampled flag shouldn't be propagated as well.
       if (debug === '1') {
@@ -93,36 +160,20 @@ export class B3Propagator implements HttpTextPropagator {
   }
 
   extract(context: Context, carrier: unknown, getter: GetterFunction): Context {
-    const traceIdHeader = getter(carrier, X_B3_TRACE_ID);
-    const spanIdHeader = getter(carrier, X_B3_SPAN_ID);
-    const parentSpanIdHeader = getter(carrier, X_B3_PARENT_SPAN_ID);
-    const sampledHeader = getter(carrier, X_B3_SAMPLED);
-    const flagsHeader = getter(carrier, X_B3_FLAGS);
-
-    const traceIdHeaderValue = parseHeader(traceIdHeader);
-    const spanId = parseHeader(spanIdHeader);
-    const parentSpanId = parseHeader(parentSpanIdHeader);
-    const options = parseHeader(sampledHeader);
-    const debugHeaderValue = parseHeader(flagsHeader);
-    const debug = debugHeaderValue === '1';
-    const isSampled = VALID_SAMPLED_VALUES.includes(options);
-    const traceFlags =
-      debug || isSampled ? TraceFlags.SAMPLED : TraceFlags.NONE;
+    const traceId = this._getTraceId(carrier, getter);
+    const spanId = this._getSpanId(carrier, getter);
+    const parentSpanId = this._getParentSpanId(carrier, getter);
+    const traceFlags = this._getTraceFlags(carrier, getter);
+    const debug = this._getDebug(carrier, getter);
 
     if (
-      typeof traceIdHeaderValue !== 'string' ||
-      typeof spanId !== 'string' ||
-      (typeof parentSpanIdHeader === 'string' && !isValidSpanId(parentSpanId))
+      isValidTraceId(traceId) &&
+      isValidSpanId(spanId) &&
+      isValidParentSpanID(parentSpanId) &&
+      isValidSampledValue(traceFlags)
     ) {
-      return context;
-    }
-
-    context = context.setValue(PARENT_SPAN_ID_KEY, parentSpanId);
-    context = context.setValue(DEBUG_FLAG_KEY, debug ? '1' : undefined);
-
-    const traceId = traceIdHeaderValue.padStart(32, '0');
-
-    if (isValidTraceId(traceId) && isValidSpanId(spanId)) {
+      context = context.setValue(PARENT_SPAN_ID_KEY, parentSpanId);
+      context = context.setValue(DEBUG_FLAG_KEY, debug);
       return setExtractedSpanContext(context, {
         traceId,
         spanId,

diff --git a/packages/opentelemetry-core/src/trace/sampler/ParentOrElseSampler.ts b/packages/opentelemetry-core/src/trace/sampler/ParentOrElseSampler.ts
@@ -41,7 +41,7 @@ export class ParentOrElseSampler implements Sampler {
     links: Link[]
   ): SamplingResult {
     // Respect the parent sampling decision if there is one
-    if (parentContext) {
+    if (parentContext && parentContext.traceFlags !== undefined) {
       return {
         decision:
           (TraceFlags.SAMPLED & parentContext.traceFlags) === TraceFlags.SAMPLED

diff --git a/packages/opentelemetry-core/src/trace/sampler/ProbabilitySampler.ts b/packages/opentelemetry-core/src/trace/sampler/ProbabilitySampler.ts
@@ -31,7 +31,7 @@ export class ProbabilitySampler implements Sampler {
   shouldSample(parentContext?: SpanContext): SamplingResult {
     // Respect the parent sampling decision if there is one.
     // TODO(#1284): add an option to ignore parent regarding to spec.
-    if (parentContext) {
+    if (parentContext && parentContext.traceFlags !== undefined) {
       return {
         decision:
           (TraceFlags.SAMPLED & parentContext.traceFlags) === TraceFlags.SAMPLED