refactor(sidekiq): isolate ActiveJob support

Author: solnic

sentry-sidekiq/lib/sentry-sidekiq.rb

@@ -20,19 +20,11 @@ def self.job_extension
     end
 
     Sentry.register_job_extension(job_extension)
-
-    if defined?(::Rails::Railtie)
-      class Railtie < ::Rails::Railtie
-        config.after_initialize do
-          next unless Sentry.initialized? && defined?(::Sentry::Rails)
-
-          Sentry.configuration.rails.skippable_job_adapters << JobExtension::ACTIVE_JOB_ADAPTER
-        end
-      end
-    end
   end
 end
 
+require "sentry/sidekiq/railtie" if defined?(::Rails::Railtie)
+
 Sidekiq.configure_server do |config|
   config.error_handlers << Sentry::Sidekiq::ErrorHandler.new
   config.server_middleware do |chain|

sentry-sidekiq/lib/sentry/sidekiq/active_job_support/job_extension.rb

@@ -0,0 +1,106 @@
+# frozen_string_literal: true
+
+module Sentry
+  module Sidekiq
+    module ActiveJobSupport
+      module JobExtension
+        ACTIVE_JOB_ADAPTER = "ActiveJob::QueueAdapters::SidekiqAdapter"
+        SIDEKIQ_JOB_KEY = "_sentry_sidekiq_job"
+
+        def handles?(job)
+          super || active_job?(job)
+        end
+
+        def scope_data(job:)
+          return super unless active_job?(job)
+
+          sidekiq_job = handed_over_sidekiq_job(job)
+          sidekiq_job ? super(job: sidekiq_job) : {}
+        end
+
+        def messaging_data(job:)
+          return super unless active_job?(job)
+
+          sidekiq_job = handed_over_sidekiq_job(job)
+          return {} unless sidekiq_job
+
+          super(job: sidekiq_job).merge(body_size: body_size(sidekiq_job)).compact
+        end
+
+        def retry_info(job:)
+          return super unless active_job?(job)
+
+          sidekiq_job = handed_over_sidekiq_job(job)
+          super(job: sidekiq_job) if sidekiq_job
+        end
+
+        def extract_trace_headers(job:, payload:)
+          return super if sidekiq_job?(payload)
+
+          sidekiq_job = sidekiq_job_in(payload)
+          super(job: job, payload: sidekiq_job) if sidekiq_job
+        end
+
+        def transaction_name(job:)
+          active_job?(job) ? "#{self.class::SIDEKIQ_NAME}/#{job.class.name}" : super
+        end
+
+        def span_origin(job:)
+          active_job?(job) ? self.class::SPAN_ORIGIN : super
+        end
+
+        def integration(job:)
+          active_job?(job) ? Sentry::Sidekiq : super
+        end
+
+        def report_retried_exception?(job:, exception:)
+          active_job?(job) ? false : super
+        end
+
+        def performed_by_active_job?(sidekiq_job)
+          !!sidekiq_job["wrapped"] && active_job_owns_adapter?
+        end
+
+        def active_job_args(sidekiq_job, scope:)
+          payload, *arguments = sidekiq_job["args"]
+          return sidekiq_job["args"] unless payload.is_a?(Hash)
+
+          sentry_data = (payload["_sentry"] || {}).merge("delivery_queue" => sidekiq_job["queue"])
+          sentry_data["execution_scope_id"] = scope.object_id
+          [payload.merge(
+            SIDEKIQ_JOB_KEY => sidekiq_job,
+            "_sentry" => sentry_data
+          ), *arguments]
+        end
+
+        private
+
+        def active_job?(job)
+          return false unless defined?(::ActiveJob::Base) && job.is_a?(::ActiveJob::Base)
+
+          job.class.queue_adapter.class.name == ACTIVE_JOB_ADAPTER
+        end
+
+        def active_job_owns_adapter?
+          return false unless defined?(::Sentry::Rails::ActiveJobExtensions)
+          return false unless ::Sentry::Rails.respond_to?(:active_job_owns_adapter?)
+
+          ::Sentry::Rails.active_job_owns_adapter?(ACTIVE_JOB_ADAPTER)
+        end
+
+        def handed_over_sidekiq_job(job)
+          return unless active_job?(job) && job.respond_to?(:_sentry_job_data)
+
+          sidekiq_job_in(job._sentry_job_data)
+        end
+
+        def sidekiq_job_in(active_job_payload)
+          return unless active_job_payload.is_a?(Hash)
+
+          sidekiq_job = active_job_payload[SIDEKIQ_JOB_KEY]
+          sidekiq_job if sidekiq_job?(sidekiq_job)
+        end
+      end
+    end
+  end
+end

