refactor(sidekiq): extract shared building blocks

sentry-rails/lib/sentry/rails.rb

@@ -13,5 +13,28 @@ module Sentry
   module Rails
     extend Integrable
     register_integration name: "rails", version: Sentry::Rails::VERSION
+
+    # Returns whether sentry-rails' ActiveJob integration reports jobs run
+    # through the given queue adapter.
+    #
+    # When it returns true, the integration captures errors from those jobs,
+    # records their queue.process transactions and queue.publish spans,
+    # propagates traces, and consults extensions registered with
+    # {Sentry.register_job_extension}. A backend gem should then not report the
+    # same ActiveJob jobs itself.
+    #
+    # The result depends on {Sentry.initialized?} and the current configuration,
+    # so it can change after {Sentry.init}. Call it when a job is handled rather
+    # than caching the result at load time.
+    #
+    # @param adapter_class_name [String, Class] the ActiveJob queue adapter's
+    #   class name, such as "ActiveJob::QueueAdapters::SidekiqAdapter", or the
+    #   adapter class itself. It is compared against
+    #   +config.rails.skippable_job_adapters+.
+    # @return [Boolean] true if the SDK is initialized and the adapter is not
+    #   listed in +config.rails.skippable_job_adapters+, false otherwise
+    def self.active_job_owns_adapter?(adapter_class_name)
+      Sentry.initialized? && !Sentry.configuration.rails.skippable_job_adapters.include?(adapter_class_name.to_s)
+    end
   end
 end

sentry-rails/lib/sentry/rails/active_job.rb

@@ -13,6 +13,7 @@ module ActiveJobExtensions
 
       def self.prepended(base)
         base.attr_accessor :_sentry
+        base.attr_reader :_sentry_job_data
       end
 
       def perform_now
@@ -23,20 +24,25 @@ def perform_now
           SentryReporter.record(
             self,
             trace_headers: data["trace_propagation_headers"],
-            user: data["user"]
+            user: data["user"],
+            payload: _sentry_job_data
           ) { super }
         end
       end
 
       def serialize
         payload = super
         return payload if !Sentry.initialized? || already_supported_by_sentry_integration?
+        return payload if SentryReporter.instrumented_by_extension?(self)
 
         begin
           sentry_data = {}
           if Sentry.configuration.rails.active_job_propagate_traces
             headers = Sentry.get_trace_propagation_headers
-            sentry_data["trace_propagation_headers"] = headers if headers && !headers.empty?
+            if headers && !headers.empty?
+              sentry_data["trace_propagation_headers"] = headers
+              payload = SentryReporter.call_hook(self, :inject_trace_headers, payload: payload, headers: headers)
+            end
           end
 
           if Sentry.configuration.data_collection.user_info
@@ -58,14 +64,15 @@ def deserialize(job_data)
         return if !Sentry.initialized? || already_supported_by_sentry_integration?
 
         begin
+          @_sentry_job_data = job_data
           self._sentry = job_data[SENTRY_PAYLOAD_KEY]
         rescue StandardError => e
           Sentry.sdk_logger&.error("sentry-rails: failed to extract _sentry payload: #{e.class}: #{e.message}\n  #{Array(e.backtrace).first(5).join("\n  ")}")
         end
       end
 
       def already_supported_by_sentry_integration?
-        Sentry.configuration.rails.skippable_job_adapters.include?(self.class.queue_adapter.class.to_s)
+        !Sentry::Rails.active_job_owns_adapter?(self.class.queue_adapter.class)
       end
 
       class SentryReporter
@@ -78,13 +85,28 @@ class SentryReporter
         # normally derived from the adapter itself; this is only the last resort.
         MESSAGING_SYSTEM_FALLBACK = "activejob"
 
+        MESSAGING_DATA_ATTRIBUTES = {
+          system: Sentry::Span::DataConventions::MESSAGING_SYSTEM,
+          message_id: Sentry::Span::DataConventions::MESSAGING_MESSAGE_ID,
+          destination: Sentry::Span::DataConventions::MESSAGING_DESTINATION_NAME,
+          receive_latency: Sentry::Span::DataConventions::MESSAGING_MESSAGE_RECEIVE_LATENCY
+        }.freeze
+
         EVENT_HANDLERS = {
           "enqueue_retry.active_job" => :retry_handler
         }
 
         class << self
           include ErrorReporterContext
 
