Add sucker_punch integration

Appraisals

@@ -121,6 +121,7 @@ if RUBY_VERSION >= '2.2.2' && RUBY_PLATFORM != 'java'
     gem 'activerecord'
     gem 'sidekiq'
     gem 'aws-sdk'
+    gem 'sucker_punch'
   end
 else
   appraise 'contrib-old' do
@@ -134,5 +135,6 @@ else
     gem 'activerecord', '3.2.22.5'
     gem 'sidekiq', '4.0.0'
     gem 'aws-sdk', '~> 2.0'
+    gem 'sucker_punch'
   end
 end

Rakefile

@@ -51,7 +51,7 @@ namespace :test do
     t.test_files = FileList['test/contrib/rails/**/*disable_env*_test.rb']
   end
 
-  [:elasticsearch, :http, :redis, :sinatra, :sidekiq, :rack, :faraday, :grape, :aws].each do |contrib|
+  [:elasticsearch, :http, :redis, :sinatra, :sidekiq, :rack, :faraday, :grape, :aws, :sucker_punch].each do |contrib|
     Rake::TestTask.new(contrib) do |t|
       t.libs << %w[test lib]
       t.test_files = FileList["test/contrib/#{contrib}/*_test.rb"]

gemfiles/contrib.gemfile

@@ -13,5 +13,6 @@ gem "sqlite3"
 gem "activerecord"
 gem "sidekiq"
 gem "aws-sdk"
+gem "sucker_punch"
 
 gemspec path: "../"

gemfiles/contrib_old.gemfile

@@ -12,5 +12,6 @@ gem "sqlite3"
 gem "activerecord", "3.2.22.5"
 gem "sidekiq", "4.0.0"
 gem "aws-sdk", "~> 2.0"
+gem "sucker_punch"
 
 gemspec path: "../"

lib/ddtrace/contrib/sucker_punch/exception_handler.rb

@@ -0,0 +1,26 @@
+require 'sucker_punch'
+
+module Datadog
+  module Contrib
+    module SuckerPunch
+      # Patches `sucker_punch` exception handling
+      module ExceptionHandler
+        METHOD = ->(e, *) { raise(e) }
+
+        module_function
+
+        def patch!
+          ::SuckerPunch.class_eval do
+            class << self
+              alias_method :__exception_handler, :exception_handler
+
+              def exception_handler
+                ::Datadog::Contrib::SuckerPunch::ExceptionHandler::METHOD
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

lib/ddtrace/contrib/sucker_punch/instrumentation.rb

@@ -0,0 +1,60 @@
+require 'sucker_punch'
+
+module Datadog
+  module Contrib
+    module SuckerPunch
+      # Defines instrumentation patches for the `sucker_punch` gem
+      module Instrumentation
+        module_function
+
+        # rubocop:disable Metrics/MethodLength
+        def patch!
+          ::SuckerPunch::Job::ClassMethods.class_eval do
+            alias_method :__run_perform_without_datadog, :__run_perform
+            def __run_perform(*args)
+              pin = Datadog::Pin.get_from(::SuckerPunch)
+              pin.tracer.provider.context = Datadog::Context.new
+
+              __with_instrumentation('sucker_punch.perform') do |span|
+                span.resource = "PROCESS #{self}"
+                __run_perform_without_datadog(*args)
+              end
+            rescue => e
+              ::SuckerPunch.__exception_handler.call(e, self, args)
+            end
+
+            alias_method :__perform_async, :perform_async
+            def perform_async(*args)
+              __with_instrumentation('sucker_punch.perform_async') do |span|
+                span.resource = "ENQUEUE #{self}"
+                __perform_async(*args)
+              end
+            end
+
+            alias_method :__perform_in, :perform_in
+            def perform_in(interval, *args)
+              __with_instrumentation('sucker_punch.perform_in') do |span|
+                span.resource = "ENQUEUE #{self}"
+                span.set_tag('sucker_punch.perform_in', interval)
+                __perform_in(interval, *args)
+              end
+            end
+
+            private
+
+            def __with_instrumentation(name)
+              pin = Datadog::Pin.get_from(::SuckerPunch)
+
+              pin.tracer.trace(name) do |span|
+                span.service = pin.service
+                span.span_type = pin.app_type
+                span.set_tag('sucker_punch.queue', to_s)
+                yield span
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

lib/ddtrace/contrib/sucker_punch/patcher.rb

@@ -0,0 +1,50 @@
+module Datadog
+  module Contrib
+    module SuckerPunch
+      SERVICE = 'sucker_punch'.freeze
+      COMPATIBLE_WITH = Gem::Version.new('2.0.0')
+
+      # Responsible for hooking the instrumentation into `sucker_punch`
+      module Patcher
+        @patched = false
+
+        module_function
+
+        def patch
+          return @patched if patched? || !compatible?
+
+          require 'ddtrace/ext/app_types'
+          require_relative 'exception_handler'
+          require_relative 'instrumentation'
+
+          add_pin!
+          ExceptionHandler.patch!
+          Instrumentation.patch!
+
+          @patched = true
+        rescue => e
+          Datadog::Tracer.log.error("Unable to apply SuckerPunch integration: #{e}")
+          @patched
+        end
+
+        def patched?
+          @patched
+        end
+
+        def compatible?
+          return unless defined?(::SuckerPunch::VERSION)
+
+          Gem::Version.new(::SuckerPunch::VERSION) >= COMPATIBLE_WITH
+        end
+
+        def add_pin!
+          Pin.new(SERVICE, app_type: Ext::AppTypes::WORKER).tap do |pin|
+            pin.onto(::SuckerPunch)
+          end
+        end
+
+        private_class_method :compatible?, :add_pin!
+      end
+    end
+  end
+end

