diff options
author | GitLab Bot <gitlab-bot@gitlab.com> | 2021-11-18 16:16:36 +0300 |
---|---|---|
committer | GitLab Bot <gitlab-bot@gitlab.com> | 2021-11-18 16:16:36 +0300 |
commit | 311b0269b4eb9839fa63f80c8d7a58f32b8138a0 (patch) | |
tree | 07e7870bca8aed6d61fdcc810731c50d2c40af47 /lib/gitlab/database/load_balancing/setup.rb | |
parent | 27909cef6c4170ed9205afa7426b8d3de47cbb0c (diff) |
Add latest changes from gitlab-org/gitlab@14-5-stable-eev14.5.0-rc42
Diffstat (limited to 'lib/gitlab/database/load_balancing/setup.rb')
-rw-r--r-- | lib/gitlab/database/load_balancing/setup.rb | 87 |
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 |