sentry-sidekiq/lib/sentry/sidekiq/active_job_support/server_middleware.rb

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+module Sentry
+  module Sidekiq
+    module ActiveJobSupport
+      module ServerMiddleware
+        private
+
+        def execution_trace_headers(job, sidekiq_job)
+          if job_extension.performed_by_active_job?(job) && job["args"].first.is_a?(Hash)
+            headers = job["args"].first.dig("_sentry", "trace_propagation_headers")
+            return headers if headers && !headers.empty?
+          end
+
+          super
+        end
+
+        def perform_job(job, sidekiq_job, scope:)
+          return super unless job_extension.performed_by_active_job?(job)
+
+          args = job["args"]
+          begin
+            job["args"] = job_extension.active_job_args(sidekiq_job, scope: scope)
+            super
+          ensure
+            job["args"] = args
+          end
+        end
+      end
+    end
+  end
+end

sentry-sidekiq/lib/sentry/sidekiq/job_extension.rb

@@ -2,14 +2,15 @@
 
 require "sentry/sidekiq/retry_policy"
 require "sentry/sidekiq/context_filter"
+require "sentry/sidekiq/active_job_support/job_extension"
 
 module Sentry
   module Sidekiq
     class JobExtension < Sentry::JobExtensions::Extension
+      prepend ActiveJobSupport::JobExtension
+
       SPAN_ORIGIN = "auto.queue.sidekiq"
       SIDEKIQ_NAME = "Sidekiq"
-      ACTIVE_JOB_ADAPTER = "ActiveJob::QueueAdapters::SidekiqAdapter"
-      SIDEKIQ_JOB_KEY = "_sentry_sidekiq_job"
       ENQUEUED_AT_IN_MILLISECONDS = ::Gem::Version.new(::Sidekiq::VERSION) >= ::Gem::Version.new("8.0.0")
 
       def initialize(retry_policy: RetryPolicy.build)
@@ -18,23 +19,7 @@ def initialize(retry_policy: RetryPolicy.build)
       end
 
       def handles?(job)
-        sidekiq_job?(job) || active_job?(job)
-      end
-
-      def performed_by_active_job?(sidekiq_job)
-        !!sidekiq_job["wrapped"] && active_job_owns_adapter?
-      end
-
-      def active_job_args(sidekiq_job, scope:)
-        payload, *arguments = sidekiq_job["args"]
-        return sidekiq_job["args"] unless payload.is_a?(Hash)
-
-        sentry_data = (payload["_sentry"] || {}).merge("delivery_queue" => sidekiq_job["queue"])
-        sentry_data["execution_scope_id"] = scope.object_id
-        [payload.merge(
-          SIDEKIQ_JOB_KEY => sidekiq_job,
-          "_sentry" => sentry_data
-        ), *arguments]
+        sidekiq_job?(job)
       end
 
       def error_context_transaction_name(context)
@@ -45,28 +30,25 @@ def error_context_transaction_name(context)
       end
 
       def scope_data(job:)
-        sidekiq_job = sidekiq_job_for(job)
-        return {} unless sidekiq_job
+        return {} unless sidekiq_job?(job)
 
         data = {
-          contexts: { sidekiq: ContextFilter.new(sidekiq_job).filtered },
-          tags: { queue: sidekiq_job["queue"], jid: sidekiq_job["jid"] }.merge(job_tags(sidekiq_job))
+          contexts: { sidekiq: ContextFilter.new(job).filtered },
+          tags: { queue: job["queue"], jid: job["jid"] }.merge(job_tags(job))
         }
-        user = sidekiq_job["sentry_user"]
+        user = job["sentry_user"]
         data[:user] = user if user.is_a?(Hash)
         data
       end
 
       def messaging_data(job:)
