Add event_sample_rate to integrations

Rakefile

@@ -46,7 +46,7 @@ namespace :spec do
   end
 
   RSpec::Core::RakeTask.new(:contrib) do |t|
-    t.pattern = 'spec/**/contrib/{configurable,integration,patchable,patcher,registerable,configuration/*}_spec.rb'
+    t.pattern = 'spec/**/contrib/{configurable,integration,patchable,patcher,registerable,sampling,configuration/*}_spec.rb'
   end
 
   [

lib/ddtrace/contrib/configuration/settings.rb

@@ -9,6 +9,7 @@ class Settings
 
         option :service_name
         option :tracer, default: Datadog.tracer
+        option :event_sample_rate
 
         def initialize(options = {})
           configure(options)

lib/ddtrace/contrib/delayed_job/plugin.rb

@@ -1,4 +1,5 @@
 require 'delayed/plugin'
+require 'ddtrace/contrib/sampling'
 require 'ddtrace/contrib/delayed_job/ext'
 
 module Datadog
@@ -18,6 +19,7 @@ def self.instrument(job, &block)
                      end
 
           tracer.trace(Ext::SPAN_JOB, service: configuration[:service_name], resource: job_name) do |span|
+            Contrib::Sampling.set_event_sample_rate(span, configuration[:event_sample_rate])
             span.set_tag(Ext::TAG_ID, job.id)
             span.set_tag(Ext::TAG_QUEUE, job.queue) if job.queue
             span.set_tag(Ext::TAG_PRIORITY, job.priority)

lib/ddtrace/contrib/racecar/event.rb

@@ -1,3 +1,4 @@
+require 'ddtrace/contrib/sampling'
 require 'ddtrace/contrib/active_support/notifications/event'
 require 'ddtrace/contrib/racecar/ext'
 
@@ -36,6 +37,7 @@ def process(span, event, _id, payload)
             span.service = configuration[:service_name]
             span.resource = payload[:consumer_class]
 
+            Contrib::Sampling.set_event_sample_rate(span, configuration[:event_sample_rate])
             span.set_tag(Ext::TAG_TOPIC, payload[:topic])
             span.set_tag(Ext::TAG_CONSUMER, payload[:consumer_class])
             span.set_tag(Ext::TAG_PARTITION, payload[:partition])

lib/ddtrace/contrib/rack/middlewares.rb

@@ -1,6 +1,7 @@
 require 'ddtrace/ext/app_types'
 require 'ddtrace/ext/http'
 require 'ddtrace/propagation/http_propagator'
+require 'ddtrace/contrib/sampling'
 require 'ddtrace/contrib/rack/ext'
 require 'ddtrace/contrib/rack/request_queue'
 
@@ -25,7 +26,7 @@ def initialize(app)
         end
 
         def compute_queue_time(env, tracer)
-          return unless Datadog.configuration[:rack][:request_queuing]
+          return unless configuration[:request_queuing]
 
           # parse the request queue time
           request_start = Datadog::Contrib::Rack::QueueTime.get_request_start(env)
@@ -34,25 +35,25 @@ def compute_queue_time(env, tracer)
           tracer.trace(
             Ext::SPAN_HTTP_SERVER_QUEUE,
             start_time: request_start,
-            service: Datadog.configuration[:rack][:web_service_name]
+            service: configuration[:web_service_name]
           )
         end
 
         def call(env)
           # retrieve integration settings
-          tracer = Datadog.configuration[:rack][:tracer]
+          tracer = configuration[:tracer]
 
           # [experimental] create a root Span to keep track of frontend web servers
           # (i.e. Apache, nginx) if the header is properly set
           frontend_span = compute_queue_time(env, tracer)
 
           trace_options = {
-            service: Datadog.configuration[:rack][:service_name],
+            service: configuration[:service_name],
             resource: nil,
             span_type: Datadog::Ext::HTTP::TYPE
           }
 
-          if Datadog.configuration[:rack][:distributed_tracing]
+          if configuration[:distributed_tracing]
             context = HTTPPropagator.extract(env)
             tracer.provider.context = context if context.trace_id
           end
@@ -110,7 +111,7 @@ def call(env)
         end
 
         def resource_name_for(env, status)
-          if Datadog.configuration[:rack][:middleware_names] && env['RESPONSE_MIDDLEWARE']
+          if configuration[:middleware_names] && env['RESPONSE_MIDDLEWARE']
             "#{env['RESPONSE_MIDDLEWARE']}##{env['REQUEST_METHOD']}"
           else
             "#{env['REQUEST_METHOD']} #{status}".strip
@@ -134,13 +135,19 @@ def set_request_tags!(request_span, env, status, headers, response, original_env
           response_headers = parse_response_headers(headers || {})
 
           request_span.resource ||= resource_name_for(env, status)
+
+          # Set event sample rate, if available.
+          Contrib::Sampling.set_event_sample_rate(request_span, configuration[:event_sample_rate])
+
           if request_span.get_tag(Datadog::Ext::HTTP::METHOD).nil?
             request_span.set_tag(Datadog::Ext::HTTP::METHOD, env['REQUEST_METHOD'])
           end
+
           if request_span.get_tag(Datadog::Ext::HTTP::URL).nil?
-            options = Datadog.configuration[:rack][:quantize]
+            options = configuration[:quantize]
             request_span.set_tag(Datadog::Ext::HTTP::URL, Datadog::Quantization::HTTP.url(url, options))
           end
+
           if request_span.get_tag(Datadog::Ext::HTTP::BASE_URL).nil?
             request_obj = ::Rack::Request.new(env)
 
@@ -153,6 +160,7 @@ def set_request_tags!(request_span, env, status, headers, response, original_env
 
             request_span.set_tag(Datadog::Ext::HTTP::BASE_URL, base_url)
           end
+
           if request_span.get_tag(Datadog::Ext::HTTP::STATUS_CODE).nil? && status
             request_span.set_tag(Datadog::Ext::HTTP::STATUS_CODE, status)
           end
@@ -182,6 +190,10 @@ def set_request_tags!(request_span, env, status, headers, response, original_env
           If you need the Rack request span, try using `Datadog.tracer.active_span`.
           This key will be removed in version 1.0).freeze
 
+        def configuration
+          Datadog.configuration[:rack]
+        end
+
         def add_deprecation_warnings(env)
           env.instance_eval do
             unless instance_variable_defined?(:@patched_with_datadog_warnings)
@@ -221,7 +233,7 @@ def without_datadog_warnings
 
         def parse_request_headers(env)
           {}.tap do |result|
-            whitelist = Datadog.configuration[:rack][:headers][:request] || []
+            whitelist = configuration[:headers][:request] || []
             whitelist.each do |header|
               rack_header = header_to_rack_header(header)
               if env.key?(rack_header)
@@ -233,7 +245,7 @@ def parse_request_headers(env)
 
         def parse_response_headers(headers)
           {}.tap do |result|
-            whitelist = Datadog.configuration[:rack][:headers][:response] || []
+            whitelist = configuration[:headers][:response] || []
             whitelist.each do |header|
               if headers.key?(header)
                 result[Datadog::Ext::HTTP::ResponseHeaders.to_tag(header)] = headers[header]

lib/ddtrace/contrib/rake/instrumentation.rb

@@ -1,3 +1,4 @@
+require 'ddtrace/contrib/sampling'
 require 'ddtrace/contrib/rake/ext'
 
 module Datadog
@@ -41,6 +42,7 @@ def shutdown_tracer!
 
           def annotate_invoke!(span, args)
             span.resource = name
+            Contrib::Sampling.set_event_sample_rate(span, configuration[:event_sample_rate])
             span.set_tag(Ext::TAG_TASK_ARG_NAMES, arg_names)
             span.set_tag(Ext::TAG_INVOKE_ARGS, quantize_args(args)) unless args.nil?
           rescue StandardError => e

lib/ddtrace/contrib/resque/patcher.rb

@@ -19,23 +19,13 @@ def patch
           do_once(:resque) do
             begin
               require_relative 'resque_job'
-              add_pin
               get_option(:workers).each { |worker| worker.extend(ResqueJob) }
             rescue StandardError => e
               Datadog::Tracer.log.error("Unable to apply Resque integration: #{e}")
             end
           end
         end
 
-        def add_pin
-          Pin.new(
-            get_option(:service_name),
-            app: Ext::APP,
-            app_type: Datadog::Ext::AppTypes::WORKER,
-            tracer: get_option(:tracer)
-          ).onto(::Resque)
-        end
-
         def get_option(option)
           Datadog.configuration[:resque].get_option(option)
         end

lib/ddtrace/contrib/resque/resque_job.rb

@@ -1,5 +1,6 @@
 require 'ddtrace/ext/app_types'
 require 'ddtrace/sync_writer'
+require 'ddtrace/contrib/sampling'
 require 'ddtrace/contrib/sidekiq/ext'
 require 'resque'
 
@@ -9,13 +10,13 @@ module Resque
       # Uses Resque job hooks to create traces
       module ResqueJob
         def around_perform(*_)
-          pin = Pin.get_from(::Resque)
-          return yield unless pin && pin.tracer
-          pin.tracer.trace(Ext::SPAN_JOB, service: pin.service) do |span|
+          return yield unless datadog_configuration && tracer
+
+          tracer.trace(Ext::SPAN_JOB, span_options) do |span|
             span.resource = name
-            span.span_type = pin.app_type
+            span.span_type = Datadog::Ext::AppTypes::WORKER
+            Contrib::Sampling.set_event_sample_rate(span, datadog_configuration[:event_sample_rate])
             yield
-            span.service = pin.service
           end
         end
 
@@ -28,21 +29,45 @@ def on_failure_shutdown_tracer(*_)
         end
 
         def shutdown_tracer_when_forked!
-          pin = Datadog::Pin.get_from(Resque)
-          pin.tracer.shutdown! if pin && pin.tracer && pin.config && pin.config[:forked]
+          tracer.shutdown! if forked?
+        end
+
+        private
+
+        def forked?
+          pin = Datadog::Pin.get_from(::Resque)
+          return false unless pin
+          pin.config[:forked] == true
+        end
+
+        def span_options
+          { service: datadog_configuration[:service_name] }
+        end
+
+        def tracer
+          datadog_configuration.tracer
+        end
+
+        def datadog_configuration
+          Datadog.configuration[:resque]
         end
       end
     end
   end
 end
 
 Resque.after_fork do
-  # get the current tracer
-  pin = Datadog::Pin.get_from(Resque)
-  next unless pin && pin.tracer
-  pin.config ||= {}
-  pin.config[:forked] = true
-
-  # clean the state so no CoW happens
-  pin.tracer.provider.context = nil
+  configuration = Datadog.configuration[:resque]
+  next if configuration.nil?
+
+  # Add a pin, marking the job as forked.
+  # Used to trigger shutdown in forks for performance reasons.
+  Datadog::Pin.new(
+    configuration[:service_name],
+    config: { forked: true }
+  ).onto(::Resque)
+
+  # Clean the state so no CoW happens
+  next if configuration[:tracer].nil?
+  configuration[:tracer].provider.context = nil
 end

lib/ddtrace/contrib/sampling.rb

@@ -0,0 +1,14 @@
+require 'ddtrace/ext/priority'
+
+module Datadog
+  module Contrib
+    # Defines sampling behavior for integrations
+    module Sampling
+      module_function
+
+      def set_event_sample_rate(span, sample_rate)
+        span.set_metric(Datadog::Ext::Priority::TAG_EVENT_SAMPLE_RATE, sample_rate) unless sample_rate.nil?
+      end
+    end
+  end
+end

lib/ddtrace/contrib/shoryuken/tracer.rb

@@ -1,17 +1,20 @@
+require 'ddtrace/contrib/sampling'
+
 module Datadog
   module Contrib
     module Shoryuken
       # Tracer is a Shoryuken server-side middleware which traces executed jobs
       class Tracer
         def initialize(options = {})
-          @tracer = options[:tracer] || Datadog.configuration[:shoryuken][:tracer]
-          @shoryuken_service = options[:service_name] || Datadog.configuration[:shoryuken][:service_name]
+          @tracer = options[:tracer] || configuration[:tracer]
+          @shoryuken_service = options[:service_name] || configuration[:service_name]
           set_service_info(@shoryuken_service)
         end
 
         def call(worker_instance, queue, sqs_msg, body)
           @tracer.trace(Ext::SPAN_JOB, service: @shoryuken_service, span_type: Datadog::Ext::AppTypes::WORKER) do |span|
             span.resource = worker_instance.class.name
+            Contrib::Sampling.set_event_sample_rate(span, configuration[:event_sample_rate])
             span.set_tag(Ext::TAG_JOB_ID, sqs_msg.message_id)
             span.set_tag(Ext::TAG_JOB_QUEUE, queue)
             span.set_tag(Ext::TAG_JOB_ATTRIBUTES, sqs_msg.attributes) if sqs_msg.respond_to?(:attributes)
@@ -23,8 +26,12 @@ def call(worker_instance, queue, sqs_msg, body)
 
         private
 
+        def configuration
+          Datadog.configuration[:shoryuken]
+        end
+
         def set_service_info(service)
-          return if @tracer.services[service]
+          return if @tracer.nil? || @tracer.services[service]
           @tracer.set_service_info(
             service,
             Ext::APP,

lib/ddtrace/contrib/sidekiq/client_tracer.rb

@@ -1,4 +1,5 @@
 require 'ddtrace/contrib/sidekiq/tracing'
+require 'ddtrace/contrib/sampling'
 
 module Datadog
   module Contrib
@@ -9,26 +10,31 @@ class ClientTracer
 
         def initialize(options = {})
           super
-          @sidekiq_service = options[:client_service_name] || Datadog.configuration[:sidekiq][:client_service_name]
+          @sidekiq_service = options[:client_service_name] || configuration[:client_service_name]
+          set_service_info(@sidekiq_service)
         end
 
         # Client middleware arguments are documented here:
         #   https://github.com/mperham/sidekiq/wiki/Middleware#client-middleware
         def call(worker_class, job, queue, redis_pool)
-          service = @sidekiq_service
-          set_service_info(service)
-
           resource = job_resource(job)
 
-          @tracer.trace(Ext::SPAN_PUSH, service: service) do |span|
+          @tracer.trace(Ext::SPAN_PUSH, service: @sidekiq_service) do |span|
             span.resource = resource
+            Contrib::Sampling.set_event_sample_rate(span, configuration[:event_sample_rate])
             span.set_tag(Ext::TAG_JOB_ID, job['jid'])
             span.set_tag(Ext::TAG_JOB_QUEUE, job['queue'])
             span.set_tag(Ext::TAG_JOB_WRAPPER, job['class']) if job['wrapped']
 
             yield
           end
         end
+
+        private
+
+        def configuration
+          Datadog.configuration[:sidekiq]
+        end
       end
     end
   end