lib/sidekiq/tracer/server_middleware.rb in sidekiq-instrumentation-1.2.1 vs lib/sidekiq/tracer/server_middleware.rb in sidekiq-instrumentation-1.2.2.ejacobs1

@@ -3,15 +3,16 @@
 module Sidekiq
   module Tracer
     class ServerMiddleware
       include Commons
 
-      attr_reader :tracer, :active_span
+      attr_reader :tracer, :active_span, :after_trace
 
-      def initialize(tracer:, active_span:)
+      def initialize(tracer:, active_span:, after_trace:)
         @tracer = tracer
         @active_span = active_span
+        @after_trace = after_trace
       end
 
       # rubocop:disable Metrics/MethodLength
       def call(_worker, job, _queue)
         parent_span_context = extract(job)
@@ -25,9 +26,11 @@
           begin
             yield
           rescue StandardError => e
             tag_errors(scope.span, e) if scope.span
             raise
+          ensure
+            after_trace&.call(scope.span) if scope.span
           end
         end
       end
       # rubocop:enable Metrics/MethodLength