+          def job_extensions
+            @job_extensions ||= Sentry.job_extensions
+          end
+
+          def call_hook(job, hook, **arguments)
+            job_extensions.call(job_extensions.for(job), hook, job: job, **arguments)
+          end
+
           def producer_callback_registered?
             @producer_callback_registered ||= false
           end
@@ -93,9 +115,15 @@ def producer_callback_registered!
             @producer_callback_registered = true
           end
 
+          def instrumented_by_extension?(job)
+            call_hook(job, :instruments_enqueue?) == true
+          end
+
           def record_producer_span(job, &enqueue)
             return yield if !Sentry.initialized? || job.already_supported_by_sentry_integration?
 
+            return yield if instrumented_by_extension?(job)
+
             enqueued = false
 
             run_enqueue = lambda do
@@ -104,9 +132,11 @@ def record_producer_span(job, &enqueue)
             end
 
             begin
-              Sentry.with_child_span(op: "queue.publish", description: job.class.name) do |span|
+              description = call_hook(job, :publish_description) || job.class.name
+              Sentry.with_child_span(op: "queue.publish", description: description) do |span|
                 if span
-                  span.set_origin(SPAN_ORIGIN)
+                  span.set_origin(call_hook(job, :span_origin) || SPAN_ORIGIN)
+                  set_extension_messaging_data(span, job)
                   span.set_data(Sentry::Span::DataConventions::MESSAGING_SYSTEM, messaging_system(job))
                   span.set_data(Sentry::Span::DataConventions::MESSAGING_MESSAGE_ID, job.job_id)
                   span.set_data(Sentry::Span::DataConventions::MESSAGING_DESTINATION_NAME, job.queue_name)
@@ -131,15 +161,17 @@ def log_producer_span_error(e)
             Sentry.sdk_logger&.error("sentry-rails: failed to record producer span: #{e.class}: #{e.message}\n  #{Array(e.backtrace).first(5).join("\n  ")}")
           end
 
-          def record(job, trace_headers: nil, user: nil, &block)
+          def record(job, trace_headers: nil, user: nil, payload: nil, &block)
             original_hub = Sentry.get_current_hub_internal
             Sentry.clone_hub_to_current_thread
 
             Sentry.with_scope do |scope|
               begin
                 scope.set_user(user.transform_keys(&:to_sym)) if user && !user.empty?
 
