Add latest changes from gitlab-org/gitlab@14-5-stable-eev14.5.0-rc42

1 files changed, 71 insertions, 16 deletions

diff --git a/lib/gitlab/database/load_balancing/setup.rb b/lib/gitlab/database/load_balancing/setup.rb
index 3cce839a960..ef38f42f50b 100644
--- a/lib/gitlab/database/load_balancing/setup.rb
+++ b/lib/gitlab/database/load_balancing/setup.rb
@@ -5,7 +5,7 @@ module Gitlab
     module LoadBalancing
       # Class for setting up load balancing of a specific model.
       class Setup
-        attr_reader :configuration
+        attr_reader :model, :configuration
 
         def initialize(model, start_service_discovery: false)
           @model = model
@@ -14,47 +14,102 @@ module Gitlab
         end
 
         def setup
-          disable_prepared_statements
-          setup_load_balancer
+          configure_connection
+          setup_connection_proxy
           setup_service_discovery
+          setup_feature_flag_to_model_load_balancing
         end
 
-        def disable_prepared_statements
+        def configure_connection
           db_config_object = @model.connection_db_config
-          config =
-            db_config_object.configuration_hash.merge(prepared_statements: false)
+
+          hash = db_config_object.configuration_hash.merge(
+            prepared_statements: false,
+            pool: Gitlab::Database.default_pool_size
+          )
 
           hash_config = ActiveRecord::DatabaseConfigurations::HashConfig.new(
             db_config_object.env_name,
             db_config_object.name,
-            config
+            hash
           )
 
           @model.establish_connection(hash_config)
         end
 
-        def setup_load_balancer
-          lb = LoadBalancer.new(configuration)
-
+        def setup_connection_proxy
           # We just use a simple `class_attribute` here so we don't need to
           # inject any modules and/or expose unnecessary methods.
-          @model.class_attribute(:connection)
-          @model.class_attribute(:sticking)
+          setup_class_attribute(:load_balancer, load_balancer)
+          setup_class_attribute(:connection, ConnectionProxy.new(load_balancer))
+          setup_class_attribute(:sticking, Sticking.new(load_balancer))
+        end
+
+        # TODO: This is temporary code to gradually redirect traffic to use
+        # a dedicated DB replicas, or DB primaries (depending on configuration)
+        # This implements a sticky behavior for the current request if enabled.
+        #
+        # This is needed for Phase 3 and Phase 4 of application rollout
+        # https://gitlab.com/groups/gitlab-org/-/epics/6160#progress
+        #
+        # If `GITLAB_USE_MODEL_LOAD_BALANCING` is set, its value is preferred
+        # Otherwise, a `use_model_load_balancing` FF value is used
+        def setup_feature_flag_to_model_load_balancing
+          return if active_record_base?
 
-          @model.connection = ConnectionProxy.new(lb)
-          @model.sticking = Sticking.new(lb)
+          @model.singleton_class.prepend(ModelLoadBalancingFeatureFlagMixin)
         end
 
         def setup_service_discovery
           return unless configuration.service_discovery_enabled?
 
-          lb = @model.connection.load_balancer
-          sv = ServiceDiscovery.new(lb, **configuration.service_discovery)
+          sv = ServiceDiscovery.new(load_balancer, **configuration.service_discovery)
 
           sv.perform_service_discovery
 
           sv.start if @start_service_discovery
         end
+
+        def load_balancer
+          @load_balancer ||= LoadBalancer.new(configuration)
+        end
+
+        private
+
+        def setup_class_attribute(attribute, value)
+          @model.class_attribute(attribute)
+          @model.public_send("#{attribute}=", value) # rubocop:disable GitlabSecurity/PublicSend
+        end
+
+        def active_record_base?
+          @model == ActiveRecord::Base
+        end
+
+        module ModelLoadBalancingFeatureFlagMixin
+          extend ActiveSupport::Concern
+
+          def use_model_load_balancing?
+            # Cache environment variable and return env variable first if defined
+            use_model_load_balancing_env = Gitlab::Utils.to_boolean(ENV["GITLAB_USE_MODEL_LOAD_BALANCING"])
+
+            unless use_model_load_balancing_env.nil?
+              return use_model_load_balancing_env
+            end
+
+            # Check a feature flag using RequestStore (if active)
+            return false unless Gitlab::SafeRequestStore.active?
+
+            Gitlab::SafeRequestStore.fetch(:use_model_load_balancing) do
+              Feature.enabled?(:use_model_load_balancing, default_enabled: :yaml)
+            end
+          end
+
+          # rubocop:disable Database/MultipleDatabases
+          def connection
+            use_model_load_balancing? ? super : ActiveRecord::Base.connection
+          end
+          # rubocop:enable Database/MultipleDatabases
+        end
       end
     end
   end