Change job hash support multiple wal locations

lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb

@@ -10,7 +10,7 @@ module Gitlab
 
           if load_balancing_enabled?(worker_class)
             job['worker_data_consistency'] = worker_class.get_data_consistency
-            set_data_consistency_location!(job) unless location_already_provided?(job)
+            set_data_consistency_locations!(job) unless job['wal_locations']
           else
             job['worker_data_consistency'] = ::WorkerAttributes::DEFAULT_DATA_CONSISTENCY
           end
@@ -27,18 +27,26 @@ module Gitlab
             worker_class.get_data_consistency_feature_flag_enabled?
         end
 
-        def set_data_consistency_location!(job)
+        def set_data_consistency_locations!(job)
+          # Once we add support for multiple databases to our load balancer, we would use something like this:
+          #   job['wal_locations'] = Gitlab::Database::DATABASES.transform_values do |connection|
+          #      connection.load_balancer.primary_write_location.
+          #   end
+          #
+          # TODO: Replace hardcoded database config name :main when we merge unification strategy
+          # https://gitlab.com/gitlab-org/gitlab/-/issues/336566
+          wal_location = calculate_wal_location
+          job['wal_locations'] = { main: wal_location } if wal_location
+        end
+
+        def calculate_wal_location
           if Session.current.use_primary?
-            job['database_write_location'] = load_balancer.primary_write_location
+            load_balancer.primary_write_location
           else
-            job['database_replica_location'] = load_balancer.host.database_replica_location
+            load_balancer.host.database_replica_location
           end
         end
 
-        def location_already_provided?(job)
-          job['database_replica_location'] || job['database_write_location']
-        end
-
         def load_balancer
           LoadBalancing.proxy.load_balancer
         end

lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb

@@ -29,7 +29,7 @@ module Gitlab
         private
 
         def clear
-          load_balancer.release_host
+          release_hosts
           Session.clear_session
         end
 
@@ -40,10 +40,11 @@ module Gitlab
         def select_load_balancing_strategy(worker_class, job)
           return :primary unless load_balancing_available?(worker_class)
 
-          location = job['database_write_location'] || job['database_replica_location']
-          return :primary_no_wal unless location
+          wal_locations = get_wal_locations(job)
+
+          return :primary_no_wal unless wal_locations
 
-          if replica_caught_up?(location)
+          if all_replica_caught_up?(wal_locations)
             # Happy case: we can read from a replica.
             retried_before?(worker_class, job) ? :replica_retried : :replica
           elsif can_retry?(worker_class, job)
@@ -55,6 +56,19 @@ module Gitlab
           end
         end
 
+        def get_wal_locations(job)
+          job['wal_locations'] || legacy_wal_location(job)
+        end
+
+        # Already scheduled jobs could still contain legacy database write location.
+        # TODO: remove this in the next iteration
+        # https://gitlab.com/gitlab-org/gitlab/-/issues/338213
+        def legacy_wal_location(job)
+          wal_location = job['database_write_location'] || job['database_replica_location']
+
+          { main: wal_location } if wal_location
+        end
+
         def load_balancing_available?(worker_class)
           worker_class.include?(::ApplicationWorker) &&
             worker_class.utilizes_load_balancing_capabilities? &&
@@ -75,12 +89,26 @@ module Gitlab
           job['retry_count'].nil?
         end
 
-        def load_balancer
-          LoadBalancing.proxy.load_balancer
+        def all_replica_caught_up?(wal_locations)
+          wal_locations.all? do |_config_name, location|
+            # Once we add support for multiple databases to our load balancer, we would use something like this:
+            # Gitlab::Database::DATABASES[config_name].load_balancer.select_up_to_date_host(location)
+            load_balancer.select_up_to_date_host(location)
+          end
         end
 
-        def replica_caught_up?(location)
-          load_balancer.select_up_to_date_host(location)
+        def release_hosts
+          # Once we add support for multiple databases to our load balancer, we would use something like this:
+          # connection.load_balancer.primary_write_location
+          #
+          # Gitlab::Database::DATABASES.values.each do |connection|
+          #   connection.load_balancer.release_host
+          # end
+          load_balancer.release_host
+        end
+
+        def load_balancer
+          LoadBalancing.proxy.load_balancer
         end
       end
     end

spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb

@@ -58,8 +58,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
       it 'does not pass database locations', :aggregate_failures do
         run_middleware
 
-        expect(job['database_replica_location']).to be_nil
-        expect(job['database_write_location']).to be_nil
+        expect(job['wal_locations']).to be_nil
       end
 
       include_examples 'job data consistency'
@@ -86,11 +85,13 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
         end
 
         it 'passes database_replica_location' do
+          expected_location = { main: location }
+
           expect(load_balancer).to receive_message_chain(:host, "database_replica_location").and_return(location)
 
           run_middleware
 
-          expect(job['database_replica_location']).to eq(location)
+          expect(job['wal_locations']).to eq(expected_location)
         end
 
         include_examples 'job data consistency'