-                scope.set_transaction_name(job.class.name, source: :task)
+                apply_extension_scope_data(scope, job)
+
+                scope.set_transaction_name(call_hook(job, :transaction_name) || job.class.name, source: :task)
                 scope.set_tags(queue: job.queue_name)
 
                 scope.set_contexts(active_job: {
@@ -153,10 +185,19 @@ def record(job, trace_headers: nil, user: nil, &block)
                   name: scope.transaction_name,
                   source: scope.transaction_source,
                   op: OP_NAME,
-                  origin: SPAN_ORIGIN
+                  origin: call_hook(job, :span_origin) || SPAN_ORIGIN
                 }
 
-                transaction = if trace_headers && !trace_headers.empty?
+                extracted_trace = false
+                if (trace_headers.nil? || trace_headers.empty?) && payload
+                  extracted_headers = call_hook(job, :extract_trace_headers, payload: payload)
+                  unless extracted_headers.nil?
+                    trace_headers = extracted_headers
+                    extracted_trace = true
+                  end
+                end
+
+                transaction = if extracted_trace || (trace_headers && !trace_headers.empty?)
                   continued = Sentry.continue_trace(trace_headers, **transaction_options)
                   Sentry.start_transaction(transaction: continued, **transaction_options)
                 else
@@ -165,6 +206,7 @@ def record(job, trace_headers: nil, user: nil, &block)
 
                 if transaction
                   set_messaging_data(transaction, job)
+                  set_extension_messaging_data(transaction, job)
                   scope.set_span(transaction)
                 end
 
@@ -174,7 +216,7 @@ def record(job, trace_headers: nil, user: nil, &block)
               rescue Exception => e # rubocop:disable Lint/RescueException
                 finish_sentry_transaction(transaction, 500)
 
-                capture_exception(job, e)
+                capture_exception(job, e) if report_exception?(job, e)
 
                 raise
               end
@@ -196,6 +238,20 @@ def set_messaging_data(transaction, job)
             end
           end
 
+          def apply_extension_scope_data(scope, job)
+            data = call_hook(job, :scope_data)
+            scope.set_contexts(data[:contexts]) if data[:contexts]
+            scope.set_tags(data[:tags]) if data[:tags]
+            scope.set_extras(data[:extra]) if data[:extra]
+            scope.set_user(data[:user]) if data[:user]
+          end
+
+          def set_extension_messaging_data(span, job)
+            call_hook(job, :messaging_data).each do |key, value|
+              span.set_data(MESSAGING_DATA_ATTRIBUTES.fetch(key), value)
+            end
+          end
+
           def messaging_system(job)
             name = job.class.queue_adapter_name if job.class.respond_to?(:queue_adapter_name)
             name = name.to_s
@@ -209,7 +265,26 @@ def messaging_system(job)
           # runs after the failed attempt bumped +executions+, so it observes the
           # same progression.
           def retry_count(job)
-            job.executions.to_i if job.respond_to?(:executions)
+            if (info = call_hook(job, :retry_info))
+              info[:count]
+            elsif job.respond_to?(:executions)
+              job.executions.to_i
+            end
+          end
+
+          def backend_will_retry?(job)
+            info = call_hook(job, :retry_info)
+            info ? info[:will_retry] : false
+          end
+
+          def report_exception?(job, exception)
+            decision = call_hook(job, :report_exception?, exception: exception)
+            decision.nil? ? !backend_will_retry?(job) : decision
+          end
+
+          def report_retried_exception?(job, exception)
+            decision = call_hook(job, :report_retried_exception?, exception: exception)
+            decision.nil? ? Sentry.configuration.rails.active_job_report_on_retry_error : decision
           end
 
           def compute_latency(job)
@@ -220,21 +295,49 @@ def compute_latency(job)
           end
 
           def capture_exception(job, e)
+            error_data = additive_error_data(
+              call_hook(job, :error_data, exception: e)
+            )
+
             options = {
-              extra: sentry_context(job),
-              tags: {
+              extra: error_data[:extra].merge(sentry_context(job)),
+              tags: error_data[:tags].merge(
                 job_id: job.job_id,
                 provider_job_id: job.provider_job_id
-              },
+              ),
               # Send synchronously: a worker process may exit before the async
               # background worker flushes its queue, which would drop the event.
               hint: { background: false }
             }
-            options[:contexts] = execution_context if Sentry.configuration.data_collection.queues
 
-            Sentry::Rails.capture_exception(e, **options)
+            contexts = error_data[:contexts]
+            contexts = contexts.merge(matching_keys(execution_context, contexts)) if Sentry.configuration.data_collection.queues
+            options[:contexts] = contexts if Sentry.configuration.data_collection.queues || !contexts.empty?
+
+            integration = call_hook(job, :integration) || Sentry::Rails
+            integration.capture_exception(e, **options)
+          end
+
+          def additive_error_data(data)
+            scope = Sentry.get_current_scope
+
+            {
+              contexts: matching_keys(data.fetch(:contexts, {}), scope.contexts).to_h do |name, context|
+                scope_context = scope.contexts.fetch(name, {})
+                context = matching_keys(context, scope_context)
+                [name, context.merge(scope_context)]
+              end,
+              tags: matching_keys(data.fetch(:tags, {}), scope.tags).reject { |key, _value| scope.tags.key?(key) },
+              extra: matching_keys(data.fetch(:extra, {}), scope.extra).reject { |key, _value| scope.extra.key?(key) }
+            }
+          end
+
+          def matching_keys(data, reference)
+            keys = reference.keys.to_h { |key| [key.to_s, key] }
+            data.transform_keys { |key| keys.fetch(key.to_s, key) }
           end
 
+
           def register_event_handlers
             EVENT_HANDLERS.each do |name, handler|
               subscribers << ActiveSupport::Notifications.subscribe(name) do |*args|
@@ -255,7 +358,7 @@ def retry_handler(*args)
             handle_error_event(*args) do |job, error|
               return if error.nil?
               return if !Sentry.initialized? || job.already_supported_by_sentry_integration?
-              return unless Sentry.configuration.rails.active_job_report_on_retry_error
+              return unless report_retried_exception?(job, error)
 
               capture_exception(job, error)
             end