Propagate sampling_priority

lib/ddtrace/context.rb

@@ -14,6 +14,8 @@ module Datadog
   #
   # This data structure is thread-safe.
   class Context
+    attr_accessor :sampling_priority
+
     # Initialize a new thread-safe \Context.
     def initialize
       @mutex = Mutex.new

lib/ddtrace/contrib/faraday/middleware.rb

@@ -8,6 +8,8 @@ module Contrib
     module Faraday
       # Middleware implements a faraday-middleware for ddtrace instrumentation
       class Middleware < ::Faraday::Middleware
+        include Ext::DistributedTracing
+
         DEFAULT_ERROR_HANDLER = lambda do |env|
           Ext::HTTP::ERROR_RANGE.cover?(env[:status])
         end
@@ -54,10 +56,10 @@ def handle_response(span, env)
         end
 
         def propagate!(span, env)
-          env[:request_headers].merge!(
-            Ext::DistributedTracing::HTTP_HEADER_TRACE_ID => span.trace_id.to_s,
-            Ext::DistributedTracing::HTTP_HEADER_PARENT_ID => span.span_id.to_s
-          )
+          env[:request_headers][HTTP_HEADER_TRACE_ID] = span.trace_id.to_s
+          env[:request_headers][HTTP_HEADER_PARENT_ID] = span.span_id.to_s
+          return unless span.sampling_priority
+          env[:request_headers][HTTP_HEADER_SAMPLING_PRIORITY] = span.sampling_priority.to_s
         end
 
         def dd_pin

lib/ddtrace/contrib/http/patcher.rb

@@ -132,6 +132,9 @@ def request(req, body = nil, &block) # :yield: +response+
                   unless Datadog::Contrib::HTTP.should_skip_distributed_tracing?(pin)
                     req.add_field(Datadog::Ext::DistributedTracing::HTTP_HEADER_TRACE_ID, span.trace_id)
                     req.add_field(Datadog::Ext::DistributedTracing::HTTP_HEADER_PARENT_ID, span.span_id)
+                    if span.sampling_priority
+                      req.add_field(Datadog::Ext::DistributedTracing::HTTP_HEADER_SAMPLING_PRIORITY, span.sampling_priority)
+                    end
                   end
                 rescue StandardError => e
                   Datadog::Tracer.log.error("error preparing span for http request: #{e}")

lib/ddtrace/contrib/rack/middlewares.rb

@@ -1,21 +1,12 @@
 require 'ddtrace/ext/app_types'
 require 'ddtrace/ext/http'
-require 'ddtrace/distributed'
+require 'ddtrace/distributed_headers'
 
 module Datadog
   module Contrib
     # Rack module includes middlewares that are required to trace any framework
     # and application built on top of Rack.
     module Rack
-      # RACK headers to test when doing distributed tracing.
-      # They are slightly different from real headers as Rack uppercases everything
-
-      # Header used to transmit the trace ID.
-      HTTP_HEADER_TRACE_ID = 'HTTP_X_DATADOG_TRACE_ID'.freeze
-
-      # Header used to transmit the parent ID.
-      HTTP_HEADER_PARENT_ID = 'HTTP_X_DATADOG_PARENT_ID'.freeze
-
       # TraceMiddleware ensures that the Rack Request is properly traced
       # from the beginning to the end. The middleware adds the request span
       # in the Rack environment so that it can be retrieved by the underlying
@@ -71,16 +62,15 @@ def call(env)
           request_span = @tracer.trace('rack.request', trace_options)
 
           if @distributed_tracing_enabled
-            # Merge distributed trace ids if present
-            #
-            # Use integer values for tests, as it will catch both
-            # a non-existing header or a badly formed one.
-            trace_id, parent_id = Datadog::Distributed.parse_trace_headers(
-              env[Datadog::Contrib::Rack::HTTP_HEADER_TRACE_ID],
-              env[Datadog::Contrib::Rack::HTTP_HEADER_PARENT_ID]
-            )
-            request_span.trace_id = trace_id unless trace_id.nil?
-            request_span.parent_id = parent_id unless parent_id.nil?
+            headers = DistributedHeaders.new(env)
+
+            if headers.valid?
+              @tracer.provider.context.sampling_priority = headers.sampling_priority
+
+              request_span.trace_id = headers.trace_id
+              request_span.parent_id = headers.parent_id
+              request_span.sampling_priority = headers.sampling_priority
+            end
           end
 
           env[:datadog_rack_request_span] = request_span