-        sidekiq_job = sidekiq_job_for(job)
-        return {} unless sidekiq_job
+        return {} unless sidekiq_job?(job)
 
         {
           system: name,
-          message_id: sidekiq_job["jid"],
-          destination: sidekiq_job["queue"],
-          receive_latency: receive_latency(sidekiq_job),
-          body_size: active_job?(job) ? body_size(sidekiq_job) : nil
+          message_id: job["jid"],
+          destination: job["queue"],
+          receive_latency: receive_latency(job)
         }.compact
       end
 
@@ -75,10 +57,9 @@ def consumer_messaging_data(job)
       end
 
       def retry_info(job:)
-        sidekiq_job = sidekiq_job_for(job)
-        return unless sidekiq_job
+        return unless sidekiq_job?(job)
 
-        { count: @retry_policy.retry_count(sidekiq_job) || 0, will_retry: @retry_policy.will_retry?(sidekiq_job) }
+        { count: @retry_policy.retry_count(job) || 0, will_retry: @retry_policy.will_retry?(job) }
       end
 
       def inject_trace_headers(job:, payload:, headers:)
@@ -89,33 +70,20 @@ def inject_trace_headers(job:, payload:, headers:)
       end
 
       def extract_trace_headers(job:, payload:)
-        sidekiq_job = sidekiq_job_for_payload(payload)
-        return unless sidekiq_job
+        return unless sidekiq_job?(payload)
 
-        sidekiq_job["trace_propagation_headers"] || {}
+        payload["trace_propagation_headers"] || {}
       end
 
       def transaction_name(job:)
-        name = active_job?(job) ? job.class.name : job_class(job)
+        name = job_class(job)
         "#{SIDEKIQ_NAME}/#{name}" if name
       end
 
-      def span_origin(job:)
-        SPAN_ORIGIN if active_job?(job)
-      end
-
-      def integration(job:)
-        Sentry::Sidekiq if active_job?(job)
-      end
-
       def producer(job:)
         @producer ||= Sentry::JobExtensions::Producer::Delegated.new
       end
 
-      def report_retried_exception?(job:, exception:)
-        false if active_job?(job)
-      end
-
       private
 
       def sidekiq_job?(job)
@@ -126,40 +94,6 @@ def job_class(job)
         job["wrapped"] || job["class"] if job.is_a?(Hash)
       end
 
-      def active_job?(job)
-        return false unless defined?(::ActiveJob::Base) && job.is_a?(::ActiveJob::Base)
-
-        job.class.queue_adapter.class.name == ACTIVE_JOB_ADAPTER
-      end
-
-      def active_job_owns_adapter?
-        return false unless defined?(::Sentry::Rails::ActiveJobExtensions)
-        return false unless ::Sentry::Rails.respond_to?(:active_job_owns_adapter?)
-
-        ::Sentry::Rails.active_job_owns_adapter?(ACTIVE_JOB_ADAPTER)
-      end
-
-      def sidekiq_job_for(job)
-        sidekiq_job?(job) ? job : handed_over_sidekiq_job(job)
-      end
-
-      def handed_over_sidekiq_job(job)
-        return unless active_job?(job) && job.respond_to?(:_sentry_job_data)
-
-        sidekiq_job_in(job._sentry_job_data)
-      end
-
-      def sidekiq_job_for_payload(payload)
-        sidekiq_job?(payload) ? payload : sidekiq_job_in(payload)
-      end
-
-      def sidekiq_job_in(active_job_payload)
-        return unless active_job_payload.is_a?(Hash)
-
-        sidekiq_job = active_job_payload[SIDEKIQ_JOB_KEY]
-        sidekiq_job if sidekiq_job?(sidekiq_job)
-      end
-
       def job_tags(job)
         Array(job["tags"]).each_with_object({}) { |name, tags_hash| tags_hash[:"sidekiq.#{name}"] = true }
       end

sentry-sidekiq/lib/sentry/sidekiq/railtie.rb

@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+module Sentry
+  module Sidekiq
+    class Railtie < ::Rails::Railtie
+      config.after_initialize do
+        next unless Sentry.initialized? && defined?(::Sentry::Rails)
+
+        Sentry.configuration.rails.skippable_job_adapters << ActiveJobSupport::JobExtension::ACTIVE_JOB_ADAPTER
+      end
+    end
+  end
+end