Add trace context in Sidekiq to connect request transaction and worke…

…r in the one trace

lib/elastic_apm/spies/sidekiq.rb

@@ -29,7 +29,7 @@ class SidekiqSpy
       class Middleware
         def call(_worker, job, queue)
           name = SidekiqSpy.name_for(job)
-          transaction = ElasticAPM.start_transaction(name, 'Sidekiq')
+          transaction = ElasticAPM.start_transaction(name, 'Sidekiq', trace_context: get_trace_context_from(job))
           ElasticAPM.set_label(:queue, queue)
 
           yield
@@ -44,6 +44,35 @@ def call(_worker, job, queue)
         ensure
           ElasticAPM.end_transaction
         end
+
+        private
+
+        def get_trace_context_from(job)
+          return unless job['elastic_trace_context']
+
+          ElasticAPM::TraceContext.parse(metadata: job['elastic_trace_context'])
+        end
+      end
+
+      # @api private
+      class ClientMiddleware
+        def call(_worker_class, job, _queue, _redis_pool)
+          job['elastic_trace_context'] = elastic_trace_context
+          yield
+        end
+
+        private
+
+        def elastic_trace_context
+          return unless ElasticAPM.current_transaction
+
+          trace_context = ElasticAPM.current_transaction.trace_context
+
+          {
+            'traceparent' => trace_context.traceparent.to_header,
+            'tracestate' => trace_context.tracestate.to_header
+          }
+        end
       end
 
       def self.name_for(job)
@@ -63,6 +92,12 @@ def install_middleware
             chain.add Middleware
           end
         end
+
+        Sidekiq.configure_client do |config|
+          config.server_middleware do |chain|
+            chain.add ClientMiddleware
+          end
+        end
       end
 
       # @api private

spec/elastic_apm/spies/sidekiq_spec.rb

@@ -127,6 +127,36 @@ def perform
         expect(error.dig('exception', 'type')).to eq 'ZeroDivisionError'
       end
 
+      it 'creates a transaction with trace from job' do
+        with_agent do
+          transaction = ElasticAPM.start_transaction 'Test'
+          ElasticAPM.end_transaction
+
+          fake_trace_context = transaction.trace_context
+
+          allow_any_instance_of(ElasticAPM::Spies::SidekiqSpy::Middleware).to receive(:get_trace_context_from).and_return(
+            ElasticAPM::TraceContext.parse(
+              metadata: {
+                'traceparent' => fake_trace_context.traceparent.to_header,
+                'tracestate' => fake_trace_context.tracestate.to_header
+              }
+            )
+          )
+
+          Sidekiq::Testing.inline! do
+            HardWorker.perform_async
+          end
+        end
+
+        wait_for transactions: 2
+
+        fake_transaction, worker_transaction = @mock_intake.transactions
+
+        expect(worker_transaction).to_not be_nil
+        expect(worker_transaction['trace_id']).to eq(fake_transaction['trace_id'])
+        expect(worker_transaction['parent_id']).to eq(fake_transaction['id'])
+      end
+
       context 'ActiveJob', if: defined?(ActiveJob) do
         before :all do
           class ::ActiveJobbyJob < ActiveJob::Base
@@ -158,5 +188,32 @@ def perform
         end
       end
     end
+
+    describe Spies::SidekiqSpy::ClientMiddleware do
+      before :all do
+        Sidekiq.configure_client do |config|
+          config.client_middleware do |chain|
+            chain.add Spies::SidekiqSpy::ClientMiddleware
+          end
+        end
+      end
+
+      it 'adds trace context to job' do
+        current_transaction =
+          with_agent do
+            ElasticAPM.with_transaction do |transaction|
+              HardWorker.perform_async
+              transaction
+            end
+          end
+
+        current_job = Sidekiq::Queues['default'].first
+
+        trace_context = current_transaction.trace_context
+
+        expect(current_job['elastic_trace_context']['traceparent']).to eq(trace_context.traceparent.to_header)
+        expect(current_job['elastic_trace_context']['tracestate']).to eq(trace_context.tracestate.to_header)
+      end
+    end
   end
 end