From 771f94fc7fac48a353c66eab22bc25ccd6f18d54 Mon Sep 17 00:00:00 2001 From: Yorick Peterse Date: Thu, 7 Oct 2021 15:15:53 +0200 Subject: [PATCH] Add support for load balancing multiple databases This adds support for using the database load balancer with multiple databases. Load balancing is applied to two classes: - ActiveRecord::Base - Ci::CiDatabaseRecord Each class has its own load balancer, configuration, service discovery, etc. Load balancing for the CI class is only enabled when a CI configuration exists, as it can reuse the main load balancer when there's no dedicated CI database. Sticking technically supports multiple databases, but in practise we apply the same sticking rules to all databases. This is due to how LoadBalancing::Session is used: there is only one instance per request/Sidekiq job, and it's not aware of what database connections did what. This means that a write to database A will result in GitLab sticking to the primaries of _all_ databases. The choice for this is simple: it requires fewer code changes, and allows us to introduce multiple database support in smaller increments. One change we made to sticking is to turn the Sticking module into a class, and attach an instance to every base module that has its own load balancer. This makes it easier to apply sticking on a per-database level in the future, without having to type `Gitlab::Database::LoadBalancing::Sticking...` every time. Sticking also supports reading and writing of data using the old Redis key names. This ensures sticking continues to work during a deployment, as during this window we'll run two different versions in production. Once the code has been deployed to GitLab.com and has been confirmed to work, we'll remove support for reading/writing the old keys. Sidekiq also supports load balancing multiple databases. If a load balancer/database doesn't have any WAL data in the Sidekiq job, we treat the database as being in sync. This way we can support Sidekiq jobs using both the old and new load balancing data. See https://gitlab.com/gitlab-org/gitlab/-/issues/331776 for more details. Changelog: added --- app/models/ci/build.rb | 2 +- app/models/ci/ci_database_record.rb | 7 - app/models/ci/runner.rb | 2 +- app/models/project.rb | 2 +- app/services/ci/register_job_service.rb | 3 +- .../user_project_access_changed_service.rb | 2 +- config/initializers/load_balancing.rb | 71 +++--- lib/api/ci/helpers/runner.rb | 8 +- lib/api/helpers.rb | 5 +- lib/gitlab/checks/matching_merge_request.rb | 2 +- lib/gitlab/ci/trace.rb | 6 +- lib/gitlab/database.rb | 7 +- lib/gitlab/database/load_balancing.rb | 42 +--- .../load_balancing/action_cable_callbacks.rb | 2 +- .../load_balancing/active_record_proxy.rb | 15 -- .../database/load_balancing/configuration.rb | 7 + .../database/load_balancing/load_balancer.rb | 12 +- .../load_balancing/rack_middleware.rb | 46 ++-- lib/gitlab/database/load_balancing/setup.rb | 61 +++++ .../sidekiq_client_middleware.rb | 27 +-- .../sidekiq_server_middleware.rb | 33 +-- .../database/load_balancing/sticking.rb | 92 +++++--- spec/lib/api/ci/helpers/runner_spec.rb | 16 +- spec/lib/api/helpers_spec.rb | 8 +- .../checks/matching_merge_request_spec.rb | 16 +- spec/lib/gitlab/database/consistency_spec.rb | 10 +- .../action_cable_callbacks_spec.rb | 4 +- .../active_record_proxy_spec.rb | 27 --- .../load_balancing/configuration_spec.rb | 8 + .../load_balancing/load_balancer_spec.rb | 22 +- .../load_balancing/rack_middleware_spec.rb | 119 +++------- .../database/load_balancing/setup_spec.rb | 119 ++++++++++ .../sidekiq_client_middleware_spec.rb | 27 ++- .../sidekiq_server_middleware_spec.rb | 73 +++++- .../database/load_balancing/sticking_spec.rb | 215 +++++++++++------- .../gitlab/database/load_balancing_spec.rb | 177 +++----------- spec/lib/gitlab/database_spec.rb | 2 +- .../json/streaming_serializer_spec.rb | 2 +- .../metrics/subscribers/active_record_spec.rb | 2 +- .../sidekiq_logging/structured_logger_spec.rb | 2 +- spec/models/ci/build_spec.rb | 4 +- spec/models/ci/pipeline_spec.rb | 11 +- spec/models/ci/runner_spec.rb | 2 +- spec/models/project_feature_usage_spec.rb | 4 +- spec/models/project_spec.rb | 2 +- .../merge_requests/set_assignees_spec.rb | 2 +- .../services/ci/drop_pipeline_service_spec.rb | 3 +- spec/services/ci/register_job_service_spec.rb | 6 +- ...ser_project_access_changed_service_spec.rb | 2 +- spec/services/users/activity_service_spec.rb | 4 +- spec/support/database_load_balancing.rb | 30 ++- .../lib/gitlab/ci/ci_trace_shared_examples.rb | 16 +- .../user_refresh_from_replica_worker_spec.rb | 2 +- ...container_expiration_policy_worker_spec.rb | 2 +- 54 files changed, 750 insertions(+), 643 deletions(-) delete mode 100644 lib/gitlab/database/load_balancing/active_record_proxy.rb create mode 100644 lib/gitlab/database/load_balancing/setup.rb delete mode 100644 spec/lib/gitlab/database/load_balancing/active_record_proxy_spec.rb create mode 100644 spec/lib/gitlab/database/load_balancing/setup_spec.rb diff --git a/app/models/ci/build.rb b/app/models/ci/build.rb index 5709bc26956843..1e0ce5638bb8a8 100644 --- a/app/models/ci/build.rb +++ b/app/models/ci/build.rb @@ -1109,7 +1109,7 @@ def stick_build_if_status_changed return unless saved_change_to_status? return unless running? - ::Gitlab::Database::LoadBalancing::Sticking.stick(:build, id) + self.class.sticking.stick(:build, id) end def status_commit_hooks diff --git a/app/models/ci/ci_database_record.rb b/app/models/ci/ci_database_record.rb index e76b7a05b93a1c..e2b832a28e7742 100644 --- a/app/models/ci/ci_database_record.rb +++ b/app/models/ci/ci_database_record.rb @@ -12,13 +12,6 @@ class CiDatabaseRecord < Ci::ApplicationRecord if Gitlab::Database.has_config?(:ci) connects_to database: { writing: :ci, reading: :ci } - - # TODO: Load Balancing messes with `CiDatabaseRecord` - # returning wrong connection. To be removed once merged: - # https://gitlab.com/gitlab-org/gitlab/-/merge_requests/67773 - def self.connection - retrieve_connection - end end end end diff --git a/app/models/ci/runner.rb b/app/models/ci/runner.rb index 4c5e4837d0687a..dd8cdbcaf4ba6d 100644 --- a/app/models/ci/runner.rb +++ b/app/models/ci/runner.rb @@ -348,7 +348,7 @@ def tick_runner_queue # intention here is not to execute `Ci::RegisterJobService#execute` on # the primary database. # - ::Gitlab::Database::LoadBalancing::Sticking.stick(:runner, id) + ::Ci::Runner.sticking.stick(:runner, id) SecureRandom.hex.tap do |new_update| ::Gitlab::Workhorse.set_key_and_notify(runner_queue_key, new_update, diff --git a/app/models/project.rb b/app/models/project.rb index db257787a40b6a..6841fcddc0bf57 100644 --- a/app/models/project.rb +++ b/app/models/project.rb @@ -2400,7 +2400,7 @@ def licensed_features end def mark_primary_write_location - ::Gitlab::Database::LoadBalancing::Sticking.mark_primary_write_location(:project, self.id) + self.class.sticking.mark_primary_write_location(:project, self.id) end def toggle_ci_cd_settings!(settings_attribute) diff --git a/app/services/ci/register_job_service.rb b/app/services/ci/register_job_service.rb index 5dd586e9f22c4e..67ef4f1070956b 100644 --- a/app/services/ci/register_job_service.rb +++ b/app/services/ci/register_job_service.rb @@ -22,7 +22,8 @@ def initialize(runner) end def execute(params = {}) - db_all_caught_up = ::Gitlab::Database::LoadBalancing::Sticking.all_caught_up?(:runner, runner.id) + db_all_caught_up = + ::Ci::Runner.sticking.all_caught_up?(:runner, runner.id) @metrics.increment_queue_operation(:queue_attempt) diff --git a/app/services/user_project_access_changed_service.rb b/app/services/user_project_access_changed_service.rb index 5f48f410bf77ea..5bba986f4ad87e 100644 --- a/app/services/user_project_access_changed_service.rb +++ b/app/services/user_project_access_changed_service.rb @@ -30,7 +30,7 @@ def execute(blocking: true, priority: HIGH_PRIORITY) end end - ::Gitlab::Database::LoadBalancing::Sticking.bulk_stick(:user, @user_ids) + ::User.sticking.bulk_stick(:user, @user_ids) result end diff --git a/config/initializers/load_balancing.rb b/config/initializers/load_balancing.rb index 29700b37c7d9a0..a31b11bb2bef81 100644 --- a/config/initializers/load_balancing.rb +++ b/config/initializers/load_balancing.rb @@ -1,52 +1,33 @@ # frozen_string_literal: true -ActiveRecord::Base.singleton_class.attr_accessor :load_balancing_proxy - -Gitlab::Database.main.disable_prepared_statements - 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? - ) - - ActiveRecord::Base.load_balancing_proxy = - Gitlab::Database::LoadBalancing::ConnectionProxy.new(lb) - - # Populate service discovery immediately if it is configured - Gitlab::Database::LoadBalancing.perform_service_discovery -end - -setup.call - -# 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 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 +Gitlab::Database::LoadBalancing.base_models.each do |model| + # 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. + Gitlab::Database::LoadBalancing::Setup.new(model).setup + + # 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. + model.connection.load_balancer.disconnect!(timeout: 0) + 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 + Gitlab::Database::LoadBalancing::Setup + .new(model, start_service_discovery: true) + .setup + end end diff --git a/lib/api/ci/helpers/runner.rb b/lib/api/ci/helpers/runner.rb index b9662b822fb617..dabb6c7ab3a6c8 100644 --- a/lib/api/ci/helpers/runner.rb +++ b/lib/api/ci/helpers/runner.rb @@ -42,8 +42,7 @@ def current_runner token = params[:token] if token - ::Gitlab::Database::LoadBalancing::RackMiddleware - .stick_or_unstick(env, :runner, token) + ::Ci::Runner.sticking.stick_or_unstick_request(env, :runner, token) end strong_memoize(:current_runner) do @@ -80,8 +79,9 @@ def current_job id = params[:id] if id - ::Gitlab::Database::LoadBalancing::RackMiddleware - .stick_or_unstick(env, :build, id) + ::Ci::Build + .sticking + .stick_or_unstick_request(env, :build, id) end strong_memoize(:current_job) do diff --git a/lib/api/helpers.rb b/lib/api/helpers.rb index c189fde503a800..ff3590d6c13956 100644 --- a/lib/api/helpers.rb +++ b/lib/api/helpers.rb @@ -75,8 +75,9 @@ def current_user save_current_user_in_env(@current_user) if @current_user if @current_user - ::Gitlab::Database::LoadBalancing::RackMiddleware - .stick_or_unstick(env, :user, @current_user.id) + ::ApplicationRecord + .sticking + .stick_or_unstick_request(env, :user, @current_user.id) end @current_user diff --git a/lib/gitlab/checks/matching_merge_request.rb b/lib/gitlab/checks/matching_merge_request.rb index c86779c7a37286..e5ce862264f414 100644 --- a/lib/gitlab/checks/matching_merge_request.rb +++ b/lib/gitlab/checks/matching_merge_request.rb @@ -27,7 +27,7 @@ def match? # 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) + ::ApplicationRecord.sticking.select_valid_host(:project, @project.id) end # rubocop: disable CodeReuse/ActiveRecord diff --git a/lib/gitlab/ci/trace.rb b/lib/gitlab/ci/trace.rb index 6816d216015ed5..c8f19d65c9eafd 100644 --- a/lib/gitlab/ci/trace.rb +++ b/lib/gitlab/ci/trace.rb @@ -286,7 +286,8 @@ def trace_artifact def destroy_stream(build) if consistent_archived_trace?(build) - ::Gitlab::Database::LoadBalancing::Sticking + ::Ci::Build + .sticking .stick(LOAD_BALANCING_STICKING_NAMESPACE, build.id) end @@ -295,7 +296,8 @@ def destroy_stream(build) def read_trace_artifact(build) if consistent_archived_trace?(build) - ::Gitlab::Database::LoadBalancing::Sticking + ::Ci::Build + .sticking .unstick_or_continue_sticking(LOAD_BALANCING_STICKING_NAMESPACE, build.id) end diff --git a/lib/gitlab/database.rb b/lib/gitlab/database.rb index d56f2a4b0e3a41..b560d4cbca84d7 100644 --- a/lib/gitlab/database.rb +++ b/lib/gitlab/database.rb @@ -53,7 +53,12 @@ module Database def self.database_base_models @database_base_models ||= { - main: ::ApplicationRecord, + # Note that we use ActiveRecord::Base here and not ApplicationRecord. + # This is deliberate, as we also use these classes to apply load + # balancing to, and the load balancer must be enabled for _all_ models + # that inher from ActiveRecord::Base; not just our own models that + # inherit from ApplicationRecord. + main: ::ActiveRecord::Base, ci: ::Ci::CiDatabaseRecord.connection_class? ? ::Ci::CiDatabaseRecord : nil }.compact.freeze end diff --git a/lib/gitlab/database/load_balancing.rb b/lib/gitlab/database/load_balancing.rb index a5cfaf795f9eb0..3e322e752b7a03 100644 --- a/lib/gitlab/database/load_balancing.rb +++ b/lib/gitlab/database/load_balancing.rb @@ -18,44 +18,20 @@ module LoadBalancing ActiveRecord::ConnectionNotEstablished ].freeze - def self.proxy - ActiveRecord::Base.load_balancing_proxy + def self.base_models + @base_models ||= ::Gitlab::Database.database_base_models.values.freeze end - # Returns a Hash containing the load balancing configuration. - def self.configuration - @configuration ||= Configuration.for_model(ActiveRecord::Base) - end - - # 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? + def self.each_load_balancer + return to_enum(__method__) unless block_given? - configured? - end - - def self.configured? - configuration.load_balancing_enabled? || - configuration.service_discovery_enabled? - end - - def self.start_service_discovery - return unless configuration.service_discovery_enabled? - - ServiceDiscovery - .new(proxy.load_balancer, **configuration.service_discovery) - .start + base_models.each do |model| + yield model.connection.load_balancer + end end - def self.perform_service_discovery - return unless configuration.service_discovery_enabled? - - ServiceDiscovery - .new(proxy.load_balancer, **configuration.service_discovery) - .perform_service_discovery + def self.release_hosts + each_load_balancer(&:release_host) end DB_ROLES = [ diff --git a/lib/gitlab/database/load_balancing/action_cable_callbacks.rb b/lib/gitlab/database/load_balancing/action_cable_callbacks.rb index 4feba989a0aec4..7164976ff73b25 100644 --- a/lib/gitlab/database/load_balancing/action_cable_callbacks.rb +++ b/lib/gitlab/database/load_balancing/action_cable_callbacks.rb @@ -16,7 +16,7 @@ def self.wrapper inner.call ensure - ::Gitlab::Database::LoadBalancing.proxy.load_balancer.release_host + ::Gitlab::Database::LoadBalancing.release_hosts ::Gitlab::Database::LoadBalancing::Session.clear_session end end diff --git a/lib/gitlab/database/load_balancing/active_record_proxy.rb b/lib/gitlab/database/load_balancing/active_record_proxy.rb deleted file mode 100644 index 5462c2a3188c83..00000000000000 --- a/lib/gitlab/database/load_balancing/active_record_proxy.rb +++ /dev/null @@ -1,15 +0,0 @@ -# frozen_string_literal: true - -module Gitlab - module Database - module LoadBalancing - # Module injected into ActiveRecord::Base to allow hijacking of the - # "connection" method. - module ActiveRecordProxy - def connection - ::Gitlab::Database::LoadBalancing.proxy || super - end - end - end - end -end diff --git a/lib/gitlab/database/load_balancing/configuration.rb b/lib/gitlab/database/load_balancing/configuration.rb index 238f55fd98e746..6156515bd73bed 100644 --- a/lib/gitlab/database/load_balancing/configuration.rb +++ b/lib/gitlab/database/load_balancing/configuration.rb @@ -72,7 +72,14 @@ def pool_size Database.default_pool_size end + # 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 load_balancing_enabled? + return false if Gitlab::Runtime.rake? + hosts.any? || service_discovery_enabled? end diff --git a/lib/gitlab/database/load_balancing/load_balancer.rb b/lib/gitlab/database/load_balancing/load_balancer.rb index 5d9dbfdabbdf43..cc9ca325337848 100644 --- a/lib/gitlab/database/load_balancing/load_balancer.rb +++ b/lib/gitlab/database/load_balancing/load_balancer.rb @@ -12,22 +12,22 @@ class LoadBalancer REPLICA_SUFFIX = '_replica' - attr_reader :host_list, :configuration + attr_reader :name, :host_list, :configuration # configuration - An instance of `LoadBalancing::Configuration` that # contains the configuration details (such as the hosts) # for this load balancer. - # primary_only - If set, the replicas are ignored and the primary is - # always used. - def initialize(configuration, primary_only: false) + def initialize(configuration) @configuration = configuration - @primary_only = primary_only + @primary_only = !configuration.load_balancing_enabled? @host_list = - if primary_only + if @primary_only HostList.new([PrimaryHost.new(self)]) else HostList.new(configuration.hosts.map { |addr| Host.new(addr, self) }) end + + @name = @configuration.model.connection_db_config.name.to_sym end def primary_only? diff --git a/lib/gitlab/database/load_balancing/rack_middleware.rb b/lib/gitlab/database/load_balancing/rack_middleware.rb index 22871427bff7b5..fc035f42097cd0 100644 --- a/lib/gitlab/database/load_balancing/rack_middleware.rb +++ b/lib/gitlab/database/load_balancing/rack_middleware.rb @@ -9,21 +9,6 @@ module LoadBalancing class RackMiddleware STICK_OBJECT = 'load_balancing.stick_object' - # Unsticks or continues sticking the current request. - # - # This method also updates the Rack environment so #call can later - # determine if we still need to stick or not. - # - # env - The Rack environment. - # namespace - The namespace to use for sticking. - # id - The identifier to use for sticking. - def self.stick_or_unstick(env, namespace, id) - ::Gitlab::Database::LoadBalancing::Sticking.unstick_or_continue_sticking(namespace, id) - - env[STICK_OBJECT] ||= Set.new - env[STICK_OBJECT] << [namespace, id] - end - def initialize(app) @app = app end @@ -51,41 +36,46 @@ def call(env) # Typically this code will only be reachable for Rails requests as # Grape data is not yet available at this point. def unstick_or_continue_sticking(env) - namespaces_and_ids = sticking_namespaces_and_ids(env) + namespaces_and_ids = sticking_namespaces(env) - namespaces_and_ids.each do |namespace, id| - ::Gitlab::Database::LoadBalancing::Sticking.unstick_or_continue_sticking(namespace, id) + namespaces_and_ids.each do |(model, namespace, id)| + model.sticking.unstick_or_continue_sticking(namespace, id) end end # Determine if we need to stick after handling a request. def stick_if_necessary(env) - namespaces_and_ids = sticking_namespaces_and_ids(env) + namespaces_and_ids = sticking_namespaces(env) - namespaces_and_ids.each do |namespace, id| - ::Gitlab::Database::LoadBalancing::Sticking.stick_if_necessary(namespace, id) + namespaces_and_ids.each do |model, namespace, id| + model.sticking.stick_if_necessary(namespace, id) end end def clear - load_balancer.release_host + ::Gitlab::Database::LoadBalancing.release_hosts ::Gitlab::Database::LoadBalancing::Session.clear_session end - def load_balancer - ::Gitlab::Database::LoadBalancing.proxy.load_balancer - end - # Determines the sticking namespace and identifier based on the Rack # environment. # # For Rails requests this uses warden, but Grape and others have to # manually set the right environment variable. - def sticking_namespaces_and_ids(env) + def sticking_namespaces(env) warden = env['warden'] if warden && warden.user - [[:user, warden.user.id]] + # When sticking per user, _only_ sticking the main connection could + # result in the application trying to read data from a different + # connection, while that data isn't available yet. + # + # To prevent this from happening, we scope sticking to all the + # models that support load balancing. In the future (if we + # determined this to be OK) we may be able to relax this. + LoadBalancing.base_models.map do |model| + [model, :user, warden.user.id] + end elsif env[STICK_OBJECT].present? env[STICK_OBJECT].to_a else diff --git a/lib/gitlab/database/load_balancing/setup.rb b/lib/gitlab/database/load_balancing/setup.rb new file mode 100644 index 00000000000000..3cce839a960f73 --- /dev/null +++ b/lib/gitlab/database/load_balancing/setup.rb @@ -0,0 +1,61 @@ +# frozen_string_literal: true + +module Gitlab + module Database + module LoadBalancing + # Class for setting up load balancing of a specific model. + class Setup + attr_reader :configuration + + def initialize(model, start_service_discovery: false) + @model = model + @configuration = Configuration.for_model(model) + @start_service_discovery = start_service_discovery + end + + def setup + disable_prepared_statements + setup_load_balancer + setup_service_discovery + end + + def disable_prepared_statements + db_config_object = @model.connection_db_config + config = + db_config_object.configuration_hash.merge(prepared_statements: false) + + hash_config = ActiveRecord::DatabaseConfigurations::HashConfig.new( + db_config_object.env_name, + db_config_object.name, + config + ) + + @model.establish_connection(hash_config) + end + + def setup_load_balancer + lb = LoadBalancer.new(configuration) + + # 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) + + @model.connection = ConnectionProxy.new(lb) + @model.sticking = Sticking.new(lb) + end + + def setup_service_discovery + return unless configuration.service_discovery_enabled? + + lb = @model.connection.load_balancer + sv = ServiceDiscovery.new(lb, **configuration.service_discovery) + + sv.perform_service_discovery + + sv.start if @start_service_discovery + end + end + end + end +end diff --git a/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb b/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb index adda6c1f6e7d3e..ebf90ebf9d0ce0 100644 --- a/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb +++ b/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb @@ -30,26 +30,23 @@ def load_balancing_enabled?(worker_class) end def set_data_consistency_locations!(job) - # Once we add support for multiple databases to our load balancer, we would use something like this: - # job['wal_locations'] = Gitlab::Database.databases.transform_values do |connection| - # connection.load_balancer.primary_write_location - # end - # - job['wal_locations'] = { ::Gitlab::Database::MAIN_DATABASE_NAME.to_sym => wal_location } if wal_location - end + locations = {} - def wal_location - strong_memoize(:wal_location) do - if ::Gitlab::Database::LoadBalancing::Session.current.use_primary? - load_balancer.primary_write_location - else - load_balancer.host.database_replica_location + ::Gitlab::Database::LoadBalancing.each_load_balancer do |lb| + if (location = wal_location_for(lb)) + locations[lb.name] = location end end + + job['wal_locations'] = locations end - def load_balancer - ::Gitlab::Database::LoadBalancing.proxy.load_balancer + def wal_location_for(load_balancer) + if ::Gitlab::Database::LoadBalancing::Session.current.use_primary? + load_balancer.primary_write_location + else + load_balancer.host.database_replica_location + end end end end diff --git a/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb b/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb index 5e391c7228da7b..f0c7016032bb78 100644 --- a/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb +++ b/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb @@ -29,7 +29,7 @@ def call(worker, job, _queue) private def clear - release_hosts + LoadBalancing.release_hosts Session.clear_session end @@ -44,7 +44,7 @@ def select_load_balancing_strategy(worker_class, job) return :primary_no_wal unless wal_locations - if all_databases_has_replica_caught_up?(wal_locations) + if databases_in_sync?(wal_locations) # Happy case: we can read from a replica. retried_before?(worker_class, job) ? :replica_retried : :replica elsif can_retry?(worker_class, job) @@ -89,27 +89,18 @@ def not_yet_retried?(job) job['retry_count'].nil? end - def all_databases_has_replica_caught_up?(wal_locations) - wal_locations.all? do |_config_name, location| - # Once we add support for multiple databases to our load balancer, we would use something like this: - # Gitlab::Database.databases[config_name].load_balancer.select_up_to_date_host(location) - load_balancer.select_up_to_date_host(location) + def databases_in_sync?(wal_locations) + LoadBalancing.each_load_balancer.all? do |lb| + if (location = wal_locations[lb.name]) + lb.select_up_to_date_host(location) + else + # If there's no entry for a load balancer it means the Sidekiq + # job doesn't care for it. In this case we'll treat the load + # balancer as being in sync. + true + end end end - - def release_hosts - # Once we add support for multiple databases to our load balancer, we would use something like this: - # connection.load_balancer.primary_write_location - # - # Gitlab::Database.databases.values.each do |connection| - # connection.load_balancer.release_host - # end - load_balancer.release_host - end - - def load_balancer - LoadBalancing.proxy.load_balancer - end end end end diff --git a/lib/gitlab/database/load_balancing/sticking.rb b/lib/gitlab/database/load_balancing/sticking.rb index 8d55ade488026a..aa5808e61b0127 100644 --- a/lib/gitlab/database/load_balancing/sticking.rb +++ b/lib/gitlab/database/load_balancing/sticking.rb @@ -5,34 +5,47 @@ module Database module LoadBalancing # Module used for handling sticking connections to a primary, if # necessary. - # - # ## Examples - # - # Sticking a user to the primary: - # - # Sticking.stick_if_necessary(:user, current_user.id) - # - # To unstick if possible, or continue using the primary otherwise: - # - # Sticking.unstick_or_continue_sticking(:user, current_user.id) - module Sticking + class Sticking # The number of seconds after which a session should stop reading from # the primary. EXPIRATION = 30 + def initialize(load_balancer) + @load_balancer = load_balancer + @model = load_balancer.configuration.model + end + + # Unsticks or continues sticking the current request. + # + # This method also updates the Rack environment so #call can later + # determine if we still need to stick or not. + # + # env - The Rack environment. + # namespace - The namespace to use for sticking. + # id - The identifier to use for sticking. + # model - The ActiveRecord model to scope sticking to. + def stick_or_unstick_request(env, namespace, id) + unstick_or_continue_sticking(namespace, id) + + env[RackMiddleware::STICK_OBJECT] ||= Set.new + env[RackMiddleware::STICK_OBJECT] << [@model, namespace, id] + end + # Sticks to the primary if a write was performed. - def self.stick_if_necessary(namespace, id) + def stick_if_necessary(namespace, id) stick(namespace, id) if Session.current.performed_write? end - # Checks if we are caught-up with all the work - def self.all_caught_up?(namespace, id) + def all_caught_up?(namespace, id) location = last_write_location_for(namespace, id) return true unless location - load_balancer.select_up_to_date_host(location).tap do |found| - ActiveSupport::Notifications.instrument('caught_up_replica_pick.load_balancing', { result: found } ) + @load_balancer.select_up_to_date_host(location).tap do |found| + ActiveSupport::Notifications.instrument( + 'caught_up_replica_pick.load_balancing', + { result: found } + ) unstick(namespace, id) if found end @@ -43,7 +56,7 @@ def self.all_caught_up?(namespace, id) # in another thread. # # Returns true if one host was selected. - def self.select_caught_up_replicas(namespace, id) + def select_caught_up_replicas(namespace, id) location = last_write_location_for(namespace, id) # Unlike all_caught_up?, we return false if no write location exists. @@ -51,33 +64,36 @@ def self.select_caught_up_replicas(namespace, id) # write location. If no such location exists, err on the side of caution. return false unless location - load_balancer.select_up_to_date_host(location).tap do |selected| + @load_balancer.select_up_to_date_host(location).tap do |selected| unstick(namespace, id) if selected end end # Sticks to the primary if necessary, otherwise unsticks an object (if # it was previously stuck to the primary). - def self.unstick_or_continue_sticking(namespace, id) - Session.current.use_primary! unless all_caught_up?(namespace, id) + def unstick_or_continue_sticking(namespace, id) + return if all_caught_up?(namespace, id) + + Session.current.use_primary! end # Select a replica that has caught up with the primary. If one has not been # found, stick to the primary. - def self.select_valid_host(namespace, id) - replica_selected = select_caught_up_replicas(namespace, id) + def select_valid_host(namespace, id) + replica_selected = + select_caught_up_replicas(namespace, id) Session.current.use_primary! unless replica_selected end # Starts sticking to the primary for the given namespace and id, using # the latest WAL pointer from the primary. - def self.stick(namespace, id) + def stick(namespace, id) mark_primary_write_location(namespace, id) Session.current.use_primary! end - def self.bulk_stick(namespace, ids) + def bulk_stick(namespace, ids) with_primary_write_location do |location| ids.each do |id| set_write_location_for(namespace, id, location) @@ -87,45 +103,49 @@ def self.bulk_stick(namespace, ids) Session.current.use_primary! end - def self.with_primary_write_location - location = load_balancer.primary_write_location + def with_primary_write_location + location = @load_balancer.primary_write_location return if location.blank? yield(location) end - def self.mark_primary_write_location(namespace, id) + def mark_primary_write_location(namespace, id) with_primary_write_location do |location| set_write_location_for(namespace, id, location) end end - # Stops sticking to the primary. - def self.unstick(namespace, id) + def unstick(namespace, id) Gitlab::Redis::SharedState.with do |redis| redis.del(redis_key_for(namespace, id)) + redis.del(old_redis_key_for(namespace, id)) end end - def self.set_write_location_for(namespace, id, location) + def set_write_location_for(namespace, id, location) Gitlab::Redis::SharedState.with do |redis| redis.set(redis_key_for(namespace, id), location, ex: EXPIRATION) + redis.set(old_redis_key_for(namespace, id), location, ex: EXPIRATION) end end - def self.last_write_location_for(namespace, id) + def last_write_location_for(namespace, id) Gitlab::Redis::SharedState.with do |redis| - redis.get(redis_key_for(namespace, id)) + redis.get(redis_key_for(namespace, id)) || + redis.get(old_redis_key_for(namespace, id)) end end - def self.redis_key_for(namespace, id) - "database-load-balancing/write-location/#{namespace}/#{id}" + def redis_key_for(namespace, id) + name = @load_balancer.name + + "database-load-balancing/write-location/#{name}/#{namespace}/#{id}" end - def self.load_balancer - LoadBalancing.proxy.load_balancer + def old_redis_key_for(namespace, id) + "database-load-balancing/write-location/#{namespace}/#{id}" end end end diff --git a/spec/lib/api/ci/helpers/runner_spec.rb b/spec/lib/api/ci/helpers/runner_spec.rb index 99f2db544a5067..cc871d66d4008e 100644 --- a/spec/lib/api/ci/helpers/runner_spec.rb +++ b/spec/lib/api/ci/helpers/runner_spec.rb @@ -15,8 +15,8 @@ it 'handles sticking of a build when a build ID is specified' do allow(helper).to receive(:params).and_return(id: build.id) - expect(Gitlab::Database::LoadBalancing::RackMiddleware) - .to receive(:stick_or_unstick) + expect(ApplicationRecord.sticking) + .to receive(:stick_or_unstick_request) .with({}, :build, build.id) helper.current_job @@ -25,8 +25,8 @@ it 'does not handle sticking if no build ID was specified' do allow(helper).to receive(:params).and_return({}) - expect(Gitlab::Database::LoadBalancing::RackMiddleware) - .not_to receive(:stick_or_unstick) + expect(ApplicationRecord.sticking) + .not_to receive(:stick_or_unstick_request) helper.current_job end @@ -44,8 +44,8 @@ it 'handles sticking of a runner if a token is specified' do allow(helper).to receive(:params).and_return(token: runner.token) - expect(Gitlab::Database::LoadBalancing::RackMiddleware) - .to receive(:stick_or_unstick) + expect(ApplicationRecord.sticking) + .to receive(:stick_or_unstick_request) .with({}, :runner, runner.token) helper.current_runner @@ -54,8 +54,8 @@ it 'does not handle sticking if no token was specified' do allow(helper).to receive(:params).and_return({}) - expect(Gitlab::Database::LoadBalancing::RackMiddleware) - .not_to receive(:stick_or_unstick) + expect(ApplicationRecord.sticking) + .not_to receive(:stick_or_unstick_request) helper.current_runner end diff --git a/spec/lib/api/helpers_spec.rb b/spec/lib/api/helpers_spec.rb index c966527ca5b391..37e040a422bee3 100644 --- a/spec/lib/api/helpers_spec.rb +++ b/spec/lib/api/helpers_spec.rb @@ -35,8 +35,8 @@ def app it 'handles sticking when a user could be found' do allow_any_instance_of(API::Helpers).to receive(:initial_current_user).and_return(user) - expect(Gitlab::Database::LoadBalancing::RackMiddleware) - .to receive(:stick_or_unstick).with(any_args, :user, 42) + expect(ApplicationRecord.sticking) + .to receive(:stick_or_unstick_request).with(any_args, :user, 42) get 'user' @@ -46,8 +46,8 @@ def app it 'does not handle sticking if no user could be found' do allow_any_instance_of(API::Helpers).to receive(:initial_current_user).and_return(nil) - expect(Gitlab::Database::LoadBalancing::RackMiddleware) - .not_to receive(:stick_or_unstick) + expect(ApplicationRecord.sticking) + .not_to receive(:stick_or_unstick_request) get 'user' diff --git a/spec/lib/gitlab/checks/matching_merge_request_spec.rb b/spec/lib/gitlab/checks/matching_merge_request_spec.rb index dce6a27191894d..c65a1e4d6566e2 100644 --- a/spec/lib/gitlab/checks/matching_merge_request_spec.rb +++ b/spec/lib/gitlab/checks/matching_merge_request_spec.rb @@ -37,10 +37,20 @@ 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) + allow(::ApplicationRecord.sticking) + .to receive(:all_caught_up?) + .and_return(all_caught_up) + + expect(::ApplicationRecord.sticking) + .to receive(:select_valid_host) + .with(:project, project.id) + .and_call_original + + allow(::ApplicationRecord.sticking) + .to receive(:select_caught_up_replicas) + .with(:project, project.id) + .and_return(all_caught_up) end after do diff --git a/spec/lib/gitlab/database/consistency_spec.rb b/spec/lib/gitlab/database/consistency_spec.rb index b3d166d58bbf02..5055be81c880bf 100644 --- a/spec/lib/gitlab/database/consistency_spec.rb +++ b/spec/lib/gitlab/database/consistency_spec.rb @@ -7,7 +7,15 @@ Gitlab::Database::LoadBalancing::Session.current end - describe '.with_read_consistency', :db_load_balancing do + before do + Gitlab::Database::LoadBalancing::Session.clear_session + end + + after do + Gitlab::Database::LoadBalancing::Session.clear_session + end + + describe '.with_read_consistency' do it 'sticks to primary database' do expect(session).not_to be_using_primary diff --git a/spec/lib/gitlab/database/load_balancing/action_cable_callbacks_spec.rb b/spec/lib/gitlab/database/load_balancing/action_cable_callbacks_spec.rb index ebbbafb855ffa1..768855464c109b 100644 --- a/spec/lib/gitlab/database/load_balancing/action_cable_callbacks_spec.rb +++ b/spec/lib/gitlab/database/load_balancing/action_cable_callbacks_spec.rb @@ -5,7 +5,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ActionCableCallbacks, :request_store do describe '.wrapper' do it 'uses primary and then releases the connection and clears the session' do - expect(Gitlab::Database::LoadBalancing).to receive_message_chain(:proxy, :load_balancer, :release_host) + expect(Gitlab::Database::LoadBalancing).to receive(:release_hosts) expect(Gitlab::Database::LoadBalancing::Session).to receive(:clear_session) described_class.wrapper.call( @@ -18,7 +18,7 @@ context 'with an exception' do it 'releases the connection and clears the session' do - expect(Gitlab::Database::LoadBalancing).to receive_message_chain(:proxy, :load_balancer, :release_host) + expect(Gitlab::Database::LoadBalancing).to receive(:release_hosts) expect(Gitlab::Database::LoadBalancing::Session).to receive(:clear_session) expect do diff --git a/spec/lib/gitlab/database/load_balancing/active_record_proxy_spec.rb b/spec/lib/gitlab/database/load_balancing/active_record_proxy_spec.rb deleted file mode 100644 index 2ccb52f7b145ec..00000000000000 --- a/spec/lib/gitlab/database/load_balancing/active_record_proxy_spec.rb +++ /dev/null @@ -1,27 +0,0 @@ -# frozen_string_literal: true - -require 'spec_helper' - -RSpec.describe Gitlab::Database::LoadBalancing::ActiveRecordProxy do - describe '#connection' do - it 'returns a connection proxy' do - dummy = Class.new do - include Gitlab::Database::LoadBalancing::ActiveRecordProxy - end - - proxy = double(:proxy) - - expect(Gitlab::Database::LoadBalancing).to receive(:proxy) - .and_return(proxy) - - expect(dummy.new.connection).to eq(proxy) - end - - it 'returns a connection when no proxy is present' do - allow(Gitlab::Database::LoadBalancing).to receive(:proxy).and_return(nil) - - expect(ActiveRecord::Base.connection) - .to eq(ActiveRecord::Base.retrieve_connection) - end - end -end diff --git a/spec/lib/gitlab/database/load_balancing/configuration_spec.rb b/spec/lib/gitlab/database/load_balancing/configuration_spec.rb index 6621e6276a5309..3e5249a3deaa58 100644 --- a/spec/lib/gitlab/database/load_balancing/configuration_spec.rb +++ b/spec/lib/gitlab/database/load_balancing/configuration_spec.rb @@ -108,6 +108,14 @@ end describe '#load_balancing_enabled?' do + it 'returns false when running inside a Rake task' do + config = described_class.new(ActiveRecord::Base, %w[foo bar]) + + allow(Gitlab::Runtime).to receive(:rake?).and_return(true) + + expect(config.load_balancing_enabled?).to eq(false) + end + it 'returns true when hosts are configured' do config = described_class.new(ActiveRecord::Base, %w[foo bar]) diff --git a/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb b/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb index fb17ba0e121796..f3ce5563e3883f 100644 --- a/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb +++ b/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb @@ -47,16 +47,27 @@ def twice_wrapped_exception(top, middle, original) end describe '#initialize' do - it 'ignores the hosts when the primary_only option is enabled' do + it 'ignores the hosts when load balancing is disabled' do config = Gitlab::Database::LoadBalancing::Configuration .new(ActiveRecord::Base, [db_host]) - lb = described_class.new(config, primary_only: true) + + allow(config).to receive(:load_balancing_enabled?).and_return(false) + + lb = described_class.new(config) hosts = lb.host_list.hosts expect(hosts.length).to eq(1) expect(hosts.first) .to be_instance_of(Gitlab::Database::LoadBalancing::PrimaryHost) end + + it 'sets the name of the connection that is used' do + config = + Gitlab::Database::LoadBalancing::Configuration.new(ActiveRecord::Base) + lb = described_class.new(config) + + expect(lb.name).to eq(:main) + end end describe '#read' do @@ -140,10 +151,13 @@ def twice_wrapped_exception(top, middle, original) .to yield_with_args(ActiveRecord::Base.retrieve_connection) end - it 'uses the primary when the primary_only option is enabled' do + it 'uses the primary when load balancing is disabled' do config = Gitlab::Database::LoadBalancing::Configuration .new(ActiveRecord::Base) - lb = described_class.new(config, primary_only: true) + + allow(config).to receive(:load_balancing_enabled?).and_return(false) + + lb = described_class.new(config) # When no hosts are configured, we don't want to produce any warnings, as # they aren't useful/too noisy. diff --git a/spec/lib/gitlab/database/load_balancing/rack_middleware_spec.rb b/spec/lib/gitlab/database/load_balancing/rack_middleware_spec.rb index 5453e36e08fe68..af7e2a4b167c6d 100644 --- a/spec/lib/gitlab/database/load_balancing/rack_middleware_spec.rb +++ b/spec/lib/gitlab/database/load_balancing/rack_middleware_spec.rb @@ -6,12 +6,12 @@ let(:app) { double(:app) } let(:middleware) { described_class.new(app) } let(:warden_user) { double(:warden, user: double(:user, id: 42)) } - let(:single_sticking_object) { Set.new([[:user, 42]]) } + let(:single_sticking_object) { Set.new([[ActiveRecord::Base, :user, 42]]) } let(:multiple_sticking_objects) do Set.new([ - [:user, 42], - [:runner, '123456789'], - [:runner, '1234'] + [ActiveRecord::Base, :user, 42], + [ActiveRecord::Base, :runner, '123456789'], + [ActiveRecord::Base, :runner, '1234'] ]) end @@ -19,42 +19,6 @@ Gitlab::Database::LoadBalancing::Session.clear_session end - describe '.stick_or_unstick' do - it 'sticks or unsticks a single object and updates the Rack environment' do - expect(Gitlab::Database::LoadBalancing::Sticking) - .to receive(:unstick_or_continue_sticking) - .with(:user, 42) - - env = {} - - described_class.stick_or_unstick(env, :user, 42) - - expect(env[described_class::STICK_OBJECT].to_a).to eq([[:user, 42]]) - end - - it 'sticks or unsticks multiple objects and updates the Rack environment' do - expect(Gitlab::Database::LoadBalancing::Sticking) - .to receive(:unstick_or_continue_sticking) - .with(:user, 42) - .ordered - - expect(Gitlab::Database::LoadBalancing::Sticking) - .to receive(:unstick_or_continue_sticking) - .with(:runner, '123456789') - .ordered - - env = {} - - described_class.stick_or_unstick(env, :user, 42) - described_class.stick_or_unstick(env, :runner, '123456789') - - expect(env[described_class::STICK_OBJECT].to_a).to eq([ - [:user, 42], - [:runner, '123456789'] - ]) - end - end - describe '#call' do it 'handles a request' do env = {} @@ -77,7 +41,7 @@ describe '#unstick_or_continue_sticking' do it 'does not stick if no namespace and identifier could be found' do - expect(Gitlab::Database::LoadBalancing::Sticking) + expect(ApplicationRecord.sticking) .not_to receive(:unstick_or_continue_sticking) middleware.unstick_or_continue_sticking({}) @@ -86,9 +50,11 @@ it 'sticks to the primary if a warden user is found' do env = { 'warden' => warden_user } - expect(Gitlab::Database::LoadBalancing::Sticking) - .to receive(:unstick_or_continue_sticking) - .with(:user, 42) + Gitlab::Database::LoadBalancing.base_models.each do |model| + expect(model.sticking) + .to receive(:unstick_or_continue_sticking) + .with(:user, 42) + end middleware.unstick_or_continue_sticking(env) end @@ -96,7 +62,7 @@ it 'sticks to the primary if a sticking namespace and identifier is found' do env = { described_class::STICK_OBJECT => single_sticking_object } - expect(Gitlab::Database::LoadBalancing::Sticking) + expect(ApplicationRecord.sticking) .to receive(:unstick_or_continue_sticking) .with(:user, 42) @@ -106,17 +72,17 @@ it 'sticks to the primary if multiple sticking namespaces and identifiers were found' do env = { described_class::STICK_OBJECT => multiple_sticking_objects } - expect(Gitlab::Database::LoadBalancing::Sticking) + expect(ApplicationRecord.sticking) .to receive(:unstick_or_continue_sticking) .with(:user, 42) .ordered - expect(Gitlab::Database::LoadBalancing::Sticking) + expect(ApplicationRecord.sticking) .to receive(:unstick_or_continue_sticking) .with(:runner, '123456789') .ordered - expect(Gitlab::Database::LoadBalancing::Sticking) + expect(ApplicationRecord.sticking) .to receive(:unstick_or_continue_sticking) .with(:runner, '1234') .ordered @@ -127,7 +93,7 @@ describe '#stick_if_necessary' do it 'does not stick to the primary if not necessary' do - expect(Gitlab::Database::LoadBalancing::Sticking) + expect(ApplicationRecord.sticking) .not_to receive(:stick_if_necessary) middleware.stick_if_necessary({}) @@ -136,9 +102,11 @@ it 'sticks to the primary if a warden user is found' do env = { 'warden' => warden_user } - expect(Gitlab::Database::LoadBalancing::Sticking) - .to receive(:stick_if_necessary) - .with(:user, 42) + Gitlab::Database::LoadBalancing.base_models.each do |model| + expect(model.sticking) + .to receive(:stick_if_necessary) + .with(:user, 42) + end middleware.stick_if_necessary(env) end @@ -146,7 +114,7 @@ it 'sticks to the primary if a a single sticking object is found' do env = { described_class::STICK_OBJECT => single_sticking_object } - expect(Gitlab::Database::LoadBalancing::Sticking) + expect(ApplicationRecord.sticking) .to receive(:stick_if_necessary) .with(:user, 42) @@ -156,17 +124,17 @@ it 'sticks to the primary if multiple sticking namespaces and identifiers were found' do env = { described_class::STICK_OBJECT => multiple_sticking_objects } - expect(Gitlab::Database::LoadBalancing::Sticking) + expect(ApplicationRecord.sticking) .to receive(:stick_if_necessary) .with(:user, 42) .ordered - expect(Gitlab::Database::LoadBalancing::Sticking) + expect(ApplicationRecord.sticking) .to receive(:stick_if_necessary) .with(:runner, '123456789') .ordered - expect(Gitlab::Database::LoadBalancing::Sticking) + expect(ApplicationRecord.sticking) .to receive(:stick_if_necessary) .with(:runner, '1234') .ordered @@ -177,47 +145,34 @@ describe '#clear' do it 'clears the currently used host and session' do - lb = double(:lb) session = spy(:session) - allow(middleware).to receive(:load_balancer).and_return(lb) - - expect(lb).to receive(:release_host) - stub_const('Gitlab::Database::LoadBalancing::Session', session) + expect(Gitlab::Database::LoadBalancing).to receive(:release_hosts) + middleware.clear expect(session).to have_received(:clear_session) end end - describe '.load_balancer' do - it 'returns a the load balancer' do - proxy = double(:proxy) - - expect(Gitlab::Database::LoadBalancing).to receive(:proxy) - .and_return(proxy) - - expect(proxy).to receive(:load_balancer) - - middleware.load_balancer - end - end - - describe '#sticking_namespaces_and_ids' do + describe '#sticking_namespaces' do context 'using a Warden request' do it 'returns the warden user if present' do env = { 'warden' => warden_user } + ids = Gitlab::Database::LoadBalancing.base_models.map do |model| + [model, :user, 42] + end - expect(middleware.sticking_namespaces_and_ids(env)).to eq([[:user, 42]]) + expect(middleware.sticking_namespaces(env)).to eq(ids) end it 'returns an empty Array if no user was present' do warden = double(:warden, user: nil) env = { 'warden' => warden } - expect(middleware.sticking_namespaces_and_ids(env)).to eq([]) + expect(middleware.sticking_namespaces(env)).to eq([]) end end @@ -225,17 +180,17 @@ it 'returns the sticking object' do env = { described_class::STICK_OBJECT => multiple_sticking_objects } - expect(middleware.sticking_namespaces_and_ids(env)).to eq([ - [:user, 42], - [:runner, '123456789'], - [:runner, '1234'] + expect(middleware.sticking_namespaces(env)).to eq([ + [ActiveRecord::Base, :user, 42], + [ActiveRecord::Base, :runner, '123456789'], + [ActiveRecord::Base, :runner, '1234'] ]) end end context 'using a regular request' do it 'returns an empty Array' do - expect(middleware.sticking_namespaces_and_ids({})).to eq([]) + expect(middleware.sticking_namespaces({})).to eq([]) end end end diff --git a/spec/lib/gitlab/database/load_balancing/setup_spec.rb b/spec/lib/gitlab/database/load_balancing/setup_spec.rb new file mode 100644 index 00000000000000..01646bc76efbbe --- /dev/null +++ b/spec/lib/gitlab/database/load_balancing/setup_spec.rb @@ -0,0 +1,119 @@ +# frozen_string_literal: true + +require 'spec_helper' + +RSpec.describe Gitlab::Database::LoadBalancing::Setup do + describe '#setup' do + it 'sets up the load balancer' do + setup = described_class.new(ActiveRecord::Base) + + expect(setup).to receive(:disable_prepared_statements) + expect(setup).to receive(:setup_load_balancer) + expect(setup).to receive(:setup_service_discovery) + + setup.setup + end + end + + describe '#disable_prepared_statements' do + it 'disables prepared statements and reconnects to the database' do + config = double( + :config, + configuration_hash: { host: 'localhost' }, + env_name: 'test', + name: 'main' + ) + model = double(:model, connection_db_config: config) + + expect(ActiveRecord::DatabaseConfigurations::HashConfig) + .to receive(:new) + .with('test', 'main', { host: 'localhost', prepared_statements: false }) + .and_call_original + + # HashConfig doesn't implement its own #==, so we can't directly compare + # the expected value with a pre-defined one. + expect(model) + .to receive(:establish_connection) + .with(an_instance_of(ActiveRecord::DatabaseConfigurations::HashConfig)) + + described_class.new(model).disable_prepared_statements + end + end + + describe '#setup_load_balancer' do + it 'sets up the load balancer' do + model = Class.new(ActiveRecord::Base) + setup = described_class.new(model) + config = Gitlab::Database::LoadBalancing::Configuration.new(model) + lb = instance_spy(Gitlab::Database::LoadBalancing::LoadBalancer) + + allow(lb).to receive(:configuration).and_return(config) + + expect(Gitlab::Database::LoadBalancing::LoadBalancer) + .to receive(:new) + .with(setup.configuration) + .and_return(lb) + + setup.setup_load_balancer + + expect(model.connection.load_balancer).to eq(lb) + expect(model.sticking) + .to be_an_instance_of(Gitlab::Database::LoadBalancing::Sticking) + end + end + + describe '#setup_service_discovery' do + context 'when service discovery is disabled' do + it 'does nothing' do + expect(Gitlab::Database::LoadBalancing::ServiceDiscovery) + .not_to receive(:new) + + described_class.new(ActiveRecord::Base).setup_service_discovery + end + end + + context 'when service discovery is enabled' do + it 'immediately performs service discovery' 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?) + .and_return(true) + + allow(Gitlab::Database::LoadBalancing::ServiceDiscovery) + .to receive(:new) + .with(lb, setup.configuration.service_discovery) + .and_return(sv) + + expect(sv).to receive(:perform_service_discovery) + expect(sv).not_to receive(:start) + + setup.setup_service_discovery + end + + it 'starts service discovery if needed' 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?) + .and_return(true) + + allow(Gitlab::Database::LoadBalancing::ServiceDiscovery) + .to receive(:new) + .with(lb, setup.configuration.service_discovery) + .and_return(sv) + + expect(sv).to receive(:perform_service_discovery) + expect(sv).to receive(:start) + + setup.setup_service_discovery + end + end + end +end diff --git a/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb b/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb index 18d9cc9e170565..1d00cba0ecadfd 100644 --- a/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb +++ b/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb @@ -5,7 +5,6 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do let(:middleware) { described_class.new } - let(:load_balancer) { Gitlab::Database::LoadBalancing.proxy.load_balancer } let(:worker_class) { 'TestDataConsistencyWorker' } let(:job) { { "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e" } } @@ -84,9 +83,15 @@ def perform(*args) end it 'passes database_replica_location' do - expected_location = { Gitlab::Database::MAIN_DATABASE_NAME.to_sym => location } + expected_location = {} - expect(load_balancer).to receive_message_chain(:host, "database_replica_location").and_return(location) + Gitlab::Database::LoadBalancing.each_load_balancer do |lb| + expect(lb.host) + .to receive(:database_replica_location) + .and_return(location) + + expected_location[lb.name] = location + end run_middleware @@ -102,9 +107,15 @@ def perform(*args) end it 'passes primary write location', :aggregate_failures do - expected_location = { Gitlab::Database::MAIN_DATABASE_NAME.to_sym => location } + expected_location = {} - expect(load_balancer).to receive(:primary_write_location).and_return(location) + Gitlab::Database::LoadBalancing.each_load_balancer do |lb| + expect(lb) + .to receive(:primary_write_location) + .and_return(location) + + expected_location[lb.name] = location + end run_middleware @@ -136,8 +147,10 @@ def perform(*args) let(:job) { { "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e", 'wal_locations' => wal_locations } } before do - allow(load_balancer).to receive(:primary_write_location).and_return(new_location) - allow(load_balancer).to receive(:database_replica_location).and_return(new_location) + Gitlab::Database::LoadBalancing.each_load_balancer do |lb| + allow(lb).to receive(:primary_write_location).and_return(new_location) + allow(lb).to receive(:database_replica_location).and_return(new_location) + end end shared_examples_for 'does not set database location again' do |use_primary| diff --git a/spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb b/spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb index 7d3670c527110a..06efdcd8f9987a 100644 --- a/spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb +++ b/spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb @@ -4,9 +4,6 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware, :clean_gitlab_redis_queues do let(:middleware) { described_class.new } - - 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' } } @@ -15,6 +12,7 @@ skip_default_enabled_yaml_check replication_lag!(false) + Gitlab::Database::LoadBalancing::Session.clear_session end after do @@ -66,7 +64,10 @@ def perform(*args) let(:wal_locations) { { Gitlab::Database::MAIN_DATABASE_NAME.to_sym => location } } it 'does not stick to the primary', :aggregate_failures do - expect(load_balancer).to receive(:select_up_to_date_host).with(location).and_return(true) + expect(ActiveRecord::Base.connection.load_balancer) + .to receive(:select_up_to_date_host) + .with(location) + .and_return(true) run_middleware do expect(Gitlab::Database::LoadBalancing::Session.current.use_primary?).not_to be_truthy @@ -91,7 +92,12 @@ def perform(*args) let(:job) { { 'job_id' => 'a180b47c-3fd6-41b8-81e9-34da61c3400e', 'wal_locations' => wal_locations } } before do - allow(load_balancer).to receive(:select_up_to_date_host).with(location).and_return(true) + Gitlab::Database::LoadBalancing.each_load_balancer do |lb| + allow(lb) + .to receive(:select_up_to_date_host) + .with(location) + .and_return(true) + end end it_behaves_like 'replica is up to date', 'replica' @@ -101,7 +107,10 @@ def perform(*args) let(:job) { { 'job_id' => 'a180b47c-3fd6-41b8-81e9-34da61c3400e', 'dedup_wal_locations' => wal_locations } } before do - allow(load_balancer).to receive(:select_up_to_date_host).with(wal_locations[:main]).and_return(true) + allow(ActiveRecord::Base.connection.load_balancer) + .to receive(:select_up_to_date_host) + .with(wal_locations[:main]) + .and_return(true) end it_behaves_like 'replica is up to date', 'replica' @@ -111,7 +120,10 @@ def perform(*args) let(:job) { { 'job_id' => 'a180b47c-3fd6-41b8-81e9-34da61c3400e', 'database_write_location' => '0/D525E3A8' } } before do - allow(load_balancer).to receive(:select_up_to_date_host).with('0/D525E3A8').and_return(true) + allow(ActiveRecord::Base.connection.load_balancer) + .to receive(:select_up_to_date_host) + .with('0/D525E3A8') + .and_return(true) end it_behaves_like 'replica is up to date', 'replica' @@ -187,7 +199,9 @@ def perform(*args) context 'when replica is not up to date' do before do - allow(load_balancer).to receive(:select_up_to_date_host).and_return(false) + Gitlab::Database::LoadBalancing.each_load_balancer do |lb| + allow(lb).to receive(:select_up_to_date_host).and_return(false) + end end include_examples 'stick to the primary', 'primary' @@ -195,6 +209,45 @@ def perform(*args) end end + describe '#databases_in_sync?' do + it 'treats load balancers without WAL entries as in sync' do + expect(middleware.send(:databases_in_sync?, {})) + .to eq(true) + end + + it 'returns true when all load balancers are in sync' do + locations = {} + + Gitlab::Database::LoadBalancing.each_load_balancer do |lb| + locations[lb.name] = 'foo' + + expect(lb) + .to receive(:select_up_to_date_host) + .with('foo') + .and_return(true) + end + + expect(middleware.send(:databases_in_sync?, locations)) + .to eq(true) + end + + it 'returns false when the load balancers are not in sync' do + locations = {} + + Gitlab::Database::LoadBalancing.each_load_balancer do |lb| + locations[lb.name] = 'foo' + + allow(lb) + .to receive(:select_up_to_date_host) + .with('foo') + .and_return(false) + end + + expect(middleware.send(:databases_in_sync?, locations)) + .to eq(false) + end + end + def process_job(job) Sidekiq::JobRetry.new.local(worker_class, job.to_json, 'default') do worker_class.process_job(job) @@ -208,6 +261,8 @@ def run_middleware end def replication_lag!(exists) - allow(load_balancer).to receive(:select_up_to_date_host).and_return(!exists) + Gitlab::Database::LoadBalancing.each_load_balancer do |lb| + allow(lb).to receive(:select_up_to_date_host).and_return(!exists) + end end end diff --git a/spec/lib/gitlab/database/load_balancing/sticking_spec.rb b/spec/lib/gitlab/database/load_balancing/sticking_spec.rb index 1087ff6dd781d3..4b0eb8858ddc59 100644 --- a/spec/lib/gitlab/database/load_balancing/sticking_spec.rb +++ b/spec/lib/gitlab/database/load_balancing/sticking_spec.rb @@ -3,19 +3,60 @@ require 'spec_helper' RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do + let(:sticking) do + described_class.new(ActiveRecord::Base.connection.load_balancer) + end + after do Gitlab::Database::LoadBalancing::Session.clear_session end - describe '.stick_if_necessary' do + describe '#stick_or_unstick_request' do + it 'sticks or unsticks a single object and updates the Rack environment' do + expect(sticking) + .to receive(:unstick_or_continue_sticking) + .with(:user, 42) + + env = {} + + sticking.stick_or_unstick_request(env, :user, 42) + + expect(env[Gitlab::Database::LoadBalancing::RackMiddleware::STICK_OBJECT].to_a) + .to eq([[ActiveRecord::Base, :user, 42]]) + end + + it 'sticks or unsticks multiple objects and updates the Rack environment' do + expect(sticking) + .to receive(:unstick_or_continue_sticking) + .with(:user, 42) + .ordered + + expect(sticking) + .to receive(:unstick_or_continue_sticking) + .with(:runner, '123456789') + .ordered + + env = {} + + sticking.stick_or_unstick_request(env, :user, 42) + sticking.stick_or_unstick_request(env, :runner, '123456789') + + expect(env[Gitlab::Database::LoadBalancing::RackMiddleware::STICK_OBJECT].to_a).to eq([ + [ActiveRecord::Base, :user, 42], + [ActiveRecord::Base, :runner, '123456789'] + ]) + end + end + + describe '#stick_if_necessary' do 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(sticking).not_to receive(:stick) - described_class.stick_if_necessary(:user, 42) + sticking.stick_if_necessary(:user, 42) end it 'sticks to the primary if a write was performed' do @@ -23,20 +64,21 @@ .to receive(:performed_write?) .and_return(true) - expect(described_class).to receive(:stick).with(:user, 42) + expect(sticking) + .to receive(:stick) + .with(:user, 42) - described_class.stick_if_necessary(:user, 42) + sticking.stick_if_necessary(:user, 42) end end - describe '.all_caught_up?' do - let(:lb) { double(:lb) } + describe '#all_caught_up?' do + let(:lb) { ActiveRecord::Base.connection.load_balancer } let(:last_write_location) { 'foo' } before do - allow(described_class).to receive(:load_balancer).and_return(lb) - - allow(described_class).to receive(:last_write_location_for) + allow(sticking) + .to receive(:last_write_location_for) .with(:user, 42) .and_return(last_write_location) end @@ -45,13 +87,9 @@ let(:last_write_location) { nil } it 'returns true' do - allow(described_class).to receive(:last_write_location_for) - .with(:user, 42) - .and_return(nil) - expect(lb).not_to receive(:select_up_to_date_host) - expect(described_class.all_caught_up?(:user, 42)).to eq(true) + expect(sticking.all_caught_up?(:user, 42)).to eq(true) end end @@ -61,9 +99,11 @@ end it 'returns true, and unsticks' do - expect(described_class).to receive(:unstick).with(:user, 42) + expect(sticking) + .to receive(:unstick) + .with(:user, 42) - expect(described_class.all_caught_up?(:user, 42)).to eq(true) + expect(sticking.all_caught_up?(:user, 42)).to eq(true) end it 'notifies with the proper event payload' do @@ -72,7 +112,7 @@ .with('caught_up_replica_pick.load_balancing', { result: true }) .and_call_original - described_class.all_caught_up?(:user, 42) + sticking.all_caught_up?(:user, 42) end end @@ -82,7 +122,7 @@ end it 'returns false' do - expect(described_class.all_caught_up?(:user, 42)).to eq(false) + expect(sticking.all_caught_up?(:user, 42)).to eq(false) end it 'notifies with the proper event payload' do @@ -91,42 +131,43 @@ .with('caught_up_replica_pick.load_balancing', { result: false }) .and_call_original - described_class.all_caught_up?(:user, 42) + sticking.all_caught_up?(:user, 42) end end end - describe '.unstick_or_continue_sticking' do - let(:lb) { double(:lb) } - - before do - allow(described_class).to receive(:load_balancer).and_return(lb) - end + describe '#unstick_or_continue_sticking' do + let(:lb) { ActiveRecord::Base.connection.load_balancer } it 'simply returns if no write location could be found' do - allow(described_class).to receive(:last_write_location_for) + allow(sticking) + .to receive(:last_write_location_for) .with(:user, 42) .and_return(nil) expect(lb).not_to receive(:select_up_to_date_host) - described_class.unstick_or_continue_sticking(:user, 42) + sticking.unstick_or_continue_sticking(:user, 42) end it 'unsticks if all secondaries have caught up' do - allow(described_class).to receive(:last_write_location_for) + allow(sticking) + .to receive(:last_write_location_for) .with(:user, 42) .and_return('foo') allow(lb).to receive(:select_up_to_date_host).with('foo').and_return(true) - expect(described_class).to receive(:unstick).with(:user, 42) + expect(sticking) + .to receive(:unstick) + .with(:user, 42) - described_class.unstick_or_continue_sticking(:user, 42) + sticking.unstick_or_continue_sticking(:user, 42) end it 'continues using the primary if the secondaries have not yet caught up' do - allow(described_class).to receive(:last_write_location_for) + allow(sticking) + .to receive(:last_write_location_for) .with(:user, 42) .and_return('foo') @@ -135,21 +176,22 @@ expect(Gitlab::Database::LoadBalancing::Session.current) .to receive(:use_primary!) - described_class.unstick_or_continue_sticking(:user, 42) + sticking.unstick_or_continue_sticking(:user, 42) end end RSpec.shared_examples 'sticking' do before do - lb = double(:lb, primary_write_location: 'foo') - - allow(described_class).to receive(:load_balancer).and_return(lb) + allow(ActiveRecord::Base.connection.load_balancer) + .to receive(:primary_write_location) + .and_return('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') + expect(sticking) + .to receive(:set_write_location_for) + .with(:user, id, 'foo') end expect(Gitlab::Database::LoadBalancing::Session.current) @@ -159,99 +201,106 @@ end end - describe '.stick' do + describe '#stick' do it_behaves_like 'sticking' do let(:ids) { [42] } - subject { described_class.stick(:user, ids.first) } + subject { sticking.stick(:user, ids.first) } end end - describe '.bulk_stick' do + describe '#bulk_stick' do it_behaves_like 'sticking' do let(:ids) { [42, 43] } - subject { described_class.bulk_stick(:user, ids) } + subject { sticking.bulk_stick(:user, ids) } end end - describe '.mark_primary_write_location' do + describe '#mark_primary_write_location' do 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) + allow(ActiveRecord::Base.connection.load_balancer) + .to receive(:primary_write_location) + .and_return('foo') - expect(described_class).to receive(:set_write_location_for) + expect(sticking) + .to receive(:set_write_location_for) .with(:user, 42, 'foo') - described_class.mark_primary_write_location(:user, 42) + sticking.mark_primary_write_location(:user, 42) end end - describe '.unstick' do + describe '#unstick' do it 'removes the sticking data from Redis' do - described_class.set_write_location_for(:user, 4, 'foo') - described_class.unstick(:user, 4) + sticking.set_write_location_for(:user, 4, 'foo') + sticking.unstick(:user, 4) - expect(described_class.last_write_location_for(:user, 4)).to be_nil + expect(sticking.last_write_location_for(:user, 4)).to be_nil end - end - - describe '.last_write_location_for' do - it 'returns the last WAL write location for a user' do - described_class.set_write_location_for(:user, 4, 'foo') - expect(described_class.last_write_location_for(:user, 4)).to eq('foo') - end - end + it 'removes the old key' do + Gitlab::Redis::SharedState.with do |redis| + redis.set(sticking.send(:old_redis_key_for, :user, 4), 'foo', ex: 30) + end - describe '.redis_key_for' do - it 'returns a String' do - expect(described_class.redis_key_for(:user, 42)) - .to eq('database-load-balancing/write-location/user/42') + sticking.unstick(:user, 4) + expect(sticking.last_write_location_for(:user, 4)).to be_nil end end - describe '.load_balancer' do - it 'returns a the load balancer' do - proxy = double(:proxy) + describe '#last_write_location_for' do + it 'returns the last WAL write location for a user' do + sticking.set_write_location_for(:user, 4, 'foo') - expect(Gitlab::Database::LoadBalancing).to receive(:proxy) - .and_return(proxy) + expect(sticking.last_write_location_for(:user, 4)).to eq('foo') + end - expect(proxy).to receive(:load_balancer) + it 'falls back to reading the old key' do + Gitlab::Redis::SharedState.with do |redis| + redis.set(sticking.send(:old_redis_key_for, :user, 4), 'foo', ex: 30) + end - described_class.load_balancer + expect(sticking.last_write_location_for(:user, 4)).to eq('foo') end end - describe '.select_caught_up_replicas' do - let(:lb) { double(:lb) } - - before do - allow(described_class).to receive(:load_balancer).and_return(lb) + describe '#redis_key_for' do + it 'returns a String' do + expect(sticking.redis_key_for(:user, 42)) + .to eq('database-load-balancing/write-location/main/user/42') end + end + + describe '#select_caught_up_replicas' do + let(:lb) { ActiveRecord::Base.connection.load_balancer } context 'with no write location' do before do - allow(described_class).to receive(:last_write_location_for) - .with(:project, 42).and_return(nil) + allow(sticking) + .to receive(:last_write_location_for) + .with(:project, 42) + .and_return(nil) end it 'returns false and does not try to find caught up hosts' do expect(lb).not_to receive(:select_up_to_date_host) - expect(described_class.select_caught_up_replicas(:project, 42)).to be false + expect(sticking.select_caught_up_replicas(:project, 42)).to be false end end context 'with write location' do before do - allow(described_class).to receive(:last_write_location_for) - .with(:project, 42).and_return('foo') + allow(sticking) + .to receive(:last_write_location_for) + .with(:project, 42) + .and_return('foo') end it 'returns true, selects hosts, and unsticks if any secondary has caught up' do expect(lb).to receive(:select_up_to_date_host).and_return(true) - expect(described_class).to receive(:unstick).with(:project, 42) - expect(described_class.select_caught_up_replicas(:project, 42)).to be true + expect(sticking) + .to receive(:unstick) + .with(:project, 42) + expect(sticking.select_caught_up_replicas(:project, 42)).to be true end end end diff --git a/spec/lib/gitlab/database/load_balancing_spec.rb b/spec/lib/gitlab/database/load_balancing_spec.rb index 58b03286260ae8..bf5314e2c34e48 100644 --- a/spec/lib/gitlab/database/load_balancing_spec.rb +++ b/spec/lib/gitlab/database/load_balancing_spec.rb @@ -3,173 +3,48 @@ require 'spec_helper' RSpec.describe Gitlab::Database::LoadBalancing do - describe '.proxy' do - it 'returns the connection proxy' do - proxy = double(:connection_proxy) + describe '.base_models' do + it 'returns the models to apply load balancing to' do + models = described_class.base_models - allow(ActiveRecord::Base) - .to receive(:load_balancing_proxy) - .and_return(proxy) + expect(models).to include(ActiveRecord::Base) - expect(described_class.proxy).to eq(proxy) - end - end - - describe '.configuration' do - it 'returns the configuration for the load balancer' do - raw = ActiveRecord::Base.connection_db_config.configuration_hash - cfg = described_class.configuration - - # There isn't much to test here as the load balancing settings might not - # (and likely aren't) set when running tests. - expect(cfg.pool_size).to eq(raw[:pool]) - end - end - - describe '.enable_replicas?' do - context 'when hosts are specified' do - before do - allow(described_class.configuration) - .to receive(:hosts) - .and_return(%w(foo)) - 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) - - 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) - - expect(described_class.enable_replicas?).to eq(false) + if Gitlab::Database.has_config?(:ci) + expect(models).to include(Ci::CiDatabaseRecord) end end - 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) - - allow(described_class.configuration) - .to receive(:service_discovery_enabled?) - .and_return(true) - - expect(described_class.enable_replicas?).to eq(true) - end - end - - 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_replicas?).to eq(false) - end + it 'returns the models as a frozen array' do + expect(described_class.base_models).to be_frozen end end - describe '.configured?' do - it 'returns true when hosts are configured' do - allow(described_class.configuration) - .to receive(:hosts) - .and_return(%w[foo]) + describe '.each_load_balancer' do + it 'yields every load balancer to the supplied block' do + lbs = [] - expect(described_class.configured?).to eq(true) - end - - it 'returns true when service discovery is enabled' do - allow(described_class.configuration).to receive(:hosts).and_return([]) - allow(described_class.configuration) - .to receive(:service_discovery_enabled?) - .and_return(true) - - expect(described_class.configured?).to eq(true) - end - - it 'returns false when neither service discovery nor hosts are configured' 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.configured?).to eq(false) - end - end - - describe '.start_service_discovery' do - it 'does not start if service discovery is disabled' do - expect(Gitlab::Database::LoadBalancing::ServiceDiscovery) - .not_to receive(:new) + described_class.each_load_balancer do |lb| + lbs << lb + end - described_class.start_service_discovery + expect(lbs.length).to eq(described_class.base_models.length) end - it 'starts service discovery if enabled' do - allow(described_class.configuration) - .to receive(:service_discovery_enabled?) - .and_return(true) - - instance = double(:instance) - config = Gitlab::Database::LoadBalancing::Configuration - .new(ActiveRecord::Base) - lb = Gitlab::Database::LoadBalancing::LoadBalancer.new(config) - proxy = Gitlab::Database::LoadBalancing::ConnectionProxy.new(lb) + it 'returns an Enumerator when no block is given' do + res = described_class.each_load_balancer - allow(described_class) - .to receive(:proxy) - .and_return(proxy) - - expect(Gitlab::Database::LoadBalancing::ServiceDiscovery) - .to receive(:new) - .with(lb, an_instance_of(Hash)) - .and_return(instance) - - expect(instance) - .to receive(:start) - - described_class.start_service_discovery + expect(res.next) + .to be_an_instance_of(Gitlab::Database::LoadBalancing::LoadBalancer) end end - describe '.perform_service_discovery' do - it 'does nothing if service discovery is disabled' do - expect(Gitlab::Database::LoadBalancing::ServiceDiscovery) - .not_to receive(:new) - - 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) + describe '.release_hosts' do + it 'releases the host of every load balancer' do + described_class.each_load_balancer do |lb| + expect(lb).to receive(:release_host) + end - described_class.perform_service_discovery + described_class.release_hosts end end @@ -227,7 +102,7 @@ # - In each test, we listen to the SQL queries (via sql.active_record # instrumentation) while triggering real queries from the defined model. # - We assert the desinations (replica/primary) of the queries in order. - describe 'LoadBalancing integration tests', :db_load_balancing, :delete do + describe 'LoadBalancing integration tests', :database_replica, :delete do before(:all) do ActiveRecord::Schema.define do create_table :load_balancing_test, force: true do |t| diff --git a/spec/lib/gitlab/database_spec.rb b/spec/lib/gitlab/database_spec.rb index 0c1b61428f97fb..a2e7b6d27b9af1 100644 --- a/spec/lib/gitlab/database_spec.rb +++ b/spec/lib/gitlab/database_spec.rb @@ -203,7 +203,7 @@ .to eq('main') end - context 'when replicas are configured', :db_load_balancing do + context 'when replicas are configured', :database_replica do it 'returns the name for a replica' do replica = ActiveRecord::Base.connection.load_balancer.host diff --git a/spec/lib/gitlab/import_export/json/streaming_serializer_spec.rb b/spec/lib/gitlab/import_export/json/streaming_serializer_spec.rb index f54a57f806e5da..d69d775fffb492 100644 --- a/spec/lib/gitlab/import_export/json/streaming_serializer_spec.rb +++ b/spec/lib/gitlab/import_export/json/streaming_serializer_spec.rb @@ -158,7 +158,7 @@ end describe 'load balancing' do - context 'when feature flag load_balancing_for_export_workers is enabled', :db_load_balancing do + context 'when feature flag load_balancing_for_export_workers is enabled' do before do stub_feature_flags(load_balancing_for_export_workers: true) end diff --git a/spec/lib/gitlab/metrics/subscribers/active_record_spec.rb b/spec/lib/gitlab/metrics/subscribers/active_record_spec.rb index b6b6c1ffa1f10c..a8e4f039da4c59 100644 --- a/spec/lib/gitlab/metrics/subscribers/active_record_spec.rb +++ b/spec/lib/gitlab/metrics/subscribers/active_record_spec.rb @@ -195,7 +195,7 @@ def sql(query, comments: true) with_them do let(:payload) { { name: name, sql: sql(sql_query, comments: comments), connection: connection } } - context 'query using a connection to a replica', :db_load_balancing do + context 'query using a connection to a replica' do before do allow(Gitlab::Database::LoadBalancing).to receive(:db_role_for_connection).and_return(:replica) end diff --git a/spec/lib/gitlab/sidekiq_logging/structured_logger_spec.rb b/spec/lib/gitlab/sidekiq_logging/structured_logger_spec.rb index f10ca1e9216994..d801b84775b3de 100644 --- a/spec/lib/gitlab/sidekiq_logging/structured_logger_spec.rb +++ b/spec/lib/gitlab/sidekiq_logging/structured_logger_spec.rb @@ -317,7 +317,7 @@ end end - context 'when load balancing is enabled', :db_load_balancing do + context 'when load balancing is enabled' do let(:db_config_name) do ::Gitlab::Database.db_config_name(ApplicationRecord.retrieve_connection) end diff --git a/spec/models/ci/build_spec.rb b/spec/models/ci/build_spec.rb index 3df892fe2ef7b8..9ba7d0ab35967f 100644 --- a/spec/models/ci/build_spec.rb +++ b/spec/models/ci/build_spec.rb @@ -347,10 +347,10 @@ end describe '#stick_build_if_status_changed' do - it 'sticks the build if the status changed', :db_load_balancing do + it 'sticks the build if the status changed' do job = create(:ci_build, :pending) - expect(Gitlab::Database::LoadBalancing::Sticking).to receive(:stick) + expect(ApplicationRecord.sticking).to receive(:stick) .with(:build, job.id) job.update!(status: :running) diff --git a/spec/models/ci/pipeline_spec.rb b/spec/models/ci/pipeline_spec.rb index a19acbd196db62..dc2ffffa873a43 100644 --- a/spec/models/ci/pipeline_spec.rb +++ b/spec/models/ci/pipeline_spec.rb @@ -2790,7 +2790,16 @@ def create_pipeline(status, ref, sha) extra_update_queries = 4 # transition ... => :canceled, queue pop extra_generic_commit_status_validation_queries = 2 # name_uniqueness_across_types - extra_load_balancer_queries = 3 + + # The number of extra load balancing queries depends on whether or not + # we use a load balancer for CI. That in turn depends on the contents of + # database.yml, so here we support both cases. + extra_load_balancer_queries = + if Gitlab::Database.has_config?(:ci) + 6 + else + 3 + end expect(control2.count).to eq(control1.count + extra_update_queries + extra_generic_commit_status_validation_queries + extra_load_balancer_queries) end diff --git a/spec/models/ci/runner_spec.rb b/spec/models/ci/runner_spec.rb index 5364dfc972c9e2..b28eed04316fc2 100644 --- a/spec/models/ci/runner_spec.rb +++ b/spec/models/ci/runner_spec.rb @@ -397,7 +397,7 @@ def stub_redis_runner_contacted_at(value) it 'sticks the runner to the primary and calls the original method' do runner = create(:ci_runner) - expect(Gitlab::Database::LoadBalancing::Sticking).to receive(:stick) + expect(ApplicationRecord.sticking).to receive(:stick) .with(:runner, runner.id) expect(Gitlab::Workhorse).to receive(:set_key_and_notify) diff --git a/spec/models/project_feature_usage_spec.rb b/spec/models/project_feature_usage_spec.rb index 698c5374e88919..3765a2b37a7ebb 100644 --- a/spec/models/project_feature_usage_spec.rb +++ b/spec/models/project_feature_usage_spec.rb @@ -133,10 +133,8 @@ subject { project.feature_usage } - context 'database load balancing is configured', :db_load_balancing do + context 'database load balancing is configured' do before do - allow(ActiveRecord::Base).to receive(:connection).and_return(::Gitlab::Database::LoadBalancing.proxy) - ::Gitlab::Database::LoadBalancing::Session.clear_session end diff --git a/spec/models/project_spec.rb b/spec/models/project_spec.rb index f192d2971cb90d..0919ebe8ce40aa 100644 --- a/spec/models/project_spec.rb +++ b/spec/models/project_spec.rb @@ -3050,7 +3050,7 @@ def has_external_wiki let(:project) { create(:project) } it 'marks the location with project ID' do - expect(Gitlab::Database::LoadBalancing::Sticking).to receive(:mark_primary_write_location).with(:project, project.id) + expect(ApplicationRecord.sticking).to receive(:mark_primary_write_location).with(:project, project.id) project.mark_primary_write_location end diff --git a/spec/requests/api/graphql/mutations/merge_requests/set_assignees_spec.rb b/spec/requests/api/graphql/mutations/merge_requests/set_assignees_spec.rb index dec9afd1310bd6..608b36e4f15fdd 100644 --- a/spec/requests/api/graphql/mutations/merge_requests/set_assignees_spec.rb +++ b/spec/requests/api/graphql/mutations/merge_requests/set_assignees_spec.rb @@ -115,7 +115,7 @@ def run_mutation! context 'when passing append as true' do let(:mode) { Types::MutationOperationModeEnum.enum[:append] } let(:input) { { assignee_usernames: [assignee2.username], operation_mode: mode } } - let(:db_query_limit) { 21 } + let(:db_query_limit) { 22 } before do # In CE, APPEND is a NOOP as you can't have multiple assignees diff --git a/spec/services/ci/drop_pipeline_service_spec.rb b/spec/services/ci/drop_pipeline_service_spec.rb index c6a118c6083106..ddb53712d9ca5a 100644 --- a/spec/services/ci/drop_pipeline_service_spec.rb +++ b/spec/services/ci/drop_pipeline_service_spec.rb @@ -50,13 +50,14 @@ def drop_pipeline!(pipeline) end.count writes_per_build = 2 + load_balancer_queries = 3 expected_reads_count = control_count - writes_per_build create_list(:ci_build, 5, :running, pipeline: cancelable_pipeline) expect do drop_pipeline!(cancelable_pipeline) - end.not_to exceed_query_limit(expected_reads_count + (5 * writes_per_build)) + end.not_to exceed_query_limit(expected_reads_count + (5 * writes_per_build) + load_balancer_queries) end end end diff --git a/spec/services/ci/register_job_service_spec.rb b/spec/services/ci/register_job_service_spec.rb index d1a7b2a822b1a7..650353eb751acb 100644 --- a/spec/services/ci/register_job_service_spec.rb +++ b/spec/services/ci/register_job_service_spec.rb @@ -14,7 +14,7 @@ module Ci let!(:pending_job) { create(:ci_build, :pending, :queued, pipeline: pipeline) } describe '#execute' do - context 'checks database loadbalancing stickiness', :db_load_balancing do + context 'checks database loadbalancing stickiness' do subject { described_class.new(shared_runner).execute } before do @@ -22,14 +22,14 @@ module Ci end it 'result is valid if replica did caught-up' do - expect(Gitlab::Database::LoadBalancing::Sticking).to receive(:all_caught_up?) + expect(ApplicationRecord.sticking).to receive(:all_caught_up?) .with(:runner, shared_runner.id) { true } expect(subject).to be_valid end it 'result is invalid if replica did not caught-up' do - expect(Gitlab::Database::LoadBalancing::Sticking).to receive(:all_caught_up?) + expect(ApplicationRecord.sticking).to receive(:all_caught_up?) .with(:runner, shared_runner.id) { false } expect(subject).not_to be_valid diff --git a/spec/services/user_project_access_changed_service_spec.rb b/spec/services/user_project_access_changed_service_spec.rb index 3af23fafef5db3..438db6b987b13c 100644 --- a/spec/services/user_project_access_changed_service_spec.rb +++ b/spec/services/user_project_access_changed_service_spec.rb @@ -53,7 +53,7 @@ end it 'sticks all the updated users and returns the original result', :aggregate_failures do - expect(Gitlab::Database::LoadBalancing::Sticking).to receive(:bulk_stick).with(:user, [1, 2]) + expect(ApplicationRecord.sticking).to receive(:bulk_stick).with(:user, [1, 2]) expect(service.execute).to eq(10) end diff --git a/spec/services/users/activity_service_spec.rb b/spec/services/users/activity_service_spec.rb index 6c1df5c745f0ca..092c5cd3e5e240 100644 --- a/spec/services/users/activity_service_spec.rb +++ b/spec/services/users/activity_service_spec.rb @@ -91,9 +91,9 @@ context 'when last activity is in the past' do let(:user) { create(:user, last_activity_on: Date.today - 1.week) } - context 'database load balancing is configured', :db_load_balancing do + context 'database load balancing is configured' do before do - allow(ActiveRecord::Base).to receive(:connection).and_return(::Gitlab::Database::LoadBalancing.proxy) + ::Gitlab::Database::LoadBalancing::Session.clear_session end let(:service) do diff --git a/spec/support/database_load_balancing.rb b/spec/support/database_load_balancing.rb index d5573162c136fe..014575e8a820f2 100644 --- a/spec/support/database_load_balancing.rb +++ b/spec/support/database_load_balancing.rb @@ -1,20 +1,30 @@ # frozen_string_literal: true RSpec.configure do |config| - config.before(:each, :db_load_balancing) do - config = Gitlab::Database::LoadBalancing::Configuration - .new(ActiveRecord::Base, [Gitlab::Database.main.config['host']]) - lb = ::Gitlab::Database::LoadBalancing::LoadBalancer.new(config) - proxy = ::Gitlab::Database::LoadBalancing::ConnectionProxy.new(lb) + config.around(:each, :database_replica) do |example| + old_proxies = [] - allow(ActiveRecord::Base).to receive(:load_balancing_proxy).and_return(proxy) + Gitlab::Database::LoadBalancing.base_models.each do |model| + config = Gitlab::Database::LoadBalancing::Configuration + .new(model, [model.connection_db_config.configuration_hash[:host]]) + lb = Gitlab::Database::LoadBalancing::LoadBalancer.new(config) - ::Gitlab::Database::LoadBalancing::Session.clear_session + old_proxies << [model, model.connection] + + model.connection = + Gitlab::Database::LoadBalancing::ConnectionProxy.new(lb) + end + + Gitlab::Database::LoadBalancing::Session.clear_session redis_shared_state_cleanup! - end - config.after(:each, :db_load_balancing) do - ::Gitlab::Database::LoadBalancing::Session.clear_session + example.run + + Gitlab::Database::LoadBalancing::Session.clear_session redis_shared_state_cleanup! + + old_proxies.each do |(model, proxy)| + model.connection = proxy + end end end diff --git a/spec/support/shared_examples/lib/gitlab/ci/ci_trace_shared_examples.rb b/spec/support/shared_examples/lib/gitlab/ci/ci_trace_shared_examples.rb index c768ba912cce7e..8b4ecd7d5ae65c 100644 --- a/spec/support/shared_examples/lib/gitlab/ci/ci_trace_shared_examples.rb +++ b/spec/support/shared_examples/lib/gitlab/ci/ci_trace_shared_examples.rb @@ -35,8 +35,8 @@ stub_feature_flags(gitlab_ci_archived_trace_consistent_reads: trace.job.project) end - it 'calls ::Gitlab::Database::LoadBalancing::Sticking.unstick_or_continue_sticking' do - expect(::Gitlab::Database::LoadBalancing::Sticking).to receive(:unstick_or_continue_sticking) + it 'calls ::ApplicationRecord.sticking.unstick_or_continue_sticking' do + expect(::ApplicationRecord.sticking).to receive(:unstick_or_continue_sticking) .with(described_class::LOAD_BALANCING_STICKING_NAMESPACE, trace.job.id) .and_call_original @@ -49,8 +49,8 @@ stub_feature_flags(gitlab_ci_archived_trace_consistent_reads: false) end - it 'does not call ::Gitlab::Database::LoadBalancing::Sticking.unstick_or_continue_sticking' do - expect(::Gitlab::Database::LoadBalancing::Sticking).not_to receive(:unstick_or_continue_sticking) + it 'does not call ::ApplicationRecord.sticking.unstick_or_continue_sticking' do + expect(::ApplicationRecord.sticking).not_to receive(:unstick_or_continue_sticking) trace.read { |stream| stream } end @@ -305,8 +305,8 @@ stub_feature_flags(gitlab_ci_archived_trace_consistent_reads: trace.job.project) end - it 'calls ::Gitlab::Database::LoadBalancing::Sticking.stick' do - expect(::Gitlab::Database::LoadBalancing::Sticking).to receive(:stick) + it 'calls ::ApplicationRecord.sticking.stick' do + expect(::ApplicationRecord.sticking).to receive(:stick) .with(described_class::LOAD_BALANCING_STICKING_NAMESPACE, trace.job.id) .and_call_original @@ -319,8 +319,8 @@ stub_feature_flags(gitlab_ci_archived_trace_consistent_reads: false) end - it 'does not call ::Gitlab::Database::LoadBalancing::Sticking.stick' do - expect(::Gitlab::Database::LoadBalancing::Sticking).not_to receive(:stick) + it 'does not call ::ApplicationRecord.sticking.stick' do + expect(::ApplicationRecord.sticking).not_to receive(:stick) subject end diff --git a/spec/workers/authorized_project_update/user_refresh_from_replica_worker_spec.rb b/spec/workers/authorized_project_update/user_refresh_from_replica_worker_spec.rb index 0da5834377336a..da4b726c0b5eed 100644 --- a/spec/workers/authorized_project_update/user_refresh_from_replica_worker_spec.rb +++ b/spec/workers/authorized_project_update/user_refresh_from_replica_worker_spec.rb @@ -44,7 +44,7 @@ end end - context 'with load balancing enabled', :db_load_balancing do + context 'with load balancing enabled' do it 'reads from the replica database' do expect(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_replicas_for_read_queries).and_call_original diff --git a/spec/workers/container_expiration_policy_worker_spec.rb b/spec/workers/container_expiration_policy_worker_spec.rb index 9f370b10f6a0c2..ebf8004115193d 100644 --- a/spec/workers/container_expiration_policy_worker_spec.rb +++ b/spec/workers/container_expiration_policy_worker_spec.rb @@ -156,7 +156,7 @@ subject end - context 'with load balancing enabled', :db_load_balancing do + context 'with load balancing enabled' do it 'reads the counts from the replica' do expect(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_replicas_for_read_queries).and_call_original -- GitLab