Address reviewer feedback

- extract method utilizes_load_balancing_capabilities?
- Fix application_worker specs

app/workers/concerns/application_worker.rb

@@ -55,7 +55,7 @@ module ApplicationWorker
     def perform_async(*args)
       # Worker execution for workers with data_consistency set to :delayed or :sticky
       # will be delayed to give replication enough time to complete
-      if delayed_execution? && data_consistency_delayed_execution_feature_flag_enabled?
+      if utilizes_load_balancing_capabilities? && data_consistency_delayed_execution_feature_flag_enabled?
         perform_in(delay_interval, *args)
       else
         super
@@ -132,11 +132,5 @@ module ApplicationWorker
     def delay_interval
       DEFAULT_DELAY_INTERVAL.seconds
     end
-
-    private
-
-    def delayed_execution?
-      get_data_consistency != :always
-    end
   end
 end

app/workers/concerns/worker_attributes.rb

@@ -99,11 +99,16 @@ module WorkerAttributes
       # Since the deduplication should always take into account the latest binary replication pointer into account,
       # not the first one, the deduplication will not work with sticky or delayed.
       # Follow up issue to improve this: https://gitlab.com/gitlab-org/gitlab/-/issues/325291
-      if idempotent? && get_data_consistency != :always
+      if idempotent? && utilizes_load_balancing_capabilities?
         raise ArgumentError, "Class can't be marked as idempotent if data_consistency is not set to :always"
       end
     end
 
+    # If data_consistency is not set to :always, worker will try to utilize load balancing capabilities and use the replica
+    def utilizes_load_balancing_capabilities?
+      get_data_consistency != :always
+    end
+
     def get_data_consistency
       class_attributes[:data_consistency] || :always
     end

ee/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb

@@ -24,7 +24,7 @@ module Gitlab
 
           job['worker_data_consistency'] = worker_class.get_data_consistency
 
-          return if worker_class.get_data_consistency == :always
+          return unless worker_class.utilizes_load_balancing_capabilities?
 
           if Session.current.performed_write?
             job['database_write_location'] = load_balancer.primary_write_location

ee/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb

@@ -25,7 +25,7 @@ module Gitlab
 
         def requires_primary?(worker_class, job)
           return true unless worker_class.include?(::ApplicationWorker)
-          return true if worker_class.get_data_consistency == :always
+          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']

spec/workers/concerns/application_worker_spec.rb

@@ -177,9 +177,9 @@ RSpec.describe ApplicationWorker do
   end
 
   describe '.perform_async' do
-    context 'when workers data consistency is not :always' do
+    shared_examples_for 'worker utilizes load balancing capabilities' do |data_consistency|
       before do
-        worker.data_consistency(:sticky)
+        worker.data_consistency(data_consistency)
       end
 
       context 'when data_consistency_delayed_execution feature flag is disabled' do
@@ -200,12 +200,6 @@ RSpec.describe ApplicationWorker do
         end
       end
 
-      it 'delayed_execution? return true' do
-        expect(worker).to receive(:delayed_execution?).and_return(true)
-
-        worker.perform_async
-      end
-
       it 'call perform_in' do
         expect(worker).to receive(:perform_in).with(described_class::DEFAULT_DELAY_INTERVAL.seconds, 123)
 
@@ -213,6 +207,14 @@ RSpec.describe ApplicationWorker do
       end
     end
 
+    context 'when workers data consistency is :sticky' do
+      it_behaves_like 'worker utilizes load balancing capabilities', :sticky
+    end
+
+    context 'when workers data consistency is :delayed' do
+      it_behaves_like 'worker utilizes load balancing capabilities', :delayed
+    end
+
     context 'when workers data consistency is :always' do
       before do
         worker.data_consistency(:always)
@@ -223,12 +225,6 @@ RSpec.describe ApplicationWorker do
 
         worker.perform_async
       end
-
-      it 'delayed_execution? return false' do
-        expect(worker).to receive(:delayed_execution?).and_return(false)
-
-        worker.perform_async
-      end
     end
   end