1 files changed, 40 insertions, 23 deletions

diff --git a/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb b/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb
index 9bd0adf8dbd..0551750568a 100644
--- a/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb
+++ b/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb
@@ -7,8 +7,18 @@ module Gitlab
         JobReplicaNotUpToDate = Class.new(StandardError)
 
         def call(worker, job, _queue)
-          if requires_primary?(worker.class, job)
+          worker_class = worker.class
+          strategy = select_load_balancing_strategy(worker_class, job)
+
+          job['load_balancing_strategy'] = strategy.to_s
+
+          if use_primary?(strategy)
             Session.current.use_primary!
+          elsif strategy == :retry
+            raise JobReplicaNotUpToDate, "Sidekiq job #{worker_class} JID-#{job['jid']} couldn't use the replica."\
+              "  Replica was not up to date."
+          else
+            # this means we selected an up-to-date replica, but there is nothing to do in this case.
           end
 
           yield
@@ -23,31 +33,42 @@ module Gitlab
           Session.clear_session
         end
 
-        def requires_primary?(worker_class, job)
-          return true unless worker_class.include?(::ApplicationWorker)
-          return true unless worker_class.utilizes_load_balancing_capabilities?
-          return true unless worker_class.get_data_consistency_feature_flag_enabled?
-
-          location = job['database_write_location'] || job['database_replica_location']
+        def use_primary?(strategy)
+          strategy.start_with?('primary')
+        end
 
-          return true unless location
+        def select_load_balancing_strategy(worker_class, job)
+          return :primary unless load_balancing_available?(worker_class)
 
-          job_data_consistency = worker_class.get_data_consistency
-          job[:data_consistency] = job_data_consistency.to_s
+          location = job['database_write_location'] || job['database_replica_location']
+          return :primary_no_wal unless location
 
           if replica_caught_up?(location)
-            job[:database_chosen] = 'replica'
-            false
-          elsif job_data_consistency == :delayed && not_yet_retried?(job)
-            job[:database_chosen] = 'retry'
-            raise JobReplicaNotUpToDate, "Sidekiq job #{worker_class} JID-#{job['jid']} couldn't use the replica."\
-               "  Replica was not up to date."
+            # Happy case: we can read from a replica.
+            retried_before?(worker_class, job) ? :replica_retried : :replica
+          elsif can_retry?(worker_class, job)
+            # Optimistic case: The worker allows retries and we have retries left.
+            :retry
           else
-            job[:database_chosen] = 'primary'
-            true
+            # Sad case: we need to fall back to the primary.
+            :primary
           end
         end
 
+        def load_balancing_available?(worker_class)
+          worker_class.include?(::ApplicationWorker) &&
+            worker_class.utilizes_load_balancing_capabilities? &&
+            worker_class.get_data_consistency_feature_flag_enabled?
+        end
+
+        def can_retry?(worker_class, job)
+          worker_class.get_data_consistency == :delayed && not_yet_retried?(job)
+        end
+
+        def retried_before?(worker_class, job)
+          worker_class.get_data_consistency == :delayed && !not_yet_retried?(job)
+        end
+
         def not_yet_retried?(job)
           # if `retry_count` is `nil` it indicates that this job was never retried
           # the `0` indicates that this is a first retry
@@ -59,11 +80,7 @@ module Gitlab
         end
 
         def replica_caught_up?(location)
-          if Feature.enabled?(:sidekiq_load_balancing_rotate_up_to_date_replica)
-            load_balancer.select_up_to_date_host(location)
-          else
-            load_balancer.host.caught_up?(location)
-          end
+          load_balancer.select_up_to_date_host(location)
         end
       end
     end