116 lines
3.9 KiB
Ruby
116 lines
3.9 KiB
Ruby
# frozen_string_literal: true
|
|
|
|
module Gitlab
|
|
module Database
|
|
module LoadBalancing
|
|
# Class for setting up load balancing of a specific model.
|
|
class Setup
|
|
attr_reader :model, :configuration
|
|
|
|
def initialize(model, start_service_discovery: false)
|
|
@model = model
|
|
@configuration = Configuration.for_model(model)
|
|
@start_service_discovery = start_service_discovery
|
|
end
|
|
|
|
def setup
|
|
configure_connection
|
|
setup_connection_proxy
|
|
setup_service_discovery
|
|
setup_feature_flag_to_model_load_balancing
|
|
end
|
|
|
|
def configure_connection
|
|
db_config_object = @model.connection_db_config
|
|
|
|
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,
|
|
hash
|
|
)
|
|
|
|
@model.establish_connection(hash_config)
|
|
end
|
|
|
|
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.
|
|
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.singleton_class.prepend(ModelLoadBalancingFeatureFlagMixin)
|
|
end
|
|
|
|
def setup_service_discovery
|
|
return unless configuration.service_discovery_enabled?
|
|
|
|
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
|
|
end
|