From 0f252622c23c4513237434c536e04fe666c3a35c Mon Sep 17 00:00:00 2001
From: Neel Shah <neelshah.sa@gmail.com>
Date: Thu, 28 Sep 2023 15:07:32 +0200
Subject: [PATCH] Fix Sidekiq tracing headers not being overwritten in case of
 schedules and retries
 [#2118](https://github.com/getsentry/sentry-ruby/pull/2118)

https://github.com/getsentry/sentry-ruby/pull/1774 added
`SentryContextClientMiddleware` to the server middleware chain too.

Sidekiq's scheduler pushes to the client on the server again for
schedules and retries which causes our trace propagation to be broken
for this case.
https://github.com/sidekiq/sidekiq/blob/aadc77a172f1490fb141c7936d2801ca3af925ef/lib/sidekiq/scheduled.rb#L39

Prioritize taking the trace propagation headers from the job whenever
they exist.
---
 CHANGELOG.md                                  |  1 +
 .../sidekiq/sentry_context_middleware.rb      |  2 +-
 .../sidekiq/sentry_context_middleware_spec.rb | 19 +++++++++++++++++++
 3 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7977f6e3b..1388a7da2 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -13,6 +13,7 @@
 - Always send envelope trace header from dynamic sampling context [#2113](https://github.com/getsentry/sentry-ruby/pull/2113)
 - Improve `TestHelper`'s setup/teardown helpers ([#2116](https://github.com/getsentry/sentry-ruby/pull/2116))
   - Fixes [#2103](https://github.com/getsentry/sentry-ruby/issues/2103)
+- Fix Sidekiq tracing headers not being overwritten in case of schedules and retries [#2118](https://github.com/getsentry/sentry-ruby/pull/2118)
 
 ## 5.11.0
 
diff --git a/sentry-sidekiq/lib/sentry/sidekiq/sentry_context_middleware.rb b/sentry-sidekiq/lib/sentry/sidekiq/sentry_context_middleware.rb
index fcb6e0ada..b6cde1ea5 100644
--- a/sentry-sidekiq/lib/sentry/sidekiq/sentry_context_middleware.rb
+++ b/sentry-sidekiq/lib/sentry/sidekiq/sentry_context_middleware.rb
@@ -59,7 +59,7 @@ def call(_worker_class, job, _queue, _redis_pool)
 
         user = Sentry.get_current_scope.user
         job["sentry_user"] = user unless user.empty?
-        job["trace_propagation_headers"] = Sentry.get_trace_propagation_headers
+        job["trace_propagation_headers"] ||= Sentry.get_trace_propagation_headers
         yield
       end
     end
diff --git a/sentry-sidekiq/spec/sentry/sidekiq/sentry_context_middleware_spec.rb b/sentry-sidekiq/spec/sentry/sidekiq/sentry_context_middleware_spec.rb
index 3a2b48de0..4a8542291 100644
--- a/sentry-sidekiq/spec/sentry/sidekiq/sentry_context_middleware_spec.rb
+++ b/sentry-sidekiq/spec/sentry/sidekiq/sentry_context_middleware_spec.rb
@@ -129,5 +129,24 @@
       expect(headers["sentry-trace"]).to eq(transaction.to_sentry_trace)
       expect(headers["baggage"]).to eq(transaction.to_baggage)
     end
+
+    # sidekiq pushes the same job to the queue again from the server for schedules and retries
+    it "keeps the same trace_propagation_headers linked to the transaction when queued multiple times" do
+      client.push('queue' => 'default', 'class' => HappyWorker, 'args' => [])
+
+      # push without span transaction to simulate pushing on the server
+      Sentry.get_current_scope.clear
+      client.push(queue.first.item)
+
+      q = queue.to_a
+      expect(q.size).to be(2)
+      first_headers = q.first["trace_propagation_headers"]
+      expect(first_headers["sentry-trace"]).to eq(transaction.to_sentry_trace)
+      expect(first_headers["baggage"]).to eq(transaction.to_baggage)
+
+      second_headers = q.second["trace_propagation_headers"]
+      expect(second_headers["sentry-trace"]).to eq(transaction.to_sentry_trace)
+      expect(second_headers["baggage"]).to eq(transaction.to_baggage)
+    end
   end
 end