@@ -102,18 +103,20 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
         end
 
         it 'passes primary write location', :aggregate_failures do
+          expected_location = { main: location }
+
           expect(load_balancer).to receive(:primary_write_location).and_return(location)
 
           run_middleware
 
-          expect(job['database_write_location']).to eq(location)
+          expect(job['wal_locations']).to eq(expected_location)
         end
 
         include_examples 'job data consistency'
       end
     end
 
-    shared_examples_for 'database location was already provided' do |provided_database_location, other_location|
+    shared_examples_for 'database location was already provided' do
       shared_examples_for 'does not set database location again' do |use_primary|
         before do
           allow(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_primary?).and_return(use_primary)
@@ -122,14 +125,13 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
         it 'does not set database locations again' do
           run_middleware
 
-          expect(job[provided_database_location]).to eq(old_location)
-          expect(job[other_location]).to be_nil
+          expect(job['wal_locations']).to eq({ main: old_location })
         end
       end
 
       let(:old_location) { '0/D525E3A8' }
       let(:new_location) { 'AB/12345' }
-      let(:job) { { "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e", provided_database_location => old_location } }
+      let(:job) { { "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e", 'wal_locations' => { main: old_location } } }
 
       before do
         allow(load_balancer).to receive(:primary_write_location).and_return(new_location)
@@ -159,12 +161,8 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
       include_examples 'does not pass database locations'
     end
 
-    context 'database write location was already provided' do
-      include_examples 'database location was already provided', 'database_write_location', 'database_replica_location'
-    end
-
-    context 'database replica location was already provided' do
-      include_examples 'database location was already provided', 'database_replica_location', 'database_write_location'
+    context 'database wal location was already provided' do
+      include_examples 'database location was already provided'
     end
 
     context 'when worker data consistency is :always' do

spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb

@@ -62,9 +62,11 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware do
       include_examples 'load balancing strategy', expected_strategy
     end
 
-    shared_examples_for 'replica is up to date' do |location, expected_strategy|
+    shared_examples_for 'replica is up to date' do |expected_strategy|
+      let(:wal_locations) { { main: '0/D525E3A8' } }
+
       it 'does not stick to the primary', :aggregate_failures do
-        expect(middleware).to receive(:replica_caught_up?).with(location).and_return(true)
+        expect(load_balancer).to receive(:select_up_to_date_host).with(wal_locations[:main]).and_return(true)
 
         run_middleware do
           expect(Gitlab::Database::LoadBalancing::Session.current.use_primary?).not_to be_truthy
@@ -85,30 +87,30 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware do
         include_examples 'stick to the primary', 'primary'
       end
 
-      context 'when database replica location is set' do
-        let(:job) { { 'job_id' => 'a180b47c-3fd6-41b8-81e9-34da61c3400e', 'database_replica_location' => '0/D525E3A8' } }
+      context 'when database wal location is set' do
+        let(:job) { { 'job_id' => 'a180b47c-3fd6-41b8-81e9-34da61c3400e', 'wal_locations' => wal_locations } }
 
         before do
-          allow(middleware).to receive(:replica_caught_up?).and_return(true)
+          allow(load_balancer).to receive(:select_up_to_date_host).with(wal_locations[:main]).and_return(true)
         end
 
-        it_behaves_like 'replica is up to date', '0/D525E3A8', 'replica'
+        it_behaves_like 'replica is up to date', 'replica'
       end
 
-      context 'when database primary location is set' do
+      context 'when legacy wal location is set' do
         let(:job) { { 'job_id' => 'a180b47c-3fd6-41b8-81e9-34da61c3400e', 'database_write_location' => '0/D525E3A8' } }
 
         before do
-          allow(middleware).to receive(:replica_caught_up?).and_return(true)
+          allow(load_balancer).to receive(:select_up_to_date_host).with('0/D525E3A8').and_return(true)
         end
 
-        it_behaves_like 'replica is up to date', '0/D525E3A8', 'replica'
+        it_behaves_like 'replica is up to date', 'replica'
       end
 
       context 'when database location is not set' do
         let(:job) { { 'job_id' => 'a180b47c-3fd6-41b8-81e9-34da61c3400e' } }
 
-        it_behaves_like 'stick to the primary', 'primary_no_wal'
+        include_examples 'stick to the primary', 'primary_no_wal'
       end
     end
 
@@ -167,7 +169,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware do
               replication_lag!(false)
             end
 
-            it_behaves_like 'replica is up to date', '0/D525E3A8', 'replica_retried'
+            include_examples 'replica is up to date', 'replica_retried'
           end
         end
       end
@@ -178,7 +180,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware do
 
       context 'when replica is not up to date' do
         before do
-          allow(middleware).to receive(:replica_caught_up?).and_return(false)
+          allow(load_balancer).to receive(:select_up_to_date_host).and_return(false)
         end
 
         include_examples 'stick to the primary', 'primary'