B3 Header Support (#753)

* [core] Add support for injecting/extract B3 headers

* reorganize distributed tracing header parsing

* work on headers

* finish up header parsing tests

* fix linting issues

* fix existing tests

* refactor existing tests

* fix linting

* fix base16 trace id truncation and parsing

* ensure we return a context

* add more tests

* move clamp_sampling_priority to helper

* simplify logic for inject/extract

* Rename to Datadog::DistributedTracing::Headers

* fix rubocop issues

* refactor helper a bit

* add Datadog.configuration.distributed_tracing

* move log into unless block

* fix line too long

lib/ddtrace/configuration/settings.rb

@@ -1,4 +1,5 @@
 require 'ddtrace/ext/analytics'
+require 'ddtrace/ext/distributed'
 require 'ddtrace/ext/runtime'
 require 'ddtrace/configuration/options'
 
@@ -21,6 +22,24 @@ class Settings
               default: -> { env_to_bool(Ext::Runtime::Metrics::ENV_ENABLED, false) },
               lazy: true
 
+      # Look for all headers by default
+      option  :propagation_extract_style,
+              default: lambda {
+                env_to_list(Ext::DistributedTracing::PROPAGATION_EXTRACT_STYLE_ENV,
+                            [Ext::DistributedTracing::PROPAGATION_STYLE_DATADOG,
+                             Ext::DistributedTracing::PROPAGATION_STYLE_B3,
+                             Ext::DistributedTracing::PROPAGATION_STYLE_B3_SINGLE_HEADER])
+              },
+              lazy: true
+
+      # Only inject Datadog headers by default
+      option  :propagation_inject_style,
+              default: lambda {
+                env_to_list(Ext::DistributedTracing::PROPAGATION_INJECT_STYLE_ENV,
+                            [Ext::DistributedTracing::PROPAGATION_STYLE_DATADOG])
+              },
+              lazy: true
+
       option :tracer, default: Tracer.new
 
       def initialize(options = {})
@@ -36,6 +55,12 @@ def configure(options = {})
         yield(self) if block_given?
       end
 
+      def distributed_tracing
+        # TODO: Move distributed tracing configuration to it's own Settings sub-class
+        # DEV: We do this to fake `Datadog.configuration.distributed_tracing.propagation_inject_style`
+        self
+      end
+
       def runtime_metrics(options = nil)
         runtime_metrics = get_option(:tracer).writer.runtime_metrics
         return runtime_metrics if options.nil?

lib/ddtrace/distributed_tracing/headers/b3.rb

@@ -0,0 +1,44 @@
+require 'ddtrace/ext/distributed'
+require 'ddtrace/distributed_tracing/headers/headers'
+require 'ddtrace/distributed_tracing/headers/helpers'
+
+module Datadog
+  module DistributedTracing
+    module Headers
+      # B3 provides helpers to inject or extract headers for B3 style headers
+      module B3
+        include Ext::DistributedTracing
+
+        def self.inject!(context, env)
+          return if context.nil?
+
+          # DEV: We need these to be hex encoded
+          env[B3_HEADER_TRACE_ID] = context.trace_id.to_s(16)
+          env[B3_HEADER_SPAN_ID] = context.span_id.to_s(16)
+
+          unless context.sampling_priority.nil?
+            sampling_priority = DistributedTracing::Headers::Helpers.clamp_sampling_priority(context.sampling_priority)
+            env[B3_HEADER_SAMPLED] = sampling_priority.to_s
+          end
+        end
+
+        def self.extract(env)
+          # Extract values from headers
+          # DEV: B3 doesn't have "origin"
+          headers = Headers.new(env)
+          trace_id = headers.id(B3_HEADER_TRACE_ID, 16)
+          span_id = headers.id(B3_HEADER_SPAN_ID, 16)
+          # We don't need to try and convert sampled since B3 supports 0/1 (AUTO_REJECT/AUTO_KEEP)
+          sampling_priority = headers.number(B3_HEADER_SAMPLED)
+
+          # Return early if this propagation is not valid
+          return unless trace_id && span_id
+
+          ::Datadog::Context.new(trace_id: trace_id,
+                                 span_id: span_id,
+                                 sampling_priority: sampling_priority)
+        end
+      end
+    end
+  end
+end

lib/ddtrace/distributed_tracing/headers/b3_single.rb

@@ -0,0 +1,56 @@
+require 'ddtrace/ext/distributed'
+require 'ddtrace/distributed_tracing/headers/headers'
+require 'ddtrace/distributed_tracing/headers/helpers'
+
+module Datadog
+  module DistributedTracing
+    module Headers
+      # B3Single provides helpers to inject or extract headers for B3 single header style headers
+      module B3Single
+        include Ext::DistributedTracing
+
+        def self.inject!(context, env)
+          return if context.nil?
+
+          # Header format:
+          #   b3: {TraceId}-{SpanId}-{SamplingState}-{ParentSpanId}
+          # https://github.com/apache/incubator-zipkin-b3-propagation/tree/7c6e9f14d6627832bd80baa87ac7dabee7be23cf#single-header
+          # DEV: `{SamplingState}` and `{ParentSpanId`}` are optional
+
+          # DEV: We need these to be hex encoded
+          header = "#{context.trace_id.to_s(16)}-#{context.span_id.to_s(16)}"
+
+          unless context.sampling_priority.nil?
+            sampling_priority = DistributedTracing::Headers::Helpers.clamp_sampling_priority(context.sampling_priority)
+            header += "-#{sampling_priority}"
+          end
+
+          env[B3_HEADER_SINGLE] = header
+        end
+
+        def self.extract(env)
+          # Header format:
+          #   b3: {TraceId}-{SpanId}-{SamplingState}-{ParentSpanId}
+          # https://github.com/apache/incubator-zipkin-b3-propagation/tree/7c6e9f14d6627832bd80baa87ac7dabee7be23cf#single-header
+          # DEV: `{SamplingState}` and `{ParentSpanId`}` are optional
+
+          headers = Headers.new(env)
+          value = headers.header(B3_HEADER_SINGLE)
+          return if value.nil?
+
+          parts = value.split('-')
+          trace_id = headers.value_to_id(parts[0], 16) unless parts.empty?
+          span_id = headers.value_to_id(parts[1], 16) if parts.length > 1
+          sampling_priority = headers.value_to_number(parts[2]) if parts.length > 2
+
+          # Return early if this propagation is not valid
+          return unless trace_id && span_id
+
+          ::Datadog::Context.new(trace_id: trace_id,
+                                 span_id: span_id,
+                                 sampling_priority: sampling_priority)
+        end
+      end
+    end
+  end
+end

lib/ddtrace/distributed_tracing/headers/datadog.rb

@@ -0,0 +1,42 @@
+require 'ddtrace/ext/distributed'
+require 'ddtrace/distributed_tracing/headers/headers'
+
+module Datadog
+  module DistributedTracing
+    module Headers
+      # Datadog provides helpers to inject or extract headers for Datadog style headers
+      module Datadog
+        include Ext::DistributedTracing
+
+        def self.inject!(context, env)
+          return if context.nil?
+
+          env[HTTP_HEADER_TRACE_ID] = context.trace_id.to_s
+          env[HTTP_HEADER_PARENT_ID] = context.span_id.to_s
+          env[HTTP_HEADER_SAMPLING_PRIORITY] = context.sampling_priority.to_s unless context.sampling_priority.nil?
+          env[HTTP_HEADER_ORIGIN] = context.origin.to_s unless context.origin.nil?
+        end
+
+        def self.extract(env)
+          # Extract values from headers
+          headers = Headers.new(env)
+          trace_id = headers.id(HTTP_HEADER_TRACE_ID)
+          parent_id = headers.id(HTTP_HEADER_PARENT_ID)
+          origin = headers.header(HTTP_HEADER_ORIGIN)
+          sampling_priority = headers.number(HTTP_HEADER_SAMPLING_PRIORITY)
+
+          # Return early if this propagation is not valid
+          # DEV: To be valid we need to have a trace id and a parent id or when it is a synthetics trace, just the trace id
+          # DEV: `DistributedHeaders#id` will not return 0
+          return unless (trace_id && parent_id) || (origin == 'synthetics' && trace_id)
+
+          # Return new context
+          ::Datadog::Context.new(trace_id: trace_id,
+                                 span_id: parent_id,
+                                 origin: origin,
+                                 sampling_priority: sampling_priority)
+        end
+      end
+    end
+  end
+end

lib/ddtrace/distributed_tracing/headers/headers.rb

@@ -0,0 +1,68 @@
+require 'ddtrace/configuration'
+require 'ddtrace/span'
+require 'ddtrace/ext/distributed'
+
+module Datadog
+  module DistributedTracing
+    module Headers
+      # Headers provides easy access and validation methods for Rack headers
+      class Headers
+        include Ext::DistributedTracing
+
+        def initialize(env)
+          @env = env
+        end
+
+        def header(name)
+          rack_header = "http-#{name}".upcase!.tr('-', '_')
+
+          hdr = @env[rack_header]
+
+          # Only return the value if it is not an empty string
+          hdr if hdr != ''
+        end
+
+        def id(hdr, base = 10)
+          value_to_id(header(hdr), base)
+        end
+
+        def value_to_id(value, base = 10)
+          id = value_to_number(value, base)
+
+          # Return early if we could not parse a number
+          return if id.nil?
+
+          # Zero or greater than max allowed value of 2**64
+          return if id.zero? || id > Span::EXTERNAL_MAX_ID
+          id < 0 ? id + (2**64) : id
+        end
+
+        def number(hdr, base = 10)
+          value_to_number(header(hdr), base)
+        end
+
+        def value_to_number(value, base = 10)
+          # It's important to make a difference between no header,
+          # and a header defined to zero.
+          return if value.nil?
+
+          # Be sure we have a string
+          value = value.to_s
+
+          # If we are parsing base16 number then truncate to 64-bit
+          value = DistributedTracing::Headers::Helpers.truncate_base16_number(value) if base == 16
+
+          # Convert header to an integer
+          # DEV: Ruby `.to_i` will return `0` if a number could not be parsed
+          num = value.to_i(base)
+
+          # Ensure the parsed number is the same as the original string value
+          # e.g. We want to make sure to throw away `'nan'.to_i == 0`
+          return unless num.to_s(base) == value
+
+          num
+        end
+      end
+    end
+  end
+end

lib/ddtrace/distributed_tracing/headers/helpers.rb

@@ -0,0 +1,45 @@
+require 'ddtrace/configuration'
+require 'ddtrace/span'
+require 'ddtrace/ext/priority'
+
+module Datadog
+  module DistributedTracing
+    module Headers
+      # Helpers module provides common helper functions for distributed tracing headers
+      module Helpers
+        # Base provides common methods for distributed header helper classes
+        def self.clamp_sampling_priority(sampling_priority)
+          # B3 doesn't have our -1 (USER_REJECT) and 2 (USER_KEEP) priorities so convert to acceptable 0/1
+          if sampling_priority < 0
+            sampling_priority = Ext::Priority::AUTO_REJECT
+          elsif sampling_priority > 1
+            sampling_priority = Ext::Priority::AUTO_KEEP
+          end
+
+          sampling_priority
+        end
+
+        def self.truncate_base16_number(value)
+          # Lowercase if we want to parse base16 e.g. 3E8 => 3e8
+          # DEV: Ruby will parse `3E8` just fine, but to test
+          #      `num.to_s(base) == value` we need to lowercase
+          value = value.downcase
+
+          # Truncate to trailing 16 characters if length is greater than 16
+          # https://github.com/apache/incubator-zipkin/blob/21fe362899fef5c593370466bc5707d3837070c2/zipkin/src/main/java/zipkin2/storage/StorageComponent.java#L49-L53
+          # DEV: This ensures we truncate B3 128-bit trace and span ids to 64-bit
+          value = value[value.length - 16, 16] if value.length > 16
+
+          # Remove any leading zeros
+          # DEV: When we call `num.to_s(16)` later Ruby will not add leading zeros
+          #      for us so we want to make sure the comparision will work as expected
+          # DEV: regex, remove all leading zeros up until we find the last 0 in the string
+          #      or we find the first non-zero, this allows `'0000' -> '0'` and `'00001' -> '1'`
+          value = value.sub(/^0*(?=(0$)|[^0])/, '')
+
+          value
+        end
+      end
+    end
+  end
+end

lib/ddtrace/environment.rb

@@ -10,6 +10,14 @@ def env_to_bool(var, default = nil)
       def env_to_float(var, default = nil)
         ENV.key?(var) ? ENV[var].to_f : default
       end
+
+      def env_to_list(var, default = [])
+        if ENV.key?(var)
+          ENV[var].split(',').map(&:strip)
+        else
+          default
+        end
+      end
     end
   end
 end

lib/ddtrace/ext/distributed.rb

@@ -10,6 +10,19 @@ module DistributedTracing
       HTTP_HEADER_ORIGIN = 'x-datadog-origin'.freeze
       ORIGIN_KEY = '_dd.origin'.freeze
 
+      # B3 headers used for distributed tracing
+      B3_HEADER_TRACE_ID = 'x-b3-traceid'.freeze
+      B3_HEADER_SPAN_ID = 'x-b3-spanid'.freeze
+      B3_HEADER_SAMPLED = 'x-b3-sampled'.freeze
+      B3_HEADER_SINGLE = 'b3'.freeze
+
+      # Distributed tracing propagation options
+      PROPAGATION_STYLE_DATADOG = 'Datadog'.freeze
+      PROPAGATION_STYLE_B3 = 'B3'.freeze
+      PROPAGATION_STYLE_B3_SINGLE_HEADER = 'B3 single header'.freeze
+      PROPAGATION_INJECT_STYLE_ENV = 'DD_PROPAGATION_INJECT_STYLE'.freeze
+      PROPAGATION_EXTRACT_STYLE_ENV = 'DD_PROPAGATION_EXTRACT_STYLE'.freeze
+
       # gRPC metadata keys for distributed tracing. https://github.com/grpc/grpc-go/blob/v1.10.x/Documentation/grpc-metadata.md
       GRPC_METADATA_TRACE_ID = 'x-datadog-trace-id'.freeze
       GRPC_METADATA_PARENT_ID = 'x-datadog-parent-id'.freeze