Commit f38900c9 authored by Yorick Peterse's avatar Yorick Peterse Committed by Yorick Peterse

Always enable the database load balancer

By always enabling the database load balancer we reduce the amount of
different code paths, make it easier for developers to interact with the
load balancing code (e.g. by adding metrics), and make it easier to add
support for multiple databases (e.g. by not having to check if the load
balancer is enabled or not).

As part of this commit we fixed various parts of the code that didn't
expect the load balancer to always be enabled (or at least when running
tests).

Changelog: added
Co-Authored-By: default avatarDylan Griffith <dyl.griffith@gmail.com>
parent 5fab6ac5
......@@ -30,8 +30,6 @@ module AuthorizedProjectUpdate
# does not allow us to deduplicate these jobs.
# https://gitlab.com/gitlab-org/gitlab/-/issues/325291
def use_replica_if_available(&block)
return yield unless ::Gitlab::Database::LoadBalancing.enable?
::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries(&block)
end
......
......@@ -45,8 +45,6 @@ class ContainerExpirationPolicyWorker # rubocop:disable Scalability/IdempotentWo
# not perfomed with a delay
# https://gitlab.com/gitlab-org/gitlab/-/merge_requests/63635#note_603771207
def use_replica_if_available(&blk)
return yield unless ::Gitlab::Database::LoadBalancing.enable?
::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries(&blk)
end
......
......@@ -19,4 +19,4 @@ ActionCable::SubscriptionAdapter::Redis.redis_connector = lambda do |config|
end
Gitlab::ActionCable::RequestStoreCallbacks.install
Gitlab::Database::LoadBalancing::ActionCableCallbacks.install if Gitlab::Database::LoadBalancing.enable?
Gitlab::Database::LoadBalancing::ActionCableCallbacks.install
......@@ -2,27 +2,51 @@
ActiveRecord::Base.singleton_class.attr_accessor :load_balancing_proxy
if Gitlab::Database::LoadBalancing.enable?
Gitlab::Database.main.disable_prepared_statements
Gitlab::Database.main.disable_prepared_statements
Gitlab::Application.configure do |config|
config.middleware.use(Gitlab::Database::LoadBalancing::RackMiddleware)
end
Gitlab::Application.configure do |config|
config.middleware.use(Gitlab::Database::LoadBalancing::RackMiddleware)
end
# This hijacks the "connection" method to ensure both
# `ActiveRecord::Base.connection` and all models use the same load
# balancing proxy.
ActiveRecord::Base.singleton_class.prepend(Gitlab::Database::LoadBalancing::ActiveRecordProxy)
# The load balancer needs to be configured immediately, and re-configured after
# forking. This ensures queries that run before forking use the load balancer,
# and queries running after a fork don't run into any errors when using dead
# database connections.
#
# See https://gitlab.com/gitlab-org/gitlab/-/merge_requests/63485 for more
# information.
setup = proc do
lb = Gitlab::Database::LoadBalancing::LoadBalancer.new(
Gitlab::Database::LoadBalancing.configuration,
primary_only: !Gitlab::Database::LoadBalancing.enable_replicas?
)
# This hijacks the "connection" method to ensure both
# `ActiveRecord::Base.connection` and all models use the same load
# balancing proxy.
ActiveRecord::Base.singleton_class.prepend(Gitlab::Database::LoadBalancing::ActiveRecordProxy)
ActiveRecord::Base.load_balancing_proxy =
Gitlab::Database::LoadBalancing::ConnectionProxy.new(lb)
Gitlab::Database::LoadBalancing.configure_proxy
# Populate service discovery immediately if it is configured
Gitlab::Database::LoadBalancing.perform_service_discovery
end
setup.call
# This needs to be executed after fork of clustered processes
Gitlab::Cluster::LifecycleEvents.on_worker_start do
# For Host-based LB, we need to re-connect as Rails discards connections on fork
Gitlab::Database::LoadBalancing.configure_proxy
# Database queries may be run before we fork, so we must set up the load
# balancer as early as possible. When we do fork, we need to make sure all the
# hosts are disconnected.
Gitlab::Cluster::LifecycleEvents.on_before_fork do
# When forking, we don't want to wait until the connections aren't in use any
# more, as this could delay the boot cycle.
Gitlab::Database::LoadBalancing.proxy.load_balancer.disconnect!(timeout: 0)
end
# Service discovery must be started after configuring the proxy, as service
# discovery depends on this.
Gitlab::Database::LoadBalancing.start_service_discovery
end
# Service discovery only needs to run in the worker processes, as the main one
# won't be running many (if any) database queries.
Gitlab::Cluster::LifecycleEvents.on_worker_start do
setup.call
Gitlab::Database::LoadBalancing.start_service_discovery
end
......@@ -151,12 +151,8 @@ if Gitlab::Metrics.enabled? && !Rails.env.test? && !(Rails.env.development? && d
Gitlab::Application.configure do |config|
# We want to track certain metrics during the Load Balancing host resolving process.
# Because of that, we need to have metrics code available earlier for Load Balancing.
if Gitlab::Database::LoadBalancing.enable?
config.middleware.insert_before Gitlab::Database::LoadBalancing::RackMiddleware,
Gitlab::Metrics::RackMiddleware
else
config.middleware.use(Gitlab::Metrics::RackMiddleware)
end
config.middleware.insert_before Gitlab::Database::LoadBalancing::RackMiddleware,
Gitlab::Metrics::RackMiddleware
config.middleware.use(Gitlab::Middleware::RailsQueueDuration)
config.middleware.use(Gitlab::Metrics::ElasticsearchRackMiddleware)
......
......@@ -166,7 +166,7 @@ RSpec.describe Gitlab::Geo, :geo, :request_store do
describe '.connected?' do
context 'when there is a database issue' do
it 'returns false when it cannot open an active database connection' do
allow(GeoNode.connection).to receive(:active?).and_return(false)
allow(GeoNode.retrieve_connection).to receive(:active?).and_return(false)
expect(described_class.connected?).to be_falsey
end
......
......@@ -13,23 +13,21 @@ module Gitlab
end
def match?
if ::Gitlab::Database::LoadBalancing.enable?
# When a user merges a merge request, the following sequence happens:
#
# 1. Sidekiq: MergeService runs and updates the merge request in a locked state.
# 2. Gitaly: The UserMergeBranch RPC runs.
# 3. Gitaly (gitaly-ruby): This RPC calls the pre-receive hook.
# 4. Rails: This hook makes an API request to /api/v4/internal/allowed.
# 5. Rails: This API check does a SQL query for locked merge
# requests with a matching SHA.
#
# Since steps 1 and 5 will happen on different database
# sessions, replication lag could erroneously cause step 5 to
# report no matching merge requests. To avoid this, we check
# the write location to ensure the replica can make this query.
track_session_metrics do
::Gitlab::Database::LoadBalancing::Sticking.select_valid_host(:project, @project.id)
end
# When a user merges a merge request, the following sequence happens:
#
# 1. Sidekiq: MergeService runs and updates the merge request in a locked state.
# 2. Gitaly: The UserMergeBranch RPC runs.
# 3. Gitaly (gitaly-ruby): This RPC calls the pre-receive hook.
# 4. Rails: This hook makes an API request to /api/v4/internal/allowed.
# 5. Rails: This API check does a SQL query for locked merge
# requests with a matching SHA.
#
# Since steps 1 and 5 will happen on different database
# sessions, replication lag could erroneously cause step 5 to
# report no matching merge requests. To avoid this, we check
# the write location to ensure the replica can make this query.
track_session_metrics do
::Gitlab::Database::LoadBalancing::Sticking.select_valid_host(:project, @project.id)
end
# rubocop: disable CodeReuse/ActiveRecord
......
......@@ -19,19 +19,8 @@ module Gitlab
[].freeze
end
ProxyNotConfiguredError = Class.new(StandardError)
# The connection proxy to use for load balancing (if enabled).
def self.proxy
unless load_balancing_proxy = ActiveRecord::Base.load_balancing_proxy
Gitlab::ErrorTracking.track_exception(
ProxyNotConfiguredError.new(
"Attempting to access the database load balancing proxy, but it wasn't configured.\n" \
"Did you forget to call '#{self.name}.configure_proxy'?"
))
end
load_balancing_proxy
ActiveRecord::Base.load_balancing_proxy
end
# Returns a Hash containing the load balancing configuration.
......@@ -39,8 +28,11 @@ module Gitlab
@configuration ||= Configuration.for_model(ActiveRecord::Base)
end
# Returns true if load balancing is to be enabled.
def self.enable?
# Returns `true` if the use of load balancing replicas should be enabled.
#
# This is disabled for Rake tasks to ensure e.g. database migrations
# always produce consistent results.
def self.enable_replicas?
return false if Gitlab::Runtime.rake?
configured?
......@@ -59,17 +51,12 @@ module Gitlab
.start
end
# Configures proxying of requests.
def self.configure_proxy
lb = LoadBalancer.new(configuration, primary_only: !enable?)
ActiveRecord::Base.load_balancing_proxy = ConnectionProxy.new(lb)
def self.perform_service_discovery
return unless configuration.service_discovery_enabled?
# Populate service discovery immediately if it is configured
if configuration.service_discovery_enabled?
ServiceDiscovery
.new(lb, **configuration.service_discovery)
.perform_service_discovery
end
ServiceDiscovery
.new(proxy.load_balancer, **configuration.service_discovery)
.perform_service_discovery
end
DB_ROLES = [
......
......@@ -18,8 +18,6 @@ module Gitlab
# namespace - The namespace to use for sticking.
# id - The identifier to use for sticking.
def self.stick_or_unstick(env, namespace, id)
return unless ::Gitlab::Database::LoadBalancing.enable?
::Gitlab::Database::LoadBalancing::Sticking.unstick_or_continue_sticking(namespace, id)
env[STICK_OBJECT] ||= Set.new
......
......@@ -22,8 +22,6 @@ module Gitlab
# Sticks to the primary if a write was performed.
def self.stick_if_necessary(namespace, id)
return unless LoadBalancing.enable?
stick(namespace, id) if Session.current.performed_write?
end
......@@ -75,15 +73,11 @@ module Gitlab
# Starts sticking to the primary for the given namespace and id, using
# the latest WAL pointer from the primary.
def self.stick(namespace, id)
return unless LoadBalancing.enable?
mark_primary_write_location(namespace, id)
Session.current.use_primary!
end
def self.bulk_stick(namespace, ids)
return unless LoadBalancing.enable?
with_primary_write_location do |location|
ids.each do |id|
set_write_location_for(namespace, id, location)
......@@ -94,17 +88,7 @@ module Gitlab
end
def self.with_primary_write_location
return unless LoadBalancing.configured?
# Load balancing could be enabled for the Web application server,
# but it's not activated for Sidekiq. We should update Redis with
# the write location just in case load balancing is being used.
location =
if LoadBalancing.enable?
load_balancer.primary_write_location
else
Gitlab::Database.main.get_write_location(ActiveRecord::Base.connection)
end
location = load_balancer.primary_write_location
return if location.blank?
......
......@@ -171,7 +171,6 @@ module Gitlab
def read_from_replica_if_available(&block)
return yield unless ::Feature.enabled?(:load_balancing_for_export_workers, type: :development, default_enabled: :yaml)
return yield unless ::Gitlab::Database::LoadBalancing.enable?
::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries(&block)
end
......
......@@ -47,13 +47,11 @@ module Gitlab
buckets SQL_DURATION_BUCKET
end
if ::Gitlab::Database::LoadBalancing.enable?
db_role = ::Gitlab::Database::LoadBalancing.db_role_for_connection(payload[:connection])
return if db_role.blank?
db_role = ::Gitlab::Database::LoadBalancing.db_role_for_connection(payload[:connection])
return if db_role.blank?
increment_db_role_counters(db_role, payload)
observe_db_role_duration(db_role, event)
end
increment_db_role_counters(db_role, payload)
observe_db_role_duration(db_role, event)
end
def self.db_counter_payload
......@@ -64,7 +62,7 @@ module Gitlab
payload[key] = Gitlab::SafeRequestStore[key].to_i
end
if ::Gitlab::SafeRequestStore.active? && ::Gitlab::Database::LoadBalancing.enable?
if ::Gitlab::SafeRequestStore.active?
load_balancing_metric_counter_keys.each do |counter|
payload[counter] = ::Gitlab::SafeRequestStore[counter].to_i
end
......
......@@ -10,7 +10,7 @@ module Gitlab
LOG_COUNTERS = { true => :caught_up_replica_pick_ok, false => :caught_up_replica_pick_fail }.freeze
def caught_up_replica_pick(event)
return unless Gitlab::SafeRequestStore.active? && ::Gitlab::Database::LoadBalancing.enable?
return unless Gitlab::SafeRequestStore.active?
result = event.payload[:result]
counter_name = counter(result)
......@@ -20,13 +20,13 @@ module Gitlab
# we want to update Prometheus counter after the controller/action are set
def web_transaction_completed(_event)
return unless Gitlab::SafeRequestStore.active? && ::Gitlab::Database::LoadBalancing.enable?
return unless Gitlab::SafeRequestStore.active?
LOG_COUNTERS.keys.each { |result| increment_prometheus_for_result_label(result) }
end
def self.load_balancing_payload
return {} unless Gitlab::SafeRequestStore.active? && ::Gitlab::Database::LoadBalancing.enable?
return {} unless Gitlab::SafeRequestStore.active?
{}.tap do |payload|
LOG_COUNTERS.values.each do |counter|
......
......@@ -39,7 +39,7 @@ module Gitlab
# DuplicateJobs::Server should be placed at the bottom, but before the SidekiqServerMiddleware,
# so we can compare the latest WAL location against replica
chain.add ::Gitlab::SidekiqMiddleware::DuplicateJobs::Server
chain.add ::Gitlab::Database::LoadBalancing::SidekiqServerMiddleware if load_balancing_enabled?
chain.add ::Gitlab::Database::LoadBalancing::SidekiqServerMiddleware
end
end
......@@ -52,7 +52,7 @@ module Gitlab
chain.add ::Labkit::Middleware::Sidekiq::Client
# Sidekiq Client Middleware should be placed before DuplicateJobs::Client middleware,
# so we can store WAL location before we deduplicate the job.
chain.add ::Gitlab::Database::LoadBalancing::SidekiqClientMiddleware if load_balancing_enabled?
chain.add ::Gitlab::Database::LoadBalancing::SidekiqClientMiddleware
chain.add ::Gitlab::SidekiqMiddleware::DuplicateJobs::Client
chain.add ::Gitlab::SidekiqStatus::ClientMiddleware
chain.add ::Gitlab::SidekiqMiddleware::AdminMode::Client
......@@ -61,10 +61,5 @@ module Gitlab
chain.add ::Gitlab::SidekiqMiddleware::ClientMetrics
end
end
def self.load_balancing_enabled?
::Gitlab::Database::LoadBalancing.enable?
end
private_class_method :load_balancing_enabled?
end
end
......@@ -53,10 +53,7 @@ module Gitlab
def initialize
@metrics = self.class.metrics
if ::Gitlab::Database::LoadBalancing.enable?
@metrics[:sidekiq_load_balancing_count] = ::Gitlab::Metrics.counter(:sidekiq_load_balancing_count, 'Sidekiq jobs with load balancing')
end
@metrics[:sidekiq_load_balancing_count] = ::Gitlab::Metrics.counter(:sidekiq_load_balancing_count, 'Sidekiq jobs with load balancing')
end
def call(worker, job, queue)
......@@ -128,8 +125,6 @@ module Gitlab
private
def with_load_balancing_settings(job)
return unless ::Gitlab::Database::LoadBalancing.enable?
keys = %w[load_balancing_strategy worker_data_consistency]
return unless keys.all? { |k| job.key?(k) }
......
......@@ -44,10 +44,8 @@ module Peek
count[item[:transaction]] += 1
end
if ::Gitlab::Database::LoadBalancing.enable?
count[item[:db_role]] ||= 0
count[item[:db_role]] += 1
end
count[item[:db_role]] ||= 0
count[item[:db_role]] += 1
end
def setup_subscribers
......@@ -72,8 +70,6 @@ module Peek
end
def db_role(data)
return unless ::Gitlab::Database::LoadBalancing.enable?
role = ::Gitlab::Database::LoadBalancing.db_role_for_connection(data[:connection]) ||
::Gitlab::Database::LoadBalancing::ROLE_UNKNOWN
......
......@@ -230,39 +230,21 @@ RSpec.describe 'lograge', type: :request do
end
end
context 'when load balancing is enabled' do
before do
allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
end
context 'with db payload' do
context 'when RequestStore is enabled', :request_store do
it 'includes db counters for load balancing' do
subscriber.process_action(event)
expect(log_data).to include(*db_load_balancing_logging_keys)
end
end
context 'when RequestStore is disabled' do
it 'does not include db counters for load balancing' do
subscriber.process_action(event)
context 'with db payload' do
context 'when RequestStore is enabled', :request_store do
it 'includes db counters for load balancing' do
subscriber.process_action(event)
expect(log_data).not_to include(*db_load_balancing_logging_keys)
end
expect(log_data).to include(*db_load_balancing_logging_keys)
end
end
end
context 'when load balancing is disabled' do
before do
allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(false)
end
context 'when RequestStore is disabled' do
it 'does not include db counters for load balancing' do
subscriber.process_action(event)
it 'does not include db counters for load balancing' do
subscriber.process_action(event)
expect(log_data).not_to include(*db_load_balancing_logging_keys)
expect(log_data).not_to include(*db_load_balancing_logging_keys)
end
end
end
end
......
......@@ -32,10 +32,6 @@ RSpec.describe API::Helpers do
helper
end
before do
allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
end
it 'handles sticking when a user could be found' do
allow_any_instance_of(API::Helpers).to receive(:initial_current_user).and_return(user)
......
......@@ -31,35 +31,22 @@ RSpec.describe Gitlab::Checks::MatchingMergeRequest do
expect(matcher.match?).to be false
end
context 'with load balancing disabled', :request_store, :redis do
before do
expect(::Gitlab::Database::LoadBalancing).to receive(:enable?).at_least(:once).and_return(false)
expect(::Gitlab::Database::LoadBalancing::Sticking).not_to receive(:unstick_or_continue_sticking)
expect(::Gitlab::Database::LoadBalancing::Sticking).not_to receive(:select_valid_replicas)
end
it 'does not attempt to stick to primary' do
expect(subject.match?).to be true
end
it 'increments no counters' do
expect { subject.match? }
.to change { total_counter.get }.by(0)
.and change { stale_counter.get }.by(0)
end
end
context 'with load balancing enabled', :db_load_balancing do
context 'with load balancing enabled' do
let(:session) { ::Gitlab::Database::LoadBalancing::Session.current }
let(:all_caught_up) { true }
before do
Gitlab::Database::LoadBalancing::Session.clear_session
allow(::Gitlab::Database::LoadBalancing::Sticking).to receive(:all_caught_up?).and_return(all_caught_up)
expect(::Gitlab::Database::LoadBalancing::Sticking).to receive(:select_valid_host).with(:project, project.id).and_call_original
allow(::Gitlab::Database::LoadBalancing::Sticking).to receive(:select_caught_up_replicas).with(:project, project.id).and_return(all_caught_up)
end
after do
Gitlab::Database::LoadBalancing::Session.clear_session
end
shared_examples 'secondary that has caught up to a primary' do
it 'continues to use the secondary' do
expect(session.use_primary?).to be false
......
......@@ -122,7 +122,7 @@ RSpec.describe Gitlab::Database::BulkUpdate do
stub_const('ActiveRecordBasePreparedStatementsInverted', klass)
c = ActiveRecord::Base.connection.instance_variable_get(:@config)
c = ActiveRecord::Base.retrieve_connection.instance_variable_get(:@config)
inverted = c.merge(prepared_statements: !ActiveRecord::Base.connection.prepared_statements)
ActiveRecordBasePreparedStatementsInverted.establish_connection(inverted)
......
......@@ -7,7 +7,7 @@ RSpec.describe Gitlab::Database::Consistency do
Gitlab::Database::LoadBalancing::Session.current
end
describe '.with_read_consistency' do
describe '.with_read_consistency', :db_load_balancing do
it 'sticks to primary database' do
expect(session).not_to be_using_primary
......
......@@ -20,11 +20,6 @@ RSpec.describe Gitlab::Database::LoadBalancing::RackMiddleware, :redis do
end
describe '.stick_or_unstick' do
before do
allow(Gitlab::Database::LoadBalancing).to receive(:enable?)
.and_return(true)
end
it 'sticks or unsticks a single object and updates the Rack environment' do
expect(Gitlab::Database::LoadBalancing::Sticking)
.to receive(:unstick_or_continue_sticking)
......
......@@ -5,14 +5,13 @@ require 'spec_helper'
RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
let(:middleware) { described_class.new }
let(:load_balancer) { double.as_null_object }
let(:load_balancer) { Gitlab::Database::LoadBalancing.proxy.load_balancer }
let(:worker_class) { 'TestDataConsistencyWorker' }
let(:job) { { "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e" } }
before do
skip_feature_flags_yaml_validation
skip_default_enabled_yaml_check
allow(::Gitlab::Database::LoadBalancing).to receive_message_chain(:proxy, :load_balancer).and_return(load_balancer)
end
after do
......
......@@ -5,7 +5,7 @@ require 'spec_helper'
RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware do
let(:middleware) { described_class.new }
let(:load_balancer) { double.as_null_object }
let(:load_balancer) { Gitlab::Database::LoadBalancing.proxy.load_balancer }
let(:worker) { worker_class.new }
let(:job) { { "retry" => 3, "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e", 'database_replica_location' => '0/D525E3A8' } }
......@@ -13,7 +13,6 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware do
before do
skip_feature_flags_yaml_validation
skip_default_enabled_yaml_check
allow(::Gitlab::Database::LoadBalancing).to receive_message_chain(:proxy, :load_balancer).and_return(load_balancer)
replication_lag!(false)
end
......
......@@ -8,39 +8,24 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
end
describe '.stick_if_necessary' do
context 'when sticking is disabled' do
it 'does not perform any sticking' do
expect(described_class).not_to receive(:stick)
it 'does not stick if no write was performed' do
allow(Gitlab::Database::LoadBalancing::Session.current)
.to receive(:performed_write?)
.and_return(false)
described_class.stick_if_necessary(:user, 42)
end
end
context 'when sticking is enabled' do
before do
allow(Gitlab::Database::LoadBalancing).to receive(:enable?)
.and_return(true)
end
it 'does not stick if no write was performed' do
allow(Gitlab::Database::LoadBalancing::Session.current)
.to receive(:performed_write?)
.and_return(false)
expect(described_class).not_to receive(:stick)
expect(described_class).not_to receive(:stick)
described_class.stick_if_necessary(:user, 42)
end
described_class.stick_if_necessary(:user, 42)
end
it 'sticks to the primary if a write was performed' do
allow(Gitlab::Database::LoadBalancing::Session.current)
.to receive(:performed_write?)
.and_return(true)
it 'sticks to the primary if a write was performed' do
allow(Gitlab::Database::LoadBalancing::Session.current)
.to receive(:performed_write?)
.and_return(true)
expect(described_class).to receive(:stick).with(:user, 42)
expect(described_class).to receive(:stick).with(:user, 42)
described_class.stick_if_necessary(:user, 42)
end
described_class.stick_if_necessary(:user, 42)
end
end
......@@ -155,35 +140,22 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
end
RSpec.shared_examples 'sticking' do
context 'when sticking is disabled' do
it 'does not perform any sticking', :aggregate_failures do
expect(described_class).not_to receive(:set_write_location_for)
expect(Gitlab::Database::LoadBalancing::Session.current).not_to receive(:use_primary!)
before do
lb = double(:lb, primary_write_location: 'foo')
described_class.bulk_stick(:user, ids)
end
allow(described_class).to receive(:load_balancer).and_return(lb)
end
context 'when sticking is enabled' do
before do
allow(Gitlab::Database::LoadBalancing).to receive(:configured?).and_return(true)
lb = double(:lb, primary_write_location: 'foo')
allow(described_class).to receive(:load_balancer).and_return(lb)
it 'sticks an entity to the primary', :aggregate_failures do
ids.each do |id|
expect(described_class).to receive(:set_write_location_for)
.with(:user, id, 'foo')
end
it 'sticks an entity to the primary', :aggregate_failures do
ids.each do |id|
expect(described_class).to receive(:set_write_location_for)
.with(:user, id, 'foo')
end
expect(Gitlab::Database::LoadBalancing::Session.current)
.to receive(:use_primary!)
expect(Gitlab::Database::LoadBalancing::Session.current)
.to receive(:use_primary!)
subject
end
subject
end
end
......@@ -202,63 +174,15 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
end
describe '.mark_primary_write_location' do
context 'when enabled' do
before do
allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
allow(Gitlab::Database::LoadBalancing).to receive(:configured?).and_return(true)
end
it 'updates the write location with the load balancer' do
lb = double(:lb, primary_write_location: 'foo')
it 'updates the write location with the load balancer' do
lb = double(:lb, primary_write_location: 'foo')
allow(described_class).to receive(:load_balancer).and_return(lb)
expect(described_class).to receive(:set_write_location_for)
.with(:user, 42, 'foo')
described_class.mark_primary_write_location(:user, 42)
end
end
context 'when load balancing is configured but not enabled' do
before do
allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(false)
allow(Gitlab::Database::LoadBalancing).to receive(:configured?).and_return(true)
end
it 'updates the write location with the main ActiveRecord connection' do
allow(described_class).to receive(:load_balancer).and_return(nil)
expect(ActiveRecord::Base).to receive(:connection).and_call_original
expect(described_class).to receive(:set_write_location_for)
.with(:user, 42, anything)
described_class.mark_primary_write_location(:user, 42)
end
context 'when write location is nil' do
before do
allow(Gitlab::Database.main).to receive(:get_write_location).and_return(nil)
end
it 'does not update the write location' do
expect(described_class).not_to receive(:set_write_location_for)
described_class.mark_primary_write_location(:user, 42)
end
end
end
context 'when load balancing is disabled' do
before do
allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(false)
allow(Gitlab::Database::LoadBalancing).to receive(:configured?).and_return(false)
end
allow(described_class).to receive(:load_balancer).and_return(lb)
it 'updates the write location with the main ActiveRecord connection' do
expect(described_class).not_to receive(:set_write_location_for)
expect(described_class).to receive(:set_write_location_for)
.with(:user, 42, 'foo')
described_class.mark_primary_write_location(:user, 42)
end
described_class.mark_primary_write_location(:user, 42)
end
end
......
......@@ -4,38 +4,14 @@ require 'spec_helper'
RSpec.describe Gitlab::Database::LoadBalancing do
describe '.proxy' do
before do
@previous_proxy = ActiveRecord::Base.load_balancing_proxy
it 'returns the connection proxy' do
proxy = double(:connection_proxy)
ActiveRecord::Base.load_balancing_proxy = connection_proxy
end
after do
ActiveRecord::Base.load_balancing_proxy = @previous_proxy
end
context 'when configured' do
let(:connection_proxy) { double(:connection_proxy) }
it 'returns the connection proxy' do
expect(subject.proxy).to eq(connection_proxy)
end
end
context 'when not configured' do
let(:connection_proxy) { nil }
it 'returns nil' do
expect(subject.proxy).to be_nil
end
it 'tracks an error to sentry' do
expect(Gitlab::ErrorTracking).to receive(:track_exception).with(
an_instance_of(subject::ProxyNotConfiguredError)
)
allow(ActiveRecord::Base)
.to receive(:load_balancing_proxy)
.and_return(proxy)
subject.proxy
end
expect(described_class.proxy).to eq(proxy)
end
end
......@@ -50,46 +26,53 @@ RSpec.describe Gitlab::Database::LoadBalancing do
end
end
describe '.enable?' do
before do
allow(described_class.configuration)
.to receive(:hosts)
.and_return(%w(foo))
end
it 'returns false when no hosts are specified' do
allow(described_class.configuration).to receive(:hosts).and_return([])
describe '.enable_replicas?' do
context 'when hosts are specified' do
before do
allow(described_class.configuration)
.to receive(:hosts)
.and_return(%w(foo))
end
expect(described_class.enable?).to eq(false)
end
it 'returns true' do
expect(described_class.enable_replicas?).to eq(true)
end
it 'returns true when Sidekiq is being used' do
allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(true)
it 'returns true when Sidekiq is being used' do
allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(true)
expect(described_class.enable?).to eq(true)
end
expect(described_class.enable_replicas?).to eq(true)
end
it 'returns false when running inside a Rake task' do
allow(Gitlab::Runtime).to receive(:rake?).and_return(true)
it 'returns false when running inside a Rake task' do
allow(Gitlab::Runtime).to receive(:rake?).and_return(true)
expect(described_class.enable?).to eq(false)
expect(described_class.enable_replicas?).to eq(false)
end
end
it 'returns true when load balancing should be enabled' do
allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(false)
context 'when no hosts are specified but service discovery is enabled' do
it 'returns true' do
allow(described_class.configuration).to receive(:hosts).and_return([])
allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(false)
expect(described_class.enable?).to eq(true)
end
allow(described_class.configuration)
.to receive(:service_discovery_enabled?)
.and_return(true)
it 'returns true when service discovery is enabled' do
allow(described_class.configuration).to receive(:hosts).and_return([])
allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(false)
expect(described_class.enable_replicas?).to eq(true)
end
end
allow(described_class.configuration)
.to receive(:service_discovery_enabled?)
.and_return(true)
context 'when no hosts are specified and service discovery is disabled' do
it 'returns false' do
allow(described_class.configuration).to receive(:hosts).and_return([])
allow(described_class.configuration)
.to receive(:service_discovery_enabled?)
.and_return(false)
expect(described_class.enable?).to eq(true)
expect(described_class.enable_replicas?).to eq(false)
end
end
end
......@@ -121,41 +104,6 @@ RSpec.describe Gitlab::Database::LoadBalancing do
end
end
describe '.configure_proxy' do
before do
allow(ActiveRecord::Base).to receive(:load_balancing_proxy=)
end
it 'configures the connection proxy' do
described_class.configure_proxy
expect(ActiveRecord::Base).to have_received(:load_balancing_proxy=)
.with(Gitlab::Database::LoadBalancing::ConnectionProxy)
end
context 'when service discovery is enabled' do
it 'runs initial service discovery when configuring the connection proxy' do
discover = instance_spy(Gitlab::Database::LoadBalancing::ServiceDiscovery)
allow(described_class.configuration)
.to receive(:service_discovery)
.and_return({ record: 'foo' })
expect(Gitlab::Database::LoadBalancing::ServiceDiscovery)
.to receive(:new)
.with(
an_instance_of(Gitlab::Database::LoadBalancing::LoadBalancer),
an_instance_of(Hash)
)
.and_return(discover)
expect(discover).to receive(:perform_service_discovery)
described_class.configure_proxy
end
end
end
describe '.start_service_discovery' do
it 'does not start if service discovery is disabled' do
expect(Gitlab::Database::LoadBalancing::ServiceDiscovery)
......@@ -191,15 +139,41 @@ RSpec.describe Gitlab::Database::LoadBalancing do
end
end
describe '.db_role_for_connection' do
context 'when the load balancing is not configured' do
let(:connection) { ActiveRecord::Base.connection }
describe '.perform_service_discovery' do
it 'does nothing if service discovery is disabled' do
expect(Gitlab::Database::LoadBalancing::ServiceDiscovery)
.not_to receive(:new)
it 'returns primary' do
expect(described_class.db_role_for_connection(connection)).to eq(:primary)
end
described_class.perform_service_discovery
end
it 'performs service discovery when enabled' do
allow(described_class.configuration)
.to receive(:service_discovery_enabled?)
.and_return(true)
sv = instance_spy(Gitlab::Database::LoadBalancing::ServiceDiscovery)
cfg = Gitlab::Database::LoadBalancing::Configuration
.new(ActiveRecord::Base)
lb = Gitlab::Database::LoadBalancing::LoadBalancer.new(cfg)
proxy = Gitlab::Database::LoadBalancing::ConnectionProxy.new(lb)
allow(described_class)
.to receive(:proxy)
.and_return(proxy)
expect(Gitlab::Database::LoadBalancing::ServiceDiscovery)
.to receive(:new)
.with(lb, cfg.service_discovery)
.and_return(sv)
expect(sv).to receive(:perform_service_discovery)
described_class.perform_service_discovery
end
end
describe '.db_role_for_connection' do
context 'when the NullPool is used for connection' do
let(:pool) { ActiveRecord::ConnectionAdapters::NullPool.new }
let(:connection) { double(:connection, pool: pool) }
......@@ -274,10 +248,6 @@ RSpec.describe Gitlab::Database::LoadBalancing do
end
end
before do
model.singleton_class.prepend ::Gitlab::Database::LoadBalancing::ActiveRecordProxy
end
where(:queries, :include_transaction, :expected_results) do
[
# Read methods
......
......@@ -7,7 +7,7 @@ RSpec.describe Gitlab::Database::WithLockRetries do
let(:logger) { Gitlab::Database::WithLockRetries::NULL_LOGGER }
let(:subject) { described_class.new(env: env, logger: logger, allow_savepoints: allow_savepoints, timing_configuration: timing_configuration) }
let(:allow_savepoints) { true }
let(:connection) { ActiveRecord::Base.connection }
let(:connection) { ActiveRecord::Base.retrieve_connection }
let(:timing_configuration) do
[
......
......@@ -185,10 +185,23 @@ RSpec.describe Gitlab::Database do
describe '.db_config_name' do
it 'returns the db_config name for the connection' do
connection = ActiveRecord::Base.connection
model = ActiveRecord::Base
expect(described_class.db_config_name(connection)).to be_a(String)
expect(described_class.db_config_name(connection)).to eq(connection.pool.db_config.name)
# This is a ConnectionProxy
expect(described_class.db_config_name(model.connection))
.to eq('unknown')
# This is an actual connection
expect(described_class.db_config_name(model.retrieve_connection))
.to eq('main')
end
context 'when replicas are configured', :db_load_balancing do
it 'returns the name for a replica' do
replica = ActiveRecord::Base.connection.load_balancer.host
expect(described_class.db_config_name(replica)).to eq('main_replica')
end
end
end
......@@ -279,7 +292,7 @@ RSpec.describe Gitlab::Database do
expect(event).not_to be_nil
expect(event.duration).to be > 0.0
expect(event.payload).to a_hash_including(
connection: be_a(ActiveRecord::ConnectionAdapters::AbstractAdapter)
connection: be_a(Gitlab::Database::LoadBalancing::ConnectionProxy)
)
end
end
......@@ -296,7 +309,7 @@ RSpec.describe Gitlab::Database do
expect(event).not_to be_nil
expect(event.duration).to be > 0.0
expect(event.payload).to a_hash_including(
connection: be_a(ActiveRecord::ConnectionAdapters::AbstractAdapter)
connection: be_a(Gitlab::Database::LoadBalancing::ConnectionProxy)
)
end
end
......@@ -319,7 +332,7 @@ RSpec.describe Gitlab::Database do
expect(event).not_to be_nil
expect(event.duration).to be > 0.0
expect(event.payload).to a_hash_including(
connection: be_a(ActiveRecord::ConnectionAdapters::AbstractAdapter)
connection: be_a(Gitlab::Database::LoadBalancing::ConnectionProxy)
)
end
end
......@@ -340,7 +353,7 @@ RSpec.describe Gitlab::Database do
expect(event).not_to be_nil
expect(event.duration).to be > 0.0
expect(event.payload).to a_hash_including(
connection: be_a(ActiveRecord::ConnectionAdapters::AbstractAdapter)
connection: be_a(Gitlab::Database::LoadBalancing::ConnectionProxy)
)
end
end
......
......@@ -158,26 +158,15 @@ RSpec.describe Gitlab::ImportExport::Json::StreamingSerializer do
end
describe 'load balancing' do
context 'when feature flag load_balancing_for_export_workers is enabled' do
context 'when feature flag load_balancing_for_export_workers is enabled', :db_load_balancing do
before do
stub_feature_flags(load_balancing_for_export_workers: true)
end
context 'when enabled', :db_load_balancing do
it 'reads from replica' do
expect(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_replicas_for_read_queries).and_call_original
it 'reads from replica' do
expect(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_replicas_for_read_queries).and_call_original
subject.execute
end
end
context 'when disabled' do
it 'reads from primary' do
allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(false)
expect(Gitlab::Database::LoadBalancing::Session.current).not_to receive(:use_replicas_for_read_queries)
subject.execute
end
subject.execute
end
end
......
......@@ -107,69 +107,44 @@ RSpec.describe Gitlab::InstrumentationHelper do
end
end
context 'when load balancing is enabled' do
before do
allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
end
it 'includes DB counts' do
subject
expect(payload).to include(db_replica_count: 0,
db_replica_cached_count: 0,
db_primary_count: 0,
db_primary_cached_count: 0,
db_primary_wal_count: 0,
db_replica_wal_count: 0,
db_primary_wal_cached_count: 0,
db_replica_wal_cached_count: 0)
end
context 'when replica caught up search was made' do
before do
Gitlab::SafeRequestStore[:caught_up_replica_pick_ok] = 2
Gitlab::SafeRequestStore[:caught_up_replica_pick_fail] = 1
end
it 'includes DB counts' do
subject
it 'includes related metrics' do
subject
expect(payload).to include(db_replica_count: 0,
db_replica_cached_count: 0,
db_primary_count: 0,
db_primary_cached_count: 0,
db_primary_wal_count: 0,
db_replica_wal_count: 0,
db_primary_wal_cached_count: 0,
db_replica_wal_cached_count: 0)
end
expect(payload).to include(caught_up_replica_pick_ok: 2)
expect(payload).to include(caught_up_replica_pick_fail: 1)
end
context 'when replica caught up search was made' do
before do
Gitlab::SafeRequestStore[:caught_up_replica_pick_ok] = 2
Gitlab::SafeRequestStore[:caught_up_replica_pick_fail] = 1
end
context 'when only a single counter was updated' do
before do
Gitlab::SafeRequestStore[:caught_up_replica_pick_ok] = 1
Gitlab::SafeRequestStore[:caught_up_replica_pick_fail] = nil
end
it 'includes only that counter into logging' do
subject
it 'includes related metrics' do
subject
expect(payload).to include(caught_up_replica_pick_ok: 1)
expect(payload).not_to include(:caught_up_replica_pick_fail)
end
expect(payload).to include(caught_up_replica_pick_ok: 2)
expect(payload).to include(caught_up_replica_pick_fail: 1)
end
end
context 'when load balancing is disabled' do
context 'when only a single counter was updated' do
before do
allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(false)
Gitlab::SafeRequestStore[:caught_up_replica_pick_ok] = 1
Gitlab::SafeRequestStore[:caught_up_replica_pick_fail] = nil
end
it 'does not include DB counts' do
it 'includes only that counter into logging' do
subject
expect(payload).not_to include(db_replica_count: 0,
db_replica_cached_count: 0,
db_primary_count: 0,
db_primary_cached_count: 0,
db_primary_wal_count: 0,
db_replica_wal_count: 0,
db_primary_wal_cached_count: 0,
db_replica_wal_cached_count: 0)
expect(payload).to include(caught_up_replica_pick_ok: 1)
expect(payload).not_to include(:caught_up_replica_pick_fail)
end
end
end
......
......@@ -7,7 +7,7 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
let(:env) { {} }
let(:subscriber) { described_class.new }
let(:connection) { ActiveRecord::Base.connection }
let(:connection) { ActiveRecord::Base.retrieve_connection }
let(:db_config_name) { ::Gitlab::Database.db_config_name(connection) }
describe '#transaction' do
......@@ -135,7 +135,7 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
end
it_behaves_like 'record ActiveRecord metrics'
it_behaves_like 'store ActiveRecord info in RequestStore'
it_behaves_like 'store ActiveRecord info in RequestStore', :primary
end
end
......@@ -195,11 +195,7 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
with_them do
let(:payload) { { name: name, sql: sql(sql_query, comments: comments), connection: connection } }
before do
allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
end
context 'query using a connection to a replica' do
context 'query using a connection to a replica', :db_load_balancing do
before do
allow(Gitlab::Database::LoadBalancing).to receive(:db_role_for_connection).and_return(:replica)
end
......
......@@ -5,10 +5,6 @@ require 'spec_helper'
RSpec.describe Gitlab::Metrics::Subscribers::LoadBalancing, :request_store do
let(:subscriber) { described_class.new }
before do
allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
end
describe '#caught_up_replica_pick' do
shared_examples 'having payload result value' do |result, counter_name|
subject { subscriber.caught_up_replica_pick(event) }
......
......@@ -243,8 +243,22 @@ RSpec.describe Gitlab::SidekiqLogging::StructuredLogger do
expected_end_payload.merge(
'db_duration_s' => a_value >= 0.1,
'db_count' => a_value >= 1,
'db_cached_count' => 0,
'db_write_count' => 0
"db_replica_#{db_config_name}_count" => 0,
'db_replica_duration_s' => a_value >= 0,
'db_primary_count' => a_value >= 1,
"db_primary_#{db_config_name}_count" => a_value >= 1,
'db_primary_duration_s' => a_value > 0,
"db_primary_#{db_config_name}_duration_s" => a_value > 0
)
end
let(:end_payload) do
start_payload.merge(db_payload_defaults).merge(
'message' => 'TestWorker JID-da883554ee4fe414012f5f42: done: 0.0 sec',
'job_status' => 'done',
'duration_s' => 0.0,
'completed_at' => timestamp.to_f,
'cpu_s' => 1.111112
)
end
......@@ -274,59 +288,9 @@ RSpec.describe Gitlab::SidekiqLogging::StructuredLogger do
end
end
context 'when load balancing is disabled' do
before do
allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(false)
end
let(:expected_end_payload_with_db) do
expected_end_payload.merge(
'db_duration_s' => a_value >= 0.1,
'db_count' => a_value >= 1,
'db_cached_count' => 0,
'db_write_count' => 0
)
end
include_examples 'performs database queries'
end
context 'when load balancing is enabled', :db_load_balancing do
let(:db_config_name) { ::Gitlab::Database.db_config_name(ApplicationRecord.connection) }
let(:expected_db_payload_defaults) do
metrics =
::Gitlab::Metrics::Subscribers::ActiveRecord.load_balancing_metric_counter_keys +
::Gitlab::Metrics::Subscribers::ActiveRecord.load_balancing_metric_duration_keys +
::Gitlab::Metrics::Subscribers::ActiveRecord.db_counter_keys +
[:db_duration_s]
metrics.each_with_object({}) do |key, result|
result[key.to_s] = 0
end
end
let(:expected_end_payload_with_db) do
expected_end_payload.merge(expected_db_payload_defaults).merge(
'db_duration_s' => a_value >= 0.1,
'db_count' => a_value >= 1,
"db_replica_#{db_config_name}_count" => 0,
'db_replica_duration_s' => a_value >= 0,
'db_primary_count' => a_value >= 1,
"db_primary_#{db_config_name}_count" => a_value >= 1,
'db_primary_duration_s' => a_value > 0,
"db_primary_#{db_config_name}_duration_s" => a_value > 0
)
end
let(:end_payload) do
start_payload.merge(expected_db_payload_defaults).merge(
'message' => 'TestWorker JID-da883554ee4fe414012f5f42: done: 0.0 sec',
'job_status' => 'done',
'duration_s' => 0.0,
'completed_at' => timestamp.to_f,
'cpu_s' => 1.111112
)
let(:db_config_name) do
::Gitlab::Database.db_config_name(ApplicationRecord.retrieve_connection)
end
include_examples 'performs database queries'
......
......@@ -250,60 +250,30 @@ RSpec.describe Gitlab::SidekiqMiddleware::ServerMetrics do
end
end
context 'when load_balancing is enabled' do
before do
allow(::Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
end
describe '#call' do
context 'when worker declares data consistency' do
include_context 'worker declaring data consistency'
it 'increments load balancing counter with defined data consistency' do
process_job
expect(load_balancing_metric).to have_received(:increment).with(
a_hash_including(
data_consistency: :delayed,
load_balancing_strategy: 'replica'
), 1)
end
end
context 'when worker does not declare data consistency' do
it 'increments load balancing counter with default data consistency' do
process_job
expect(load_balancing_metric).to have_received(:increment).with(
a_hash_including(
data_consistency: :always,
load_balancing_strategy: 'primary'
), 1)
end
end
end
end
context 'when load_balancing is disabled' do
include_context 'worker declaring data consistency'
describe '#call' do
context 'when worker declares data consistency' do
include_context 'worker declaring data consistency'
before do
allow(::Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(false)
end
describe '#initialize' do
it 'does not set load_balancing metrics' do
expect(Gitlab::Metrics).not_to receive(:counter).with(:sidekiq_load_balancing_count, anything)
it 'increments load balancing counter with defined data consistency' do
process_job
subject
expect(load_balancing_metric).to have_received(:increment).with(
a_hash_including(
data_consistency: :delayed,
load_balancing_strategy: 'replica'
), 1)
end
end
describe '#call' do
it 'does not increment load balancing counter' do
context 'when worker does not declare data consistency' do
it 'increments load balancing counter with default data consistency' do
process_job
expect(load_balancing_metric).not_to have_received(:increment)
expect(load_balancing_metric).to have_received(:increment).with(
a_hash_including(
data_consistency: :always,
load_balancing_strategy: 'primary'
), 1)
end
end
end
......
......@@ -28,9 +28,8 @@ RSpec.describe Gitlab::SidekiqMiddleware do
stub_const('TestWorker', worker_class)
end
shared_examples "a middleware chain" do |load_balancing_enabled|
shared_examples "a middleware chain" do
before do
allow(::Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(load_balancing_enabled)
configurator.call(chain)
end
......@@ -45,10 +44,10 @@ RSpec.describe Gitlab::SidekiqMiddleware do
end
end
shared_examples "a middleware chain for mailer" do |load_balancing_enabled|
shared_examples "a middleware chain for mailer" do
let(:worker_class) { ActiveJob::QueueAdapters::SidekiqAdapter::JobWrapper }
it_behaves_like "a middleware chain", load_balancing_enabled
it_behaves_like "a middleware chain"
end
describe '.server_configurator' do
......@@ -105,25 +104,8 @@ RSpec.describe Gitlab::SidekiqMiddleware do
end
context "all optional middlewares on" do
context "when load balancing is enabled" do
before do
allow(::Gitlab::Database::LoadBalancing).to receive_message_chain(:proxy, :load_balancer, :release_host)
end
it_behaves_like "a middleware chain", true
it_behaves_like "a middleware chain for mailer", true
end
context "when load balancing is disabled" do
let(:disabled_sidekiq_middlewares) do
[
Gitlab::Database::LoadBalancing::SidekiqServerMiddleware
]
end
it_behaves_like "a middleware chain", false
it_behaves_like "a middleware chain for mailer", false
end
it_behaves_like "a middleware chain"
it_behaves_like "a middleware chain for mailer"
end
context "all optional middlewares off" do
......@@ -135,36 +117,16 @@ RSpec.describe Gitlab::SidekiqMiddleware do
)
end
context "when load balancing is enabled" do
let(:disabled_sidekiq_middlewares) do
[
Gitlab::SidekiqMiddleware::ServerMetrics,
Gitlab::SidekiqMiddleware::ArgumentsLogger,
Gitlab::SidekiqMiddleware::MemoryKiller
]
end
before do
allow(::Gitlab::Database::LoadBalancing).to receive_message_chain(:proxy, :load_balancer, :release_host)
end
it_behaves_like "a middleware chain", true
it_behaves_like "a middleware chain for mailer", true
let(:disabled_sidekiq_middlewares) do
[
Gitlab::SidekiqMiddleware::ServerMetrics,
Gitlab::SidekiqMiddleware::ArgumentsLogger,
Gitlab::SidekiqMiddleware::MemoryKiller
]
end
context "when load balancing is disabled" do
let(:disabled_sidekiq_middlewares) do
[
Gitlab::SidekiqMiddleware::ServerMetrics,
Gitlab::SidekiqMiddleware::ArgumentsLogger,
Gitlab::SidekiqMiddleware::MemoryKiller,
Gitlab::Database::LoadBalancing::SidekiqServerMiddleware
]
end
it_behaves_like "a middleware chain", false
it_behaves_like "a middleware chain for mailer", false
end
it_behaves_like "a middleware chain"
it_behaves_like "a middleware chain for mailer"
end
end
......@@ -186,30 +148,7 @@ RSpec.describe Gitlab::SidekiqMiddleware do
]
end
context "when load balancing is disabled" do
let(:disabled_sidekiq_middlewares) do
[
Gitlab::Database::LoadBalancing::SidekiqClientMiddleware
]
end
it_behaves_like "a middleware chain", false
it_behaves_like "a middleware chain for mailer", false
# Sidekiq documentation states that the worker class could be a string
# or a class reference. We should test for both
context "worker_class as string value" do
let(:worker_args) { [worker_class.to_s, { 'args' => job_args }, queue, redis_pool] }
let(:middleware_expected_args) { [worker_class.to_s, hash_including({ 'args' => job_args }), queue, redis_pool] }
it_behaves_like "a middleware chain", false
it_behaves_like "a middleware chain for mailer", false
end
end
context "when load balancing is enabled" do
it_behaves_like "a middleware chain", true
it_behaves_like "a middleware chain for mailer", true
end
it_behaves_like "a middleware chain"
it_behaves_like "a middleware chain for mailer"
end
end
......@@ -53,154 +53,82 @@ RSpec.describe Peek::Views::ActiveRecord, :request_store do
allow(connection_primary_2).to receive(:transaction_open?).and_return(true)
allow(connection_unknown).to receive(:transaction_open?).and_return(false)
allow(::Gitlab::Database).to receive(:db_config_name).and_return('the_db_config_name')
allow(Gitlab::Database::LoadBalancing).to receive(:db_role_for_connection).with(connection_replica).and_return(:replica)
allow(Gitlab::Database::LoadBalancing).to receive(:db_role_for_connection).with(connection_primary_1).and_return(:primary)
allow(Gitlab::Database::LoadBalancing).to receive(:db_role_for_connection).with(connection_primary_2).and_return(:primary)
allow(Gitlab::Database::LoadBalancing).to receive(:db_role_for_connection).with(connection_unknown).and_return(nil)
end
context 'when database load balancing is not enabled' do
it 'subscribes and store data into peek views' do
Timecop.freeze(2021, 2, 23, 10, 0) do
ActiveSupport::Notifications.publish('sql.active_record', Time.current, Time.current + 1.second, '1', event_1)
ActiveSupport::Notifications.publish('sql.active_record', Time.current, Time.current + 2.seconds, '2', event_2)
ActiveSupport::Notifications.publish('sql.active_record', Time.current, Time.current + 3.seconds, '3', event_3)
ActiveSupport::Notifications.publish('sql.active_record', Time.current, Time.current + 4.seconds, '4', event_4)
end
it 'includes db role data and db_config_name name' do
Timecop.freeze(2021, 2, 23, 10, 0) do
ActiveSupport::Notifications.publish('sql.active_record', Time.current, Time.current + 1.second, '1', event_1)
ActiveSupport::Notifications.publish('sql.active_record', Time.current, Time.current + 2.seconds, '2', event_2)
ActiveSupport::Notifications.publish('sql.active_record', Time.current, Time.current + 3.seconds, '3', event_3)
ActiveSupport::Notifications.publish('sql.active_record', Time.current, Time.current + 4.seconds, '4', event_4)
end
expect(subject.results).to match(
calls: 4,
summary: {
"Cached" => 1,
"In a transaction" => 1
},
duration: '10000.00ms',
warnings: ["active-record duration: 10000.0 over 3000"],
details: contain_exactly(
a_hash_including(
start: be_a(Time),
cached: '',
transaction: '',
duration: 1000.0,
sql: 'SELECT * FROM users WHERE id = 10',
db_config_name: "Config name: the_db_config_name"
),
a_hash_including(
start: be_a(Time),
cached: 'Cached',
transaction: '',
duration: 2000.0,
sql: 'SELECT * FROM users WHERE id = 10',
db_config_name: "Config name: the_db_config_name"
),
a_hash_including(
start: be_a(Time),
cached: '',
transaction: 'In a transaction',
duration: 3000.0,
sql: 'UPDATE users SET admin = true WHERE id = 10',
db_config_name: "Config name: the_db_config_name"
),
a_hash_including(
start: be_a(Time),
cached: '',
transaction: '',
duration: 4000.0,
sql: 'SELECT VERSION()',
db_config_name: "Config name: the_db_config_name"
)
expect(subject.results).to match(
calls: 4,
summary: {
"Cached" => 1,
"In a transaction" => 1,
"Role: Primary" => 2,
"Role: Replica" => 1,
"Role: Unknown" => 1
},
duration: '10000.00ms',
warnings: ["active-record duration: 10000.0 over 3000"],
details: contain_exactly(
a_hash_including(
start: be_a(Time),
cached: '',
transaction: '',
duration: 1000.0,
sql: 'SELECT * FROM users WHERE id = 10',
db_role: 'Role: Primary',
db_config_name: "Config name: the_db_config_name"
),
a_hash_including(
start: be_a(Time),
cached: 'Cached',
transaction: '',
duration: 2000.0,
sql: 'SELECT * FROM users WHERE id = 10',
db_role: 'Role: Replica',
db_config_name: "Config name: the_db_config_name"
),
a_hash_including(
start: be_a(Time),
cached: '',
transaction: 'In a transaction',
duration: 3000.0,
sql: 'UPDATE users SET admin = true WHERE id = 10',
db_role: 'Role: Primary',
db_config_name: "Config name: the_db_config_name"
),
a_hash_including(
start: be_a(Time),
cached: '',
transaction: '',
duration: 4000.0,
sql: 'SELECT VERSION()',
db_role: 'Role: Unknown',
db_config_name: "Config name: the_db_config_name"
)
)
end
context 'when the GITLAB_MULTIPLE_DATABASE_METRICS env var is disabled' do
before do
stub_env('GITLAB_MULTIPLE_DATABASE_METRICS', nil)
end
it 'does not include db_config_name field' do
ActiveSupport::Notifications.publish('sql.active_record', Time.current, Time.current + 1.second, '1', event_1)
expect(subject.results[:details][0][:db_config_name]).to be_nil
end
end
)
end
context 'when database load balancing is enabled' do
context 'when the GITLAB_MULTIPLE_DATABASE_METRICS env var is disabled' do
before do
allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
allow(Gitlab::Database::LoadBalancing).to receive(:db_role_for_connection).with(connection_replica).and_return(:replica)
allow(Gitlab::Database::LoadBalancing).to receive(:db_role_for_connection).with(connection_primary_1).and_return(:primary)
allow(Gitlab::Database::LoadBalancing).to receive(:db_role_for_connection).with(connection_primary_2).and_return(:primary)
allow(Gitlab::Database::LoadBalancing).to receive(:db_role_for_connection).with(connection_unknown).and_return(nil)
stub_env('GITLAB_MULTIPLE_DATABASE_METRICS', nil)
end
it 'includes db role data and db_config_name name' do
Timecop.freeze(2021, 2, 23, 10, 0) do
ActiveSupport::Notifications.publish('sql.active_record', Time.current, Time.current + 1.second, '1', event_1)
ActiveSupport::Notifications.publish('sql.active_record', Time.current, Time.current + 2.seconds, '2', event_2)
ActiveSupport::Notifications.publish('sql.active_record', Time.current, Time.current + 3.seconds, '3', event_3)
ActiveSupport::Notifications.publish('sql.active_record', Time.current, Time.current + 4.seconds, '4', event_4)
end
expect(subject.results).to match(
calls: 4,
summary: {
"Cached" => 1,
"In a transaction" => 1,
"Role: Primary" => 2,
"Role: Replica" => 1,
"Role: Unknown" => 1
},
duration: '10000.00ms',
warnings: ["active-record duration: 10000.0 over 3000"],
details: contain_exactly(
a_hash_including(
start: be_a(Time),
cached: '',
transaction: '',
duration: 1000.0,
sql: 'SELECT * FROM users WHERE id = 10',
db_role: 'Role: Primary',
db_config_name: "Config name: the_db_config_name"
),
a_hash_including(
start: be_a(Time),
cached: 'Cached',
transaction: '',
duration: 2000.0,
sql: 'SELECT * FROM users WHERE id = 10',
db_role: 'Role: Replica',
db_config_name: "Config name: the_db_config_name"
),
a_hash_including(
start: be_a(Time),
cached: '',
transaction: 'In a transaction',
duration: 3000.0,
sql: 'UPDATE users SET admin = true WHERE id = 10',
db_role: 'Role: Primary',
db_config_name: "Config name: the_db_config_name"
),
a_hash_including(
start: be_a(Time),
cached: '',
transaction: '',
duration: 4000.0,
sql: 'SELECT VERSION()',
db_role: 'Role: Unknown',
db_config_name: "Config name: the_db_config_name"
)
)
)
end
context 'when the GITLAB_MULTIPLE_DATABASE_METRICS env var is disabled' do
before do
stub_env('GITLAB_MULTIPLE_DATABASE_METRICS', nil)
end
it 'does not include db_config_name field' do
ActiveSupport::Notifications.publish('sql.active_record', Time.current, Time.current + 1.second, '1', event_1)
it 'does not include db_config_name field' do
ActiveSupport::Notifications.publish('sql.active_record', Time.current, Time.current + 1.second, '1', event_1)
expect(subject.results[:details][0][:db_config_name]).to be_nil
end
expect(subject.results[:details][0][:db_config_name]).to be_nil
end
end
end
......@@ -194,7 +194,7 @@ RSpec.describe ApplicationRecord do
end
context 'with database load balancing' do
let(:session) { double(:session) }
let(:session) { Gitlab::Database::LoadBalancing::Session.new }
before do
allow(::Gitlab::Database::LoadBalancing::Session).to receive(:current).and_return(session)
......
......@@ -2816,8 +2816,9 @@ RSpec.describe Ci::Pipeline, :mailer, factory_default: :keep do
extra_update_queries = 4 # transition ... => :canceled, queue pop
extra_generic_commit_status_validation_queries = 2 # name_uniqueness_across_types
extra_load_balancer_queries = 3
expect(control2.count).to eq(control1.count + extra_update_queries + extra_generic_commit_status_validation_queries)
expect(control2.count).to eq(control1.count + extra_update_queries + extra_generic_commit_status_validation_queries + extra_load_balancer_queries)
end
end
......
......@@ -383,9 +383,6 @@ RSpec.describe Ci::Runner do
it 'sticks the runner to the primary and calls the original method' do
runner = create(:ci_runner)
allow(Gitlab::Database::LoadBalancing).to receive(:enable?)
.and_return(true)
expect(Gitlab::Database::LoadBalancing::Sticking).to receive(:stick)
.with(:runner, runner.id)
......
......@@ -47,8 +47,6 @@ RSpec.describe UserProjectAccessChangedService do
let(:service) { UserProjectAccessChangedService.new([1, 2]) }
before do
allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
expect(AuthorizedProjectsWorker).to receive(:bulk_perform_and_wait)
.with([[1], [2]])
.and_return(10)
......
......@@ -2,8 +2,6 @@
RSpec.configure do |config|
config.before(:each, :db_load_balancing) do
allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
config = Gitlab::Database::LoadBalancing::Configuration
.new(ActiveRecord::Base, [Gitlab::Database.main.config['host']])
lb = ::Gitlab::Database::LoadBalancing::LoadBalancer.new(config)
......
......@@ -39,17 +39,25 @@ RSpec.shared_context 'structured_logger' do
)
end
let(:db_payload_defaults) do
metrics =
::Gitlab::Metrics::Subscribers::ActiveRecord.load_balancing_metric_counter_keys +
::Gitlab::Metrics::Subscribers::ActiveRecord.load_balancing_metric_duration_keys +
::Gitlab::Metrics::Subscribers::ActiveRecord.db_counter_keys +
[:db_duration_s]
metrics.each_with_object({}) do |key, result|
result[key.to_s] = 0
end
end
let(:end_payload) do
start_payload.merge(
start_payload.merge(db_payload_defaults).merge(
'message' => 'TestWorker JID-da883554ee4fe414012f5f42: done: 0.0 sec',
'job_status' => 'done',
'duration_s' => 0.0,
'completed_at' => timestamp.to_f,
'cpu_s' => 1.111112,
'db_duration_s' => 0.0,
'db_cached_count' => 0,
'db_count' => 0,
'db_write_count' => 0
'cpu_s' => 1.111112
)
end
......
......@@ -15,6 +15,7 @@ RSpec.shared_context 'server metrics with mocked prometheus' do
let(:redis_seconds_metric) { double('redis seconds metric') }
let(:elasticsearch_seconds_metric) { double('elasticsearch seconds metric') }
let(:elasticsearch_requests_total) { double('elasticsearch calls total metric') }
let(:load_balancing_metric) { double('load balancing metric') }
before do
allow(Gitlab::Metrics).to receive(:histogram).and_call_original
......@@ -31,6 +32,7 @@ RSpec.shared_context 'server metrics with mocked prometheus' do
allow(Gitlab::Metrics).to receive(:counter).with(:sidekiq_jobs_retried_total, anything).and_return(retried_total_metric)
allow(Gitlab::Metrics).to receive(:counter).with(:sidekiq_redis_requests_total, anything).and_return(redis_requests_total)
allow(Gitlab::Metrics).to receive(:counter).with(:sidekiq_elasticsearch_requests_total, anything).and_return(elasticsearch_requests_total)
allow(Gitlab::Metrics).to receive(:counter).with(:sidekiq_load_balancing_count, anything).and_return(load_balancing_metric)
allow(Gitlab::Metrics).to receive(:gauge).with(:sidekiq_running_jobs, anything, {}, :all).and_return(running_jobs_metric)
allow(Gitlab::Metrics).to receive(:gauge).with(:sidekiq_concurrency, anything, {}, :all).and_return(concurrency_metric)
......
......@@ -4,14 +4,20 @@ RSpec.shared_examples 'store ActiveRecord info in RequestStore' do |db_role|
let(:db_config_name) { ::Gitlab::Database.db_config_names.first }
let(:expected_payload_defaults) do
result = {}
metrics =
::Gitlab::Metrics::Subscribers::ActiveRecord.load_balancing_metric_counter_keys +
::Gitlab::Metrics::Subscribers::ActiveRecord.load_balancing_metric_duration_keys +
::Gitlab::Metrics::Subscribers::ActiveRecord.db_counter_keys
metrics.each_with_object({}) do |key, result|
metrics.each do |key|
result[key] = 0
end
::Gitlab::Metrics::Subscribers::ActiveRecord.load_balancing_metric_duration_keys.each do |key|
result[key] = 0.0
end
result
end
def transform_hash(hash, another_hash)
......@@ -36,8 +42,8 @@ RSpec.shared_examples 'store ActiveRecord info in RequestStore' do |db_role|
"db_primary_#{db_config_name}_cached_count": record_cached_query ? 1 : 0,
db_primary_count: record_query ? 1 : 0,
"db_primary_#{db_config_name}_count": record_query ? 1 : 0,
db_primary_duration_s: record_query ? 0.002 : 0,
"db_primary_#{db_config_name}_duration_s": record_query ? 0.002 : 0,
db_primary_duration_s: record_query ? 0.002 : 0.0,
"db_primary_#{db_config_name}_duration_s": record_query ? 0.002 : 0.0,
db_primary_wal_count: record_wal_query ? 1 : 0,
"db_primary_#{db_config_name}_wal_count": record_wal_query ? 1 : 0,
db_primary_wal_cached_count: record_wal_query && record_cached_query ? 1 : 0,
......@@ -52,19 +58,29 @@ RSpec.shared_examples 'store ActiveRecord info in RequestStore' do |db_role|
"db_replica_#{db_config_name}_cached_count": record_cached_query ? 1 : 0,
db_replica_count: record_query ? 1 : 0,
"db_replica_#{db_config_name}_count": record_query ? 1 : 0,
db_replica_duration_s: record_query ? 0.002 : 0,
"db_replica_#{db_config_name}_duration_s": record_query ? 0.002 : 0,
db_replica_duration_s: record_query ? 0.002 : 0.0,
"db_replica_#{db_config_name}_duration_s": record_query ? 0.002 : 0.0,
db_replica_wal_count: record_wal_query ? 1 : 0,
"db_replica_#{db_config_name}_wal_count": record_wal_query ? 1 : 0,
db_replica_wal_cached_count: record_wal_query && record_cached_query ? 1 : 0,
"db_replica_#{db_config_name}_wal_cached_count": record_wal_query && record_cached_query ? 1 : 0
})
else
{
transform_hash(expected_payload_defaults, {
db_count: record_query ? 1 : 0,
db_write_count: record_write_query ? 1 : 0,
db_cached_count: record_cached_query ? 1 : 0
}
db_cached_count: record_cached_query ? 1 : 0,
db_primary_cached_count: 0,
"db_primary_#{db_config_name}_cached_count": 0,
db_primary_count: 0,
"db_primary_#{db_config_name}_count": 0,
db_primary_duration_s: 0.0,
"db_primary_#{db_config_name}_duration_s": 0.0,
db_primary_wal_count: 0,
"db_primary_#{db_config_name}_wal_count": 0,
db_primary_wal_cached_count: 0,
"db_primary_#{db_config_name}_wal_cached_count": 0
})
end
expect(described_class.db_counter_payload).to eq(expected)
......@@ -89,7 +105,7 @@ RSpec.shared_examples 'store ActiveRecord info in RequestStore' do |db_role|
end
RSpec.shared_examples 'record ActiveRecord metrics in a metrics transaction' do |db_role|
let(:db_config_name) { ::Gitlab::Database.db_config_name(ApplicationRecord.connection) }
let(:db_config_name) { ::Gitlab::Database.db_config_name(ApplicationRecord.retrieve_connection) }
it 'increments only db counters' do
if record_query
......
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