diff --git a/app/models/ci/build.rb b/app/models/ci/build.rb index 5709bc2695684324ba1c606c973ec51313f65ca1..1e0ce5638bb8a8347c30b1a3aa693dac77ca61a0 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 e76b7a05b93a1cbaea0ab1eaa8e1502dd7b3958b..e2b832a28e7742a85e25560ffd4f140172ef8b57 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 4c5e4837d0687af54b1effa975073ed3e3ff3648..dd8cdbcaf4ba6dc0bef6e0d949467eb9d8f2dbdd 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 db257787a40b6a650b7d7048965d989b0b9fdba0..6841fcddc0bf570a12c3ca63a2e6827a20c6dc4c 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 5dd586e9f22c4e6a57219a48b7ff20571337123d..67ef4f1070956b5f1fb370558728171e0cbcf72a 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 5f48f410bf77eaff02b5e6fac4047bc93b1733e6..5bba986f4ad87e249a5d0705ca6660a7daea7cce 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 29700b37c7d9a0e02744791e6b5b9cc29aff3d3b..a31b11bb2bef81fb15a580e11e50f45fa059d8d0 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 b9662b822fb61767c44fc628a2210c92b37f7249..dabb6c7ab3a6c837b97970ea66ce3e473e2c0546 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 c189fde503a80090fcf53d98db7277ac24a0457b..ff3590d6c13956a51dab91e34d805aad941caeae 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 c86779c7a372863d9e0ee1c9d190c9477ca4f9a2..e5ce862264f414aab82502cb595dfe4fda0cc097 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 6816d216015ed548417e38da57880c098208808b..c8f19d65c9eafd672649c82d0c7fdbe4f7b7c7d2 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 d56f2a4b0e3a411bcb24f69d880d1a947eaaec8a..b560d4cbca84d7579db937e351e8b1633870caa5 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 a5cfaf795f9eb0f51ae9bb6ca565dd0fe1c49b18..3e322e752b7a03887c67fe2f64cf6485e5640ae1 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 4feba989a0aec432c03e4471c2398060c1bc9057..7164976ff73b256a78b3f3e69bf7515418603377 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 5462c2a3188c83a018a49fec1fffbf878ad52631..0000000000000000000000000000000000000000 --- 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 238f55fd98e74697cb3fd87e076125b970a3bb5e..6156515bd73bed3f0975671c501616d08942ba14 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 5d9dbfdabbdf43ffcbff3b1702241bf0dfe2077d..cc9ca3253378486721c52f29631c09009af4765a 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 22871427bff7b58862deb471b7c969bc46184471..fc035f42097cd0b3a2ade52dd9f4b98a0eff7cc6 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 0000000000000000000000000000000000000000..3cce839a960f73fe874de8bed964d6850317329a --- /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 adda6c1f6e7d3ee23091dbbe978e8ecc2e57840e..ebf90ebf9d0ce0c5b04f20feff79e58229d9e1af 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 5e391c7228da7be2d2fa0e468db1a3c2db1c0b90..f0c7016032bb78fb5099bb3bb7616520e7039a0b 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 8d55ade488026a0c7c5e227ffba2453c357094fc..aa5808e61b01276a8b8b3dc16d51dc70be112eea 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 99f2db544a5067a4ed2bcae662aa5a0f47d012e3..cc871d66d4008e15a30911e0efd4c559e6fb965f 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 c966527ca5b39157eb7940d0a0cd4e0930f4a80d..37e040a422bee3f43e2c23b07091e496c2fa8269 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 dce6a27191894d52c053ac04fcd1d133cb792a1c..c65a1e4d6566e25aaee60c01bc78437bddd9241e 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 b3d166d58bbf0223ee85be064f42f658272ee10e..5055be81c880bfc0f7b8de594df907c841d58f10 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 ebbbafb855ffa14c0a22dde0ef541a301611c1f3..768855464c109b4b846cd402ce980af1256f3449 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 2ccb52f7b145ecb4ad750ba12892b46df53a4277..0000000000000000000000000000000000000000 --- 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 6621e6276a53096465f22b47baa9440ae589fee8..3e5249a3deaa5873f74f6f07cc2e5d427ecdee07 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 fb17ba0e1217969f28c83a61bd23f0e6d263918c..f3ce5563e3883f0114175e246f8997b0bbc1f5d3 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 5453e36e08fe68faf3d3d4ea73d991a18107bc37..af7e2a4b167c6d1b5470d581367acd56f6467bda 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 0000000000000000000000000000000000000000..01646bc76efbbe4e3815cd65fa94d3c93c7c2a3e --- /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 18d9cc9e170565905790c0897430fdb6283b1c70..1d00cba0ecadfda2e38cb9210cc0207a9f06826b 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 7d3670c527110a2bc12a8a29d91cdd0dbd8bd733..06efdcd8f9987a93969ce7bf2685ccb5c9978454 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 1087ff6dd781d36f089690d60a1a1de2080a8938..4b0eb8858ddc59f091ab9f2c520c569a6eddbccb 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 58b03286260ae81152bd3374de10f4a87b12a896..bf5314e2c34e489911828f834e14d74bde51e5ed 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 0c1b61428f97fb714a245dfe23bdfabec5ecda7a..a2e7b6d27b9af19cb9801cea3a5b8d5e0e44b308 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 f54a57f806e5da4ba673ba3f0059d29224963219..d69d775fffb4927f87944e392dd892309c4aef82 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 b6b6c1ffa1f10c8e3ad478ef05326b1e3278632f..a8e4f039da4c5916e7c0200b3cde8412c8ac9a51 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 f10ca1e9216994db758a941eb175436b834d1f84..d801b84775b3de7b9496956361ed6220d3219a7c 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 3df892fe2ef7b80f5c0fcfcb8e91988c26b16b9f..9ba7d0ab35967f1070c5a05c1e9953c1c22c2c33 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 a19acbd196db624b591cb7ae5f4f508a77d5dc7a..dc2ffffa873a43f4fc78bcb42c09177e8302bb0a 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 5364dfc972c9e2fd0ddedf2902762523a4d8ddcb..b28eed04316fc2af2196a1ceb31c5b350304eb90 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 698c5374e889190b2b82a387ca64335a2e2ff266..3765a2b37a7ebbcd63f6a0ed7411c4b610d20150 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 f192d2971cb90dad06281955a63d7c3f83a33ef4..0919ebe8ce40aa3a8ab2e4d62b18d482488e5e44 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 dec9afd1310bd65bdfc14b86dade97349a340687..608b36e4f15fdd57a5f9be0bf1f2f6c2079d6e3a 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 c6a118c60831069211edaca79db3812cbbb88a34..ddb53712d9ca5a9e9275d509a2287994b716514d 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 d1a7b2a822b1a7b2d85e22813d149e2090159e81..650353eb751acb2b3a744b572632f7940fafa39f 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 3af23fafef5db32ab12e9223e6f5bb736ae3f4b7..438db6b987b13cde83bc796df1554e2753d0c29a 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 6c1df5c745f0cab8ecd12ee098f8bdfcfbbfc311..092c5cd3e5e240daf05408a611202c9e4e40dfdf 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 d5573162c136fe319fbb9445f9f5fe9fa3efe77e..014575e8a820f24f0cdf743723d48b3eefb0521c 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 c768ba912cce7eac234793a32b23e38263ecf944..8b4ecd7d5ae65cd3fbcc85b95ba2d5991d4e3fe4 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 0da5834377336af180469345fe2eee8a70c8a36d..da4b726c0b5eed7ff9e58e63405d40c20843534f 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 9f370b10f6a0c23775c782417a51da95796931ea..ebf8004115193d881f0c9456928135596d6bf437 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