lib/ddtrace/monkey.rb

@@ -10,6 +10,7 @@
 require 'ddtrace/contrib/redis/patcher'
 require 'ddtrace/contrib/http/patcher'
 require 'ddtrace/contrib/aws/patcher'
+require 'ddtrace/contrib/sucker_punch/patcher'
 
 module Datadog
   # Monkey is used for monkey-patching 3rd party libs.
@@ -22,6 +23,7 @@ module Monkey
       grape: true,
       faraday: true,
       aws: true,
+      sucker_punch: true,
       active_record: false
     }
     # Patchers should expose 2 methods:
@@ -35,6 +37,7 @@ module Monkey
                   grape: Datadog::Contrib::Grape::Patcher,
                   faraday: Datadog::Contrib::Faraday::Patcher,
                   aws: Datadog::Contrib::Aws::Patcher,
+                  sucker_punch: Datadog::Contrib::SuckerPunch::Patcher,
                   active_record: Datadog::Contrib::ActiveRecord::Patcher }
     @mutex = Mutex.new
 

test/contrib/sucker_punch/dummy_worker.rb

@@ -0,0 +1,9 @@
+require 'sucker_punch'
+
+class DummyWorker
+  include ::SuckerPunch::Job
+
+  def perform(action = :none)
+    1 / 0 if action == :fail
+  end
+end

test/contrib/sucker_punch/patcher_test.rb

@@ -0,0 +1,93 @@
+require 'helper'
+require 'sucker_punch'
+require 'ddtrace'
+require_relative 'dummy_worker'
+
+module Datadog
+  module Contrib
+    module SuckerPunch
+      class PatcherTest < Minitest::Test
+        def setup
+          Monkey.patch_module(:sucker_punch)
+          ::SuckerPunch::Queue.clear
+          ::SuckerPunch::RUNNING.make_true
+
+          @tracer = enable_test_tracer!
+        end
+
+        def test_two_spans_per_job
+          # One span when pushing to the queue
+          # One span for the job execution itself
+          ::DummyWorker.perform_async
+          try_wait_until { all_spans.length == 2 }
+          assert_equal(2, all_spans.length)
+        end
+
+        def test_successful_job
+          ::DummyWorker.perform_async
+          try_wait_until { all_spans.length == 2 }
+
+          span = all_spans.find { |s| s.resource[/PROCESS/] }
+          assert_equal('sucker_punch', span.service)
+          assert_equal('sucker_punch.perform', span.name)
+          assert_equal('PROCESS DummyWorker', span.resource)
+          assert_equal('DummyWorker', span.get_tag('sucker_punch.queue'))
+          refute_equal(Ext::Errors::STATUS, span.status)
+        end
+
+        def test_failed_job
+          ::DummyWorker.perform_async(:fail)
+          try_wait_until { all_spans.length == 2 }
+
+          span = all_spans.find { |s| s.resource[/PROCESS/] }
+          assert_equal('sucker_punch', span.service)
+          assert_equal('sucker_punch.perform', span.name)
+          assert_equal('PROCESS DummyWorker', span.resource)
+          assert_equal('DummyWorker', span.get_tag('sucker_punch.queue'))
+          assert_equal(Ext::Errors::STATUS, span.status)
+          assert_equal('ZeroDivisionError', span.get_tag(Ext::Errors::TYPE))
+          assert_equal('divided by 0', span.get_tag(Ext::Errors::MSG))
+        end
+
+        def test_async_enqueueing
+          ::DummyWorker.perform_async
+          try_wait_until { all_spans.any? }
+
+          span = all_spans.find { |s| s.resource[/ENQUEUE/] }
+          assert_equal('sucker_punch', span.service)
+          assert_equal('sucker_punch.perform_async', span.name)
+          assert_equal('ENQUEUE DummyWorker', span.resource)
+          assert_equal('DummyWorker', span.get_tag('sucker_punch.queue'))
+        end
+
+        def test_delayed_enqueueing
+          ::DummyWorker.perform_in(0)
+          try_wait_until { all_spans.any? }
+
+          span = all_spans.find { |s| s.resource[/ENQUEUE/] }
+          assert_equal('sucker_punch', span.service)
+          assert_equal('sucker_punch.perform_in', span.name)
+          assert_equal('ENQUEUE DummyWorker', span.resource)
+          assert_equal('DummyWorker', span.get_tag('sucker_punch.queue'))
+          assert_equal('0', span.get_tag('sucker_punch.perform_in'))
+        end
+
+        private
+
+        attr_reader :tracer
+
+        def all_spans
+          tracer.writer.spans(:keep)
+        end
+
+        def enable_test_tracer!
+          get_test_tracer.tap { |tracer| pin.tracer = tracer }
+        end
+
+        def pin
+          ::SuckerPunch.datadog_pin
+        end
+      end
+    end
+  end
+end

test/helper.rb

@@ -84,10 +84,10 @@ def write(trace, services)
     end
   end
 
-  def spans
+  def spans(action = :clear)
     @mutex.synchronize do
       spans = @spans
-      @spans = []
+      @spans = [] if action == :clear
       spans.flatten!
       # sort the spans to avoid test flakiness
       spans.sort! do |a, b|
@@ -218,3 +218,14 @@ def test_repeat
   return 300 if RUBY_PLATFORM == 'java'
   30
 end
+
+def try_wait_until(options = {})
+  attempts = options.fetch(:attempts, 10)
+  backoff = options.fetch(:backoff, 0.1)
+
+  loop do
+    break if attempts <= 0 || yield
+    sleep(backoff)
+    attempts -= 1
+  end
+end