Commit 829d5921 authored by Dylan Griffith's avatar Dylan Griffith

Merge branch 'feature-flag-rollout' into 'master'

Allow percentage rollout of load balancer of decomposed database

See merge request gitlab-org/gitlab!73631
parents 311ca606 aef44f07
......@@ -22,6 +22,7 @@
variables:
DECOMPOSED_DB: "true"
GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: "main"
GITLAB_USE_MODEL_LOAD_BALANCING: "true"
.rspec-base:
extends: .rails-job-base
......
---
name: use_model_load_balancing
introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/73631
rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/344797
milestone: '14.5'
type: development
group: group::sharding
default_enabled: false
......@@ -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
......
......@@ -8,8 +8,9 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do
setup = described_class.new(ActiveRecord::Base)
expect(setup).to receive(:configure_connection)
expect(setup).to receive(:setup_load_balancer)
expect(setup).to receive(:setup_connection_proxy)
expect(setup).to receive(:setup_service_discovery)
expect(setup).to receive(:setup_feature_flag_to_model_load_balancing)
setup.setup
end
......@@ -44,7 +45,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do
end
end
describe '#setup_load_balancer' do
describe '#setup_connection_proxy' do
it 'sets up the load balancer' do
model = Class.new(ActiveRecord::Base)
setup = described_class.new(model)
......@@ -58,7 +59,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do
.with(setup.configuration)
.and_return(lb)
setup.setup_load_balancer
setup.setup_connection_proxy
expect(model.connection.load_balancer).to eq(lb)
expect(model.sticking)
......@@ -81,7 +82,6 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do
model = ActiveRecord::Base
setup = described_class.new(model)
sv = instance_spy(Gitlab::Database::LoadBalancing::ServiceDiscovery)
lb = model.connection.load_balancer
allow(setup.configuration)
.to receive(:service_discovery_enabled?)
......@@ -89,7 +89,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do
allow(Gitlab::Database::LoadBalancing::ServiceDiscovery)
.to receive(:new)
.with(lb, setup.configuration.service_discovery)
.with(setup.load_balancer, setup.configuration.service_discovery)
.and_return(sv)
expect(sv).to receive(:perform_service_discovery)
......@@ -102,7 +102,6 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do
model = ActiveRecord::Base
setup = described_class.new(model, start_service_discovery: true)
sv = instance_spy(Gitlab::Database::LoadBalancing::ServiceDiscovery)
lb = model.connection.load_balancer
allow(setup.configuration)
.to receive(:service_discovery_enabled?)
......@@ -110,7 +109,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do
allow(Gitlab::Database::LoadBalancing::ServiceDiscovery)
.to receive(:new)
.with(lb, setup.configuration.service_discovery)
.with(setup.load_balancer, setup.configuration.service_discovery)
.and_return(sv)
expect(sv).to receive(:perform_service_discovery)
......@@ -120,4 +119,172 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do
end
end
end
describe '#setup_feature_flag_to_model_load_balancing', :reestablished_active_record_base do
using RSpec::Parameterized::TableSyntax
where do
{
"with model LB enabled it picks a dedicated CI connection" => {
env_GITLAB_USE_MODEL_LOAD_BALANCING: 'true',
env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: nil,
request_store_active: false,
ff_use_model_load_balancing: nil,
expectations: {
main: { read: 'main_replica', write: 'main' },
ci: { read: 'ci_replica', write: 'ci' }
}
},
"with model LB enabled and re-use of primary connection it uses CI connection for reads" => {
env_GITLAB_USE_MODEL_LOAD_BALANCING: 'true',
env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: 'main',
request_store_active: false,
ff_use_model_load_balancing: nil,
expectations: {
main: { read: 'main_replica', write: 'main' },
ci: { read: 'ci_replica', write: 'main' }
}
},
"with model LB disabled it fallbacks to use main" => {
env_GITLAB_USE_MODEL_LOAD_BALANCING: 'false',
env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: nil,
request_store_active: false,
ff_use_model_load_balancing: nil,
expectations: {
main: { read: 'main_replica', write: 'main' },
ci: { read: 'main_replica', write: 'main' }
}
},
"with model LB disabled, but re-use configured it fallbacks to use main" => {
env_GITLAB_USE_MODEL_LOAD_BALANCING: 'false',
env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: 'main',
request_store_active: false,
ff_use_model_load_balancing: nil,
expectations: {
main: { read: 'main_replica', write: 'main' },
ci: { read: 'main_replica', write: 'main' }
}
},
"with FF disabled without RequestStore it uses main" => {
env_GITLAB_USE_MODEL_LOAD_BALANCING: nil,
env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: nil,
request_store_active: false,
ff_use_model_load_balancing: false,
expectations: {
main: { read: 'main_replica', write: 'main' },
ci: { read: 'main_replica', write: 'main' }
}
},
"with FF enabled without RequestStore sticking of FF does not work, so it fallbacks to use main" => {
env_GITLAB_USE_MODEL_LOAD_BALANCING: nil,
env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: nil,
request_store_active: false,
ff_use_model_load_balancing: true,
expectations: {
main: { read: 'main_replica', write: 'main' },
ci: { read: 'main_replica', write: 'main' }
}
},
"with FF disabled with RequestStore it uses main" => {
env_GITLAB_USE_MODEL_LOAD_BALANCING: nil,
env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: nil,
request_store_active: true,
ff_use_model_load_balancing: false,
expectations: {
main: { read: 'main_replica', write: 'main' },
ci: { read: 'main_replica', write: 'main' }
}
},
"with FF enabled with RequestStore it sticks FF and uses CI connection" => {
env_GITLAB_USE_MODEL_LOAD_BALANCING: nil,
env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: nil,
request_store_active: true,
ff_use_model_load_balancing: true,
expectations: {
main: { read: 'main_replica', write: 'main' },
ci: { read: 'ci_replica', write: 'ci' }
}
},
"with re-use and FF enabled with RequestStore it sticks FF and uses CI connection for reads" => {
env_GITLAB_USE_MODEL_LOAD_BALANCING: nil,
env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: 'main',
request_store_active: true,
ff_use_model_load_balancing: true,
expectations: {
main: { read: 'main_replica', write: 'main' },
ci: { read: 'ci_replica', write: 'main' }
}
}
}
end
with_them do
let(:ci_class) do
Class.new(ActiveRecord::Base) do
def self.name
'Ci::ApplicationRecordTemporary'
end
establish_connection ActiveRecord::DatabaseConfigurations::HashConfig.new(
Rails.env,
'ci',
ActiveRecord::Base.connection_db_config.configuration_hash
)
end
end
let(:models) do
{
main: ActiveRecord::Base,
ci: ci_class
}
end
around do |example|
if request_store_active
Gitlab::WithRequestStore.with_request_store do
RequestStore.clear!
example.run
end
else
example.run
end
end
before do
# Rewrite `class_attribute` to use rspec mocking and prevent modifying the objects
allow_next_instance_of(described_class) do |setup|
allow(setup).to receive(:configure_connection)
allow(setup).to receive(:setup_class_attribute) do |attribute, value|
allow(setup.model).to receive(attribute) { value }
end
end
stub_env('GITLAB_USE_MODEL_LOAD_BALANCING', env_GITLAB_USE_MODEL_LOAD_BALANCING)
stub_env('GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci', env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci)
stub_feature_flags(use_model_load_balancing: ff_use_model_load_balancing)
end
it 'results match expectations' do
result = models.transform_values do |model|
# Make load balancer to force init with a dedicated replicas connections
described_class.new(model).tap do |subject|
subject.configuration.hosts = [subject.configuration.replica_db_config.host]
subject.setup
end
load_balancer = model.connection.load_balancer
{
read: load_balancer.read { |connection| connection.pool.db_config.name },
write: load_balancer.read_write { |connection| connection.pool.db_config.name }
}
end
expect(result).to eq(expectations)
end
end
end
end
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment