Add latest changes from gitlab-org/gitlab@master

3 files changed, 57 insertions, 17 deletions

diff --git a/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb b/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb
index 0e36ebbc3ee..acb4842db31 100644
--- a/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb
+++ b/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb
@@ -4,13 +4,15 @@ module Gitlab
   module Database
     module LoadBalancing
       class SidekiqClientMiddleware
+        include Gitlab::Utils::StrongMemoize
+
         def call(worker_class, job, _queue, _redis_pool)
           # Mailers can't be constantized
           worker_class = worker_class.to_s.safe_constantize
 
           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,16 +29,25 @@ module Gitlab
             worker_class.get_data_consistency_feature_flag_enabled?
         end
 
-        def set_data_consistency_location!(job)
-          if Session.current.use_primary?
-            job['database_write_location'] = load_balancer.primary_write_location
-          else
-            job['database_replica_location'] = load_balancer.host.database_replica_location
-          end
+        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
+          job['wal_locations'] = { main: wal_location } if wal_location
         end
 
-        def location_already_provided?(job)
-          job['database_replica_location'] || job['database_write_location']
+        def wal_location
+          strong_memoize(:wal_location) do
+            if Session.current.use_primary?
+              load_balancer.primary_write_location
+            else
+              load_balancer.host.database_replica_location
+            end
+          end
         end
 
         def load_balancer
diff --git a/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb b/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb
index 0551750568a..2555f693a5d 100644
--- a/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb
+++ b/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_databases_has_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_databases_has_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
diff --git a/lib/gitlab/import_export/project/import_export.yml b/lib/gitlab/import_export/project/import_export.yml
index 21d9470cfb1..51fd84960c5 100644
--- a/lib/gitlab/import_export/project/import_export.yml
+++ b/lib/gitlab/import_export/project/import_export.yml
@@ -128,6 +128,7 @@ excluded_attributes:
     - :name
     - :path
     - :namespace_id
+    - :project_namespace_id
     - :creator_id
     - :pool_repository_id
     - :import_url