lib/ddtrace/distributed_headers.rb

@@ -0,0 +1,43 @@
+require 'ddtrace/span'
+require 'ddtrace/ext/distributed'
+
+module Datadog
+  # DistributedHeaders provides easy access and validation to headers
+  class DistributedHeaders
+    include Ext::DistributedTracing
+
+    def initialize(env)
+      @env = env
+    end
+
+    def valid?
+      trace_id && parent_id && sampling_priority
+    end
+
+    def trace_id
+      value = header(HTTP_HEADER_TRACE_ID).to_i
+      return if value <= 0 || value >= Span::MAX_ID
+      value
+    end
+
+    def parent_id
+      value = header(HTTP_HEADER_PARENT_ID).to_i
+      return if value <= 0 || value >= Span::MAX_ID
+      value
+    end
+
+    def sampling_priority
+      value = header(HTTP_HEADER_SAMPLING_PRIORITY).to_f
+      return unless SAMPLING_PRIORITY_RANGE.include?(value)
+      value
+    end
+
+    private
+
+    def header(name)
+      rack_header = "http-#{name}".upcase!.tr('-', '_')
+
+      @env[rack_header]
+    end
+  end
+end

lib/ddtrace/ext/distributed.rb

@@ -5,6 +5,8 @@ module DistributedTracing
       # These are cross-language (eg: Python, Go and other implementations should honor these)
       HTTP_HEADER_TRACE_ID = 'x-datadog-trace-id'.freeze
       HTTP_HEADER_PARENT_ID = 'x-datadog-parent-id'.freeze
+      HTTP_HEADER_SAMPLING_PRIORITY = 'x-datadog-sampling-priority'.freeze
+      SAMPLING_PRIORITY_RANGE = 0...2
     end
   end
 end

lib/ddtrace/span.rb

@@ -22,7 +22,7 @@ class Span
                   :start_time, :end_time,
                   :span_id, :trace_id, :parent_id,
                   :status, :sampled,
-                  :tracer, :context
+                  :tracer, :context, :sampling_priority
 
     attr_reader :parent
 
@@ -46,6 +46,7 @@ def initialize(tracer, name, options = {})
       @span_id = Datadog::Utils.next_id
       @parent_id = options.fetch(:parent_id, 0)
       @trace_id = options.fetch(:trace_id, Datadog::Utils.next_id)
+      @sampling_priority = options[:sampling_priority]
 
       @context = options.fetch(:context, nil)
 
@@ -160,6 +161,7 @@ def parent=(parent)
         @parent_id = parent.span_id
         @service ||= parent.service
         @sampled = parent.sampled
+        @sampling_priority = parent.sampling_priority
       end
     end
 

lib/ddtrace/tracer.rb

@@ -162,24 +162,12 @@ def set_tags(tags)
     end
 
     # Guess context and parent from child_of entry.
-    def guess_context_and_parent(options = {})
-      child_of = options.fetch(:child_of, nil) # can be context or span
-
-      ctx = nil
-      parent = nil
-      unless child_of.nil?
-        if child_of.respond_to?(:current_span)
-          ctx = child_of
-          parent = child_of.current_span
-        elsif child_of.is_a?(Datadog::Span)
-          parent = child_of
-          ctx = child_of.context
-        end
-      end
+    def guess_context_and_parent(child_of)
+      return [call_context, nil] unless child_of
 
-      ctx ||= call_context
+      return [child_of, child_of.current_span] if child_of.is_a?(Context)
 
-      [ctx, parent]
+      [child_of.context, child_of]
     end
 
     # Return a span that will trace an operation called \name. This method allows
@@ -202,7 +190,7 @@ def start_span(name, options = {})
         [:service, :resource, :span_type].include?(k)
       end
 
-      ctx, parent = guess_context_and_parent(options)
+      ctx, parent = guess_context_and_parent(options[:child_of])
       opts[:context] = ctx unless ctx.nil?
 
       span = Span.new(self, name, opts)