ddtrace 0.9.2 → 0.10.0

data/lib/ddtrace/propagation/distributed_headers.rb

@@ -0,0 +1,48 @@
+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?
+      # Sampling priority is optional.
+      trace_id && parent_id
+    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
+      hdr = header(HTTP_HEADER_SAMPLING_PRIORITY)
+      # It's important to make a difference between no header,
+      # and a header defined to zero.
+      return unless hdr
+      value = hdr.to_i
+      return if value < 0
+      value
+    end
+
+    private
+
+    def header(name)
+      rack_header = "http-#{name}".upcase!.tr('-', '_')
+
+      @env[rack_header]
+    end
+  end
+end

data/lib/ddtrace/propagation/http_propagator.rb

@@ -0,0 +1,28 @@
+require 'ddtrace/context'
+require 'ddtrace/ext/distributed'
+require 'ddtrace/propagation/distributed_headers'
+
+module Datadog
+  # HTTPPropagator helps extracting and injecting HTTP headers.
+  module HTTPPropagator
+    include Ext::DistributedTracing
+
+    # inject! popolates the env with span ID, trace ID and sampling priority
+    def self.inject!(context, env)
+      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
+      env.delete(HTTP_HEADER_SAMPLING_PRIORITY) unless context.sampling_priority
+    end
+
+    # extract returns a context containing the span ID, trace ID and
+    # sampling priority defined in env.
+    def self.extract(env)
+      headers = DistributedHeaders.new(env)
+      return Datadog::Context.new unless headers.valid?
+      Datadog::Context.new(trace_id: headers.trace_id,
+                           span_id: headers.parent_id,
+                           sampling_priority: headers.sampling_priority)
+    end
+  end
+end

data/lib/ddtrace/registry.rb

@@ -0,0 +1,42 @@
+require_relative 'registry/registerable'
+
+module Datadog
+  # Registry provides insertion/retrieval capabilities for integrations
+  class Registry
+    include Enumerable
+
+    Entry = Struct.new(:name, :klass, :auto_patch)
+
+    def initialize
+      @data = {}
+      @mutex = Mutex.new
+    end
+
+    def add(name, klass, auto_patch = false)
+      @mutex.synchronize do
+        @data[name] = Entry.new(name, klass, auto_patch).freeze
+      end
+    end
+
+    def each
+      @mutex.synchronize do
+        @data.each { |_, entry| yield(entry) }
+      end
+    end
+
+    def [](name)
+      @mutex.synchronize do
+        entry = @data[name]
+        entry.klass if entry
+      end
+    end
+
+    def to_h
+      @mutex.synchronize do
+        @data.each_with_object({}) do |(_, entry), hash|
+          hash[entry.name] = entry.auto_patch
+        end
+      end
+    end
+  end
+end

data/lib/ddtrace/registry/registerable.rb

@@ -0,0 +1,20 @@
+module Datadog
+  class Registry
+    # Registerable provides a convenience method for self-registering
+    module Registerable
+      def self.included(base)
+        base.singleton_class.send(:include, ClassMethods)
+      end
+
+      # ClassMethods
+      module ClassMethods
+        def register_as(name, options = {})
+          registry = options.fetch(:registry, Datadog.registry)
+          auto_patch = options.fetch(:auto_patch, false)
+
+          registry.add(name, self, auto_patch)
+        end
+      end
+    end
+  end
+end

data/lib/ddtrace/sampler.rb

@@ -1,3 +1,5 @@
+require 'forwardable'
+
 module Datadog
   # \Sampler performs client-side trace sampling.
   class Sampler
@@ -42,8 +44,66 @@ module Datadog
     end
 
     def sample(span)
-      span.sampled = ((span.trace_id * KNUTH_FACTOR) % Datadog::Span::MAX_ID) <= @sampling_id_threshold
       span.set_metric(SAMPLE_RATE_METRIC_KEY, @sample_rate)
