diff options
Diffstat (limited to 'lib/gitlab/database/load_balancing/setup.rb')
-rw-r--r-- | lib/gitlab/database/load_balancing/setup.rb | 77 |
1 files changed, 66 insertions, 11 deletions
diff --git a/lib/gitlab/database/load_balancing/setup.rb b/lib/gitlab/database/load_balancing/setup.rb index a44ace9da26..bcb40e45dac 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 @@ -15,8 +15,9 @@ module Gitlab def setup configure_connection - setup_load_balancer + setup_connection_proxy setup_service_discovery + setup_feature_flag_to_model_load_balancing end def configure_connection @@ -36,28 +37,82 @@ module Gitlab @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(: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 + + def sticking + use_model_load_balancing? ? super : ActiveRecord::Base.sticking + end + # rubocop:enable Database/MultipleDatabases + end end end end |