Add latest changes from gitlab-org/gitlab@15-1-stable-eev15.1.0-rc42

4 files changed, 38 insertions, 13 deletions

diff --git a/lib/gitlab/sidekiq_middleware/duplicate_jobs/duplicate_job.rb b/lib/gitlab/sidekiq_middleware/duplicate_jobs/duplicate_job.rb
index 601c8d1c3cf..7533770e254 100644
--- a/lib/gitlab/sidekiq_middleware/duplicate_jobs/duplicate_job.rb
+++ b/lib/gitlab/sidekiq_middleware/duplicate_jobs/duplicate_job.rb
@@ -63,7 +63,7 @@ module Gitlab
           read_jid = nil
           read_wal_locations = {}
 
-          Sidekiq.redis do |redis|
+          with_redis do |redis|
             redis.multi do |multi|
               multi.set(idempotency_key, jid, ex: expiry, nx: true)
               read_wal_locations = check_existing_wal_locations!(multi, expiry)
@@ -81,7 +81,7 @@ module Gitlab
         def update_latest_wal_location!
           return unless job_wal_locations.present?
 
-          Sidekiq.redis do |redis|
+          with_redis do |redis|
             redis.multi do |multi|
               job_wal_locations.each do |connection_name, location|
                 multi.eval(
@@ -100,20 +100,19 @@ module Gitlab
           strong_memoize(:latest_wal_locations) do
             read_wal_locations = {}
 
-            Sidekiq.redis do |redis|
+            with_redis do |redis|
               redis.multi do |multi|
                 job_wal_locations.keys.each do |connection_name|
                   read_wal_locations[connection_name] = multi.lindex(wal_location_key(connection_name), 0)
                 end
               end
             end
-
             read_wal_locations.transform_values(&:value).compact
           end
         end
 
         def delete!
-          Sidekiq.redis do |redis|
+          with_redis do |redis|
             redis.multi do |multi|
               multi.del(idempotency_key, deduplicated_flag_key)
               delete_wal_locations!(multi)
@@ -140,7 +139,7 @@ module Gitlab
         def set_deduplicated_flag!(expiry = duplicate_key_ttl)
           return unless reschedulable?
 
-          Sidekiq.redis do |redis|
+          with_redis do |redis|
             redis.set(deduplicated_flag_key, DEDUPLICATED_FLAG_VALUE, ex: expiry, nx: true)
           end
         end
@@ -148,7 +147,7 @@ module Gitlab
         def should_reschedule?
           return false unless reschedulable?
 
-          Sidekiq.redis do |redis|
+          with_redis do |redis|
             redis.get(deduplicated_flag_key).present?
           end
         end
@@ -272,6 +271,18 @@ module Gitlab
         def reschedulable?
           !scheduled? && options[:if_deduplicated] == :reschedule_once
         end
+
+        def with_redis
+          if Feature.enabled?(:use_primary_and_secondary_stores_for_duplicate_jobs) ||
+             Feature.enabled?(:use_primary_store_as_default_for_duplicate_jobs)
+            # TODO: Swap for Gitlab::Redis::SharedState after store transition
+            # https://gitlab.com/gitlab-com/gl-infra/scalability/-/issues/923
+            Gitlab::Redis::DuplicateJobs.with { |redis| yield redis }
+          else
+            # Keep the old behavior intact if neither feature flag is turned on
+            Sidekiq.redis { |redis| yield redis }
+          end
+        end
       end
     end
   end
diff --git a/lib/gitlab/sidekiq_middleware/duplicate_jobs/strategies/until_executed.rb b/lib/gitlab/sidekiq_middleware/duplicate_jobs/strategies/until_executed.rb
index 8c7e15364f8..347f4e61d19 100644
--- a/lib/gitlab/sidekiq_middleware/duplicate_jobs/strategies/until_executed.rb
+++ b/lib/gitlab/sidekiq_middleware/duplicate_jobs/strategies/until_executed.rb
@@ -10,6 +10,8 @@ module Gitlab
         class UntilExecuted < DeduplicatesWhenScheduling
           override :perform
           def perform(job)
+            job_deleted = false
+
             super
 
             yield
@@ -17,7 +19,10 @@ module Gitlab
             should_reschedule = duplicate_job.should_reschedule?
             # Deleting before rescheduling to make sure we don't deduplicate again.
             duplicate_job.delete!
+            job_deleted = true
             duplicate_job.reschedule if should_reschedule
+          ensure
+            duplicate_job.delete! unless job_deleted
           end
         end
       end
diff --git a/lib/gitlab/sidekiq_middleware/worker_context/client.rb b/lib/gitlab/sidekiq_middleware/worker_context/client.rb
index 7d3925e9dec..d9797e9c7c7 100644
--- a/lib/gitlab/sidekiq_middleware/worker_context/client.rb
+++ b/lib/gitlab/sidekiq_middleware/worker_context/client.rb
@@ -19,14 +19,21 @@ module Gitlab
             # This should be inside the context for the arguments so
             # that we don't override the feature category on the worker
             # with the one from the caller.
-            #
+
+            root_caller_id = Gitlab::ApplicationContext.current_context_attribute(:root_caller_id) ||
+              Gitlab::ApplicationContext.current_context_attribute(:caller_id)
+
+            context = {
+              root_caller_id: root_caller_id
+            }
+
             # We do not want to set anything explicitly in the context
             # when the feature category is 'not_owned'.
-            if worker_class.feature_category_not_owned?
-              yield
-            else
-              Gitlab::ApplicationContext.with_context(feature_category: worker_class.get_feature_category.to_s, &block)
+            unless worker_class.feature_category_not_owned?
+              context[:feature_category] = worker_class.get_feature_category.to_s
             end
+
+            Gitlab::ApplicationContext.with_context(**context, &block)
           end
         end
       end
diff --git a/lib/gitlab/sidekiq_middleware/worker_context/server.rb b/lib/gitlab/sidekiq_middleware/worker_context/server.rb
index d026f4918c6..7c2dd80c17a 100644
--- a/lib/gitlab/sidekiq_middleware/worker_context/server.rb
+++ b/lib/gitlab/sidekiq_middleware/worker_context/server.rb
@@ -12,7 +12,9 @@ module Gitlab
           # This is not a worker we know about, perhaps from a gem
           return yield unless worker_class.respond_to?(:get_worker_context)
 
-          Gitlab::ApplicationContext.with_context(feature_category: worker_class.get_feature_category.to_s) do
+          feature_category = worker_class.get_feature_category.to_s
+
+          Gitlab::ApplicationContext.with_context(feature_category: feature_category) do
             # Use the context defined on the class level as the more specific context
             wrap_in_optional_context(worker_class.get_worker_context, &block)
           end