+      span.sampled = ((span.trace_id * KNUTH_FACTOR) % Datadog::Span::MAX_ID) <= @sampling_id_threshold
+    end
+  end
+
+  # \RateByServiceSampler samples different services at different rates
+  class RateByServiceSampler < Sampler
+    DEFAULT_KEY = 'service:,env:'.freeze
+
+    def initialize(rate = 1.0, opts = {})
+      @env = opts.fetch(:env, Datadog.tracer.tags[:env])
+      @mutex = Mutex.new
+      @fallback = RateSampler.new(rate)
+      @sampler = { DEFAULT_KEY => @fallback }
+    end
+
+    def sample(span)
+      key = key_for(span)
+
+      @mutex.synchronize do
+        @sampler.fetch(key, @fallback).sample(span)
+      end
+    end
+
+    def update(rate_by_service)
+      @mutex.synchronize do
+        @sampler.delete_if { |key, _| key != DEFAULT_KEY && !rate_by_service.key?(key) }
+
+        rate_by_service.each do |key, rate|
+          @sampler[key] ||= RateSampler.new(rate)
+          @sampler[key].sample_rate = rate
+        end
+      end
+    end
+
+    private
+
+    def key_for(span)
+      "service:#{span.service},env:#{@env}"
+    end
+  end
+
+  # \PrioritySampler
+  class PrioritySampler
+    extend Forwardable
+
+    def initialize(opts = {})
+      @base_sampler = opts[:base_sampler] || RateSampler.new
+      @post_sampler = opts[:post_sampler] || RateByServiceSampler.new
     end
+
+    def sample(span)
+      span.context.sampling_priority = 0 if span.context
+      return unless @base_sampler.sample(span)
+      return unless @post_sampler.sample(span)
+      span.context.sampling_priority = 1 if span.context
+
+      true
+    end
+
+    def_delegators :@post_sampler, :update
   end
 end

data/lib/ddtrace/sync_writer.rb

@@ -0,0 +1,36 @@
+module Datadog
+  # SyncWriter flushes both services and traces synchronously
+  class SyncWriter
+    attr_reader :transport
+
+    def initialize(options = {})
+      @transport = options.fetch(:transport) do
+        HTTPTransport.new(Writer::HOSTNAME, Writer::PORT)
+      end
+    end
+
+    def write(trace, services)
+      perform_concurrently(
+        proc { flush_services(services) },
+        proc { flush_trace(trace) }
+      )
+    rescue => e
+      Tracer.log.debug(e)
+    end
+
+    private
+
+    def perform_concurrently(*tasks)
+      tasks.map { |task| Thread.new(&task) }.each(&:join)
+    end
+
+    def flush_services(services)
+      transport.send(:services, services)
+    end
+
+    def flush_trace(trace)
+      processed_traces = Pipeline.process!([trace])
+      transport.send(:traces, processed_traces)
+    end
+  end
+end

data/lib/ddtrace/tracer.rb

@@ -18,8 +18,8 @@ module Datadog
   # of these function calls and sub-requests would be encapsulated within a single trace.
   # rubocop:disable Metrics/ClassLength
   class Tracer
-    attr_reader :writer, :sampler, :services, :tags, :provider
-    attr_accessor :enabled
+    attr_reader :sampler, :services, :tags, :provider
+    attr_accessor :enabled, :writer
     attr_writer :default_service
 
     # Global, memoized, lazy initialized instance of a logger that is used within the the Datadog
@@ -71,7 +71,7 @@ module Datadog
     #
     def shutdown!
       return if !@enabled || @writer.worker.nil?
-      @writer.worker.shutdown!
+      @writer.worker.stop
     end
 
     # Return the current active \Context for this traced execution. This method is
@@ -118,11 +118,18 @@ module Datadog
       hostname = options.fetch(:hostname, nil)
       port = options.fetch(:port, nil)
       sampler = options.fetch(:sampler, nil)
+      priority_sampling = options[:priority_sampling]
 
       @enabled = enabled unless enabled.nil?
+      @sampler = sampler unless sampler.nil?
+
+      if priority_sampling
+        @sampler = PrioritySampler.new(base_sampler: @sampler)
+        @writer = Writer.new(priority_sampler: @sampler)
+      end
+
       @writer.transport.hostname = hostname unless hostname.nil?
       @writer.transport.port = port unless port.nil?
-      @sampler = sampler unless sampler.nil?
     end
 
     # Set the information about the given service. A valid example is:
@@ -162,24 +169,15 @@ module Datadog
     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)
