Added: event_sample_rate tag to Shoryuken integration.

DataDog · Jan 2, 2019 · 0bea126 · 0bea126
1 parent 053b49f
commit 0bea126
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 5 deletions.
diff --git a/lib/ddtrace/contrib/shoryuken/tracer.rb b/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,

diff --git a/spec/ddtrace/contrib/shoryuken/tracer_spec.rb b/spec/ddtrace/contrib/shoryuken/tracer_spec.rb
@@ -1,21 +1,25 @@
 require 'spec_helper'
+require 'ddtrace/contrib/sampling_examples'
+
 require 'ddtrace'
 require 'shoryuken'
 
 RSpec.describe Datadog::Contrib::Shoryuken::Tracer do
   let(:tracer) { get_test_tracer }
-  let(:options) { { tracer: tracer } }
+  let(:configuration_options) { { tracer: tracer } }
   let(:spans) { tracer.writer.spans }
   let(:span) { spans.first }
 
   before do
     Shoryuken.worker_executor = Shoryuken::Worker::InlineExecutor
 
     Datadog.configure do |c|
-      c.use :shoryuken, options
+      c.use :shoryuken, configuration_options
     end
   end
 
+  after { Datadog.registry[:shoryuken].reset_configuration! }
+
   context 'when a Shoryuken::Worker class' do
     subject(:worker_class) do
       qn = queue_name
@@ -38,6 +42,12 @@ def perform(sqs_msg, body); end
           .and_call_original
       end
 
+      # TODO: These expectations do not work because Shoryuken doesn't run middleware in tests
+      #       https://github.com/phstc/shoryuken/issues/541
+      # it_behaves_like 'event sample rate' do
+      #   before { perform_async }
+      # end
+
       it do
         expect { perform_async }.to_not raise_error
         # TODO: These expectations do not work because Shoryuken doesn't run middleware in tests