+      # call_context should not be in this code path, as start_span
+      # should never try and pick an existing context, but only get
+      # it from the parameters passed to it (child_of)
+      return [Datadog::Context.new, 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 +200,7 @@ module Datadog
         [: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)
@@ -210,6 +208,10 @@ module Datadog
         # root span
         @sampler.sample(span)
         span.set_tag('system.pid', Process.pid)
+        if ctx && ctx.trace_id && ctx.span_id
+          span.trace_id = ctx.trace_id
+          span.parent_id = ctx.span_id
+        end
       else
         # child span
         span.parent = parent # sets service, trace_id, parent_id, sampled

data/lib/ddtrace/transport.rb

@@ -19,13 +19,36 @@ module Datadog
     TRACE_COUNT_HEADER = 'X-Datadog-Trace-Count'.freeze
     RUBY_INTERPRETER = RUBY_VERSION > '1.9' ? RUBY_ENGINE + '-' + RUBY_PLATFORM : 'ruby-' + RUBY_PLATFORM
 
+    API = {
+      V4 = 'v0.4'.freeze => {
+        traces_endpoint: '/v0.4/traces'.freeze,
+        services_endpoint: '/v0.4/services'.freeze,
+        encoder: Encoding::MsgpackEncoder,
+        fallback: 'v0.3'.freeze
+      }.freeze,
+      V3 = 'v0.3'.freeze => {
+        traces_endpoint: '/v0.3/traces'.freeze,
+        services_endpoint: '/v0.3/services'.freeze,
+        encoder: Encoding::MsgpackEncoder,
+        fallback: 'v0.2'.freeze
+      }.freeze,
+      V2 = 'v0.2'.freeze => {
+        traces_endpoint: '/v0.2/traces'.freeze,
+        services_endpoint: '/v0.2/services'.freeze,
+        encoder: Encoding::JSONEncoder
+      }.freeze
+    }.freeze
+
+    private_constant :API
+
     def initialize(hostname, port, options = {})
+      api_version = options.fetch(:api_version, V3)
+
       @hostname = hostname
       @port = port
-      @traces_endpoint = '/v0.3/traces'.freeze
-      @services_endpoint = '/v0.3/services'.freeze
-      @compatibility_mode = false
-      @encoder = options.fetch(:encoder, Datadog::Encoding::MsgpackEncoder.new())
+      @api = API.fetch(api_version)
+      @encoder = options[:encoder] || @api[:encoder].new
+      @response_callback = options[:response_callback]
 
       # overwrite the Content-type with the one chosen in the Encoder
       @headers = options.fetch(:headers, {})
@@ -48,21 +71,19 @@ module Datadog
       case endpoint
       when :services
         payload = @encoder.encode_services(data)
-        status_code = post(@services_endpoint, payload)
+        status_code = post(@api[:services_endpoint], payload)
       when :traces
         count = data.length
         payload = @encoder.encode_traces(data)
-        status_code = post(@traces_endpoint, payload, count)
+        status_code = post(@api[:traces_endpoint], payload, count)
       else
         Datadog::Tracer.log.error("Unsupported endpoint: #{endpoint}")
         return nil
       end
 
-      return status_code unless downgrade?(status_code) && !@compatibility_mode
+      downgrade! && send(endpoint, data) if downgrade?(status_code)
 
-      # the API endpoint is not available so we should downgrade the connection and re-try the call
-      downgrade!
-      send(endpoint, data)
+      status_code
     end
 
     # send data to the trace-agent; the method is thread-safe
@@ -84,11 +105,13 @@ module Datadog
     # this method should target a stable API that works whatever is the agent
     # or the tracing client versions.
     def downgrade!
-      @compatibility_mode = true
-      @traces_endpoint = '/v0.2/traces'.freeze
-      @services_endpoint = '/v0.2/services'.freeze
-      @encoder = Datadog::Encoding::JSONEncoder.new()
-      @headers['Content-Type'] = @encoder.content_type
+      @mutex.synchronize do
+        fallback_version = @api.fetch(:fallback)
+
+        @api = API.fetch(fallback_version)
+        @encoder = @api[:encoder].new
+        @headers['Content-Type'] = @encoder.content_type
+      end
     end
 
     def informational?(code)
@@ -119,6 +142,8 @@ module Datadog
     # endpoint. In both cases, we're going to downgrade the transporter encoder so that
     # it will target a stable API.
     def downgrade?(code)
+      return unless @api[:fallback]
+
       code == 404 || code == 415
     end
 
@@ -141,6 +166,8 @@ module Datadog
         @mutex.synchronize { @count_server_error += 1 }
       end
 
+      process_callback(response)
+
       status_code
     rescue StandardError => e
       Datadog::Tracer.log.error(e.message)
@@ -158,5 +185,15 @@ module Datadog
         }
       end
     end
+
+    private
+
+    def process_callback(response)
+      return unless @response_callback && @response_callback.respond_to?(:call)
+
+      @response_callback.call(response)
+    rescue => e
+      Tracer.log.debug("Error processing callback: #{e}")
+    end
   end
 end