Merge branch 'move-load-balancing-proxy-to-activerecord' into 'master'

Store the load balancer proxy in ActiveRecord

See merge request gitlab-org/gitlab!65805

config/initializers/load_balancing.rb

 # frozen_string_literal: true
 
+ActiveRecord::Base.singleton_class.attr_accessor :load_balancing_proxy
+
 if Gitlab::Database::LoadBalancing.enable?
   Gitlab::Database.main.disable_prepared_statements
 
@@ -7,6 +9,11 @@ if Gitlab::Database::LoadBalancing.enable?
     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)
+
   Gitlab::Database::LoadBalancing.configure_proxy
 
   # This needs to be executed after fork of clustered processes

lib/gitlab/database/load_balancing.rb

@@ -23,7 +23,7 @@ module Gitlab
 
       # The connection proxy to use for load balancing (if enabled).
       def self.proxy
-        unless @proxy
+        unless load_balancing_proxy = ActiveRecord::Base.load_balancing_proxy
           Gitlab::ErrorTracking.track_exception(
             ProxyNotConfiguredError.new(
               "Attempting to access the database load balancing proxy, but it wasn't configured.\n" \
@@ -31,7 +31,7 @@ module Gitlab
             ))
         end
 
-        @proxy
+        load_balancing_proxy
       end
 
       # Returns a Hash containing the load balancing configuration.
@@ -107,12 +107,7 @@ module Gitlab
 
       # Configures proxying of requests.
       def self.configure_proxy(proxy = ConnectionProxy.new(hosts))
-        @proxy = proxy
-
-        # 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(ActiveRecordProxy)
+        ActiveRecord::Base.load_balancing_proxy = proxy
       end
 
       def self.active_record_models
@@ -132,7 +127,7 @@ module Gitlab
       # recognize the connection, this method returns the primary role
       # directly. In future, we may need to check for other sources.
       def self.db_role_for_connection(connection)
-        return ROLE_PRIMARY if !enable? || @proxy.blank?
+        return ROLE_PRIMARY if !enable? || proxy.blank?
 
         proxy.load_balancer.db_role_for_connection(connection)
       end

lib/gitlab/database/load_balancing/active_record_proxy.rb

@@ -7,7 +7,7 @@ module Gitlab
       # "connection" method.
       module ActiveRecordProxy
         def connection
-          LoadBalancing.proxy
+          ::Gitlab::Database::LoadBalancing.proxy
         end
       end
     end

spec/initializers/lograge_spec.rb

@@ -120,7 +120,6 @@ RSpec.describe 'lograge', type: :request do
 
   context 'with a log subscriber' do
     include_context 'parsed logs'
-    include_context 'clear DB Load Balancing configuration'
 
     let(:subscriber) { Lograge::LogSubscribers::ActionController.new }
 

spec/lib/gitlab/checks/matching_merge_request_spec.rb

@@ -49,12 +49,11 @@ RSpec.describe Gitlab::Checks::MatchingMergeRequest do
       end
     end
 
-    context 'with load balancing enabled', :request_store, :redis do
+    context 'with load balancing enabled', :db_load_balancing do
       let(:session) { ::Gitlab::Database::LoadBalancing::Session.current }
       let(:all_caught_up) { true }
 
       before do
-        expect(::Gitlab::Database::LoadBalancing).to receive(:enable?).at_least(:once).and_return(true)
         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

spec/lib/gitlab/database/load_balancing_spec.rb

@@ -3,25 +3,32 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::LoadBalancing do
-  include_context 'clear DB Load Balancing configuration'
-
   before do
     stub_env('ENABLE_LOAD_BALANCING_FOR_FOSS', 'true')
   end
 
   describe '.proxy' do
+    before do
+      @previous_proxy = ActiveRecord::Base.load_balancing_proxy
+
+      ActiveRecord::Base.load_balancing_proxy = connection_proxy
+    end
+
+    after do
+      ActiveRecord::Base.load_balancing_proxy = @previous_proxy
+    end
+
     context 'when configured' do
-      before do
-        allow(ActiveRecord::Base.singleton_class).to receive(:prepend)
-        subject.configure_proxy
-      end
+      let(:connection_proxy) { double(:connection_proxy) }
 
       it 'returns the connection proxy' do
-        expect(subject.proxy).to be_an_instance_of(subject::ConnectionProxy)
+        expect(subject.proxy).to eq(connection_proxy)
       end
     end
 
     context 'when not configured' do
+      let(:connection_proxy) { nil }
+
       it 'returns nil' do
         expect(subject.proxy).to be_nil
       end
@@ -132,7 +139,6 @@ RSpec.describe Gitlab::Database::LoadBalancing do
 
   describe '.enable?' do
     before do
-      clear_load_balancing_configuration
       allow(described_class).to receive(:hosts).and_return(%w(foo))
     end
 
@@ -173,10 +179,6 @@ RSpec.describe Gitlab::Database::LoadBalancing do
   end
 
   describe '.configured?' do
-    before do
-      clear_load_balancing_configuration
-    end
-
     it 'returns true when Sidekiq is being used' do
       allow(described_class).to receive(:hosts).and_return(%w(foo))
       allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(true)
@@ -207,12 +209,12 @@ RSpec.describe Gitlab::Database::LoadBalancing do
 
   describe '.configure_proxy' do
     it 'configures the connection proxy' do
-      allow(ActiveRecord::Base.singleton_class).to receive(:prepend)
+      allow(ActiveRecord::Base).to receive(:load_balancing_proxy=)
 
       described_class.configure_proxy
 
-      expect(ActiveRecord::Base.singleton_class).to have_received(:prepend)
-        .with(Gitlab::Database::LoadBalancing::ActiveRecordProxy)
+      expect(ActiveRecord::Base).to have_received(:load_balancing_proxy=)
+        .with(Gitlab::Database::LoadBalancing::ConnectionProxy)
     end
   end
 
@@ -315,13 +317,9 @@ RSpec.describe Gitlab::Database::LoadBalancing do
       let(:load_balancer) { described_class::LoadBalancer.new(%w(foo)) }
 
       before do
-        allow(ActiveRecord::Base.singleton_class).to receive(:prepend)
-
         allow(described_class).to receive(:enable?).and_return(true)
         allow(described_class).to receive(:proxy).and_return(proxy)
         allow(proxy).to receive(:load_balancer).and_return(load_balancer)
-
-        subject.configure_proxy(proxy)
       end
 
       context 'when the load balancer returns :replica' do
@@ -366,7 +364,7 @@ RSpec.describe Gitlab::Database::LoadBalancing do
   # - 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', :delete do
+  describe 'LoadBalancing integration tests', :db_load_balancing, :delete do
     before(:all) do
       ActiveRecord::Schema.define do
         create_table :load_balancing_test, force: true do |t|
@@ -381,30 +379,14 @@ RSpec.describe Gitlab::Database::LoadBalancing do
       end
     end
 
-    shared_context 'LoadBalancing setup' do
-      let(:development_db_config) { ActiveRecord::Base.configurations.configs_for(env_name: 'development').first.configuration_hash }
-      let(:hosts) { [development_db_config[:host]] }
-      let(:model) do
-        Class.new(ApplicationRecord) do
-          self.table_name = "load_balancing_test"
-        end
+    let(:model) do
+      Class.new(ApplicationRecord) do
+        self.table_name = "load_balancing_test"
       end
+    end
 
-      before do
-        # Preloading testing class
-        model.singleton_class.prepend ::Gitlab::Database::LoadBalancing::ActiveRecordProxy
-
-        # Setup load balancing
-        clear_load_balancing_configuration
-        allow(ActiveRecord::Base.singleton_class).to receive(:prepend)
-        subject.configure_proxy(::Gitlab::Database::LoadBalancing::ConnectionProxy.new(hosts))
-
-        original_db_config = Gitlab::Database.main.config
-        modified_db_config = original_db_config.merge(load_balancing: { hosts: hosts })
-        allow(Gitlab::Database.main).to receive(:config).and_return(modified_db_config)
-
-        ::Gitlab::Database::LoadBalancing::Session.clear_session
-      end
+    before do
+      model.singleton_class.prepend ::Gitlab::Database::LoadBalancing::ActiveRecordProxy
     end
 
     where(:queries, :include_transaction, :expected_results) do
@@ -715,8 +697,6 @@ RSpec.describe Gitlab::Database::LoadBalancing do
     end
 
     with_them do
-      include_context 'LoadBalancing setup'
-
       it 'redirects queries to the right roles' do
         roles = []
 
@@ -785,8 +765,6 @@ RSpec.describe Gitlab::Database::LoadBalancing do
       end
 
       with_them do
-        include_context 'LoadBalancing setup'
-
         it 'redirects queries to the right roles' do
           roles = []
 
@@ -805,8 +783,6 @@ RSpec.describe Gitlab::Database::LoadBalancing do
     end
 
     context 'a write inside a transaction inside fallback_to_replicas_for_ambiguous_queries block' do
-      include_context 'LoadBalancing setup'
-
       it 'raises an exception' do
         expect do
           ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do

spec/lib/gitlab/instrumentation_helper_spec.rb

@@ -102,8 +102,6 @@ RSpec.describe Gitlab::InstrumentationHelper do
     end
 
     context 'when load balancing is enabled' do
-      include_context 'clear DB Load Balancing configuration'
-
       before do
         allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
       end

spec/lib/gitlab/sidekiq_logging/structured_logger_spec.rb

@@ -228,8 +228,6 @@ RSpec.describe Gitlab::SidekiqLogging::StructuredLogger do
     end
 
     context 'when the job performs database queries' do
-      include_context 'clear DB Load Balancing configuration'
-
       before do
         allow(Time).to receive(:now).and_return(timestamp)
         allow(Process).to receive(:clock_gettime).and_call_original
@@ -293,11 +291,7 @@ RSpec.describe Gitlab::SidekiqLogging::StructuredLogger do
         include_examples 'performs database queries'
       end
 
-      context 'when load balancing is enabled' do
-        before do
-          allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
-        end
-
+      context 'when load balancing is enabled', :db_load_balancing do
         let(:db_config_name) { ::Gitlab::Database.db_config_name(ApplicationRecord.connection) }
 
         let(:expected_end_payload_with_db) do

spec/lib/gitlab/sidekiq_middleware/server_metrics_spec.rb

@@ -236,7 +236,6 @@ RSpec.describe Gitlab::SidekiqMiddleware::ServerMetrics do
 
     include_context 'server metrics with mocked prometheus'
     include_context 'server metrics call'
-    include_context 'clear DB Load Balancing configuration'
 
     shared_context 'worker declaring data consistency' do
       let(:worker_class) { LBTestWorker }

spec/models/ci/build_spec.rb

@@ -344,12 +344,9 @@ RSpec.describe Ci::Build do
   end
 
   describe '#stick_build_if_status_changed' do
-    it 'sticks the build if the status changed' do
+    it 'sticks the build if the status changed', :db_load_balancing do
       job = create(:ci_build, :pending)
 
-      allow(Gitlab::Database::LoadBalancing).to receive(:enable?)
-        .and_return(true)
-
       expect(Gitlab::Database::LoadBalancing::Sticking).to receive(:stick)
         .with(:build, job.id)
 

spec/models/project_feature_usage_spec.rb

@@ -128,19 +128,15 @@ RSpec.describe ProjectFeatureUsage, type: :model do
   end
 
   context 'ProjectFeatureUsage with DB Load Balancing', :request_store do
-    include_context 'clear DB Load Balancing configuration'
-
     describe '#log_jira_dvcs_integration_usage' do
       let!(:project) { create(:project) }
 
       subject { project.feature_usage }
 
-      context 'database load balancing is configured' do
+      context 'database load balancing is configured', :db_load_balancing do
         before do
-          # Do not pollute AR for other tests, but rather simulate effect of configure_proxy.
-          allow(ActiveRecord::Base.singleton_class).to receive(:prepend)
-          ::Gitlab::Database::LoadBalancing.configure_proxy
           allow(ActiveRecord::Base).to receive(:connection).and_return(::Gitlab::Database::LoadBalancing.proxy)
+
           ::Gitlab::Database::LoadBalancing::Session.clear_session
         end
 

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' do
+      context 'checks database loadbalancing stickiness', :db_load_balancing do
         subject { described_class.new(shared_runner).execute }
 
         before do
@@ -22,9 +22,6 @@ module Ci
         end
 
         it 'result is valid if replica did caught-up' do
-          allow(Gitlab::Database::LoadBalancing).to receive(:enable?)
-            .and_return(true)
-
           expect(Gitlab::Database::LoadBalancing::Sticking).to receive(:all_caught_up?)
             .with(:runner, shared_runner.id) { true }
 
@@ -32,9 +29,6 @@ module Ci
         end
 
         it 'result is invalid if replica did not caught-up' do
-          allow(Gitlab::Database::LoadBalancing).to receive(:enable?)
-            .and_return(true)
-
           expect(Gitlab::Database::LoadBalancing::Sticking).to receive(:all_caught_up?)
             .with(:runner, shared_runner.id) { false }
 

spec/services/users/activity_service_spec.rb

@@ -85,19 +85,14 @@ RSpec.describe Users::ActivityService do
     end
   end
 
-  context 'with DB Load Balancing', :request_store, :redis, :clean_gitlab_redis_shared_state do
-    include_context 'clear DB Load Balancing configuration'
-
+  context 'with DB Load Balancing' do
     let(:user) { create(:user, last_activity_on: last_activity_on) }
 
     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' do
+      context 'database load balancing is configured', :db_load_balancing do
         before do
-          # Do not pollute AR for other tests, but rather simulate effect of configure_proxy.
-          allow(ActiveRecord::Base.singleton_class).to receive(:prepend)
-          ::Gitlab::Database::LoadBalancing.configure_proxy
           allow(ActiveRecord::Base).to receive(:connection).and_return(::Gitlab::Database::LoadBalancing.proxy)
         end
 

spec/support/database_load_balancing.rb

+# frozen_string_literal: true
+
+RSpec.configure do |config|
+  config.before(:each, :db_load_balancing) do
+    allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
+
+    proxy = ::Gitlab::Database::LoadBalancing::ConnectionProxy.new([Gitlab::Database.main.config['host']])
+
+    allow(ActiveRecord::Base).to receive(:load_balancing_proxy).and_return(proxy)
+
+    ::Gitlab::Database::LoadBalancing::Session.clear_session
+    redis_shared_state_cleanup!
+  end
+
+  config.after(:each, :db_load_balancing) do
+    ::Gitlab::Database::LoadBalancing::Session.clear_session
+    redis_shared_state_cleanup!
+  end
+end

spec/support/shared_contexts/load_balancing_configuration_shared_context.rb

-# frozen_string_literal: true
-
-RSpec.shared_context 'clear DB Load Balancing configuration' do
-  def clear_load_balancing_configuration
-    proxy = ::Gitlab::Database::LoadBalancing.instance_variable_get(:@proxy)
-    proxy.load_balancer.release_host if proxy
-    ::Gitlab::Database::LoadBalancing.instance_variable_set(:@proxy, nil)
-
-    ::Gitlab::Database::LoadBalancing::Session.clear_session
-  end
-
-  around do |example|
-    clear_load_balancing_configuration
-
-    example.run
-
-    clear_load_balancing_configuration
-  end
-end

spec/workers/authorized_project_update/user_refresh_from_replica_worker_spec.rb

@@ -44,11 +44,7 @@ RSpec.describe AuthorizedProjectUpdate::UserRefreshFromReplicaWorker do
       end
     end
 
-    context 'with load balancing enabled' do
-      before do
-        allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
-      end
-
+    context 'with load balancing enabled', :db_load_balancing do
       it 'reads from the replica database' do
         expect(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_replicas_for_read_queries).and_call_original
 

spec/workers/container_expiration_policy_worker_spec.rb

@@ -156,11 +156,7 @@ RSpec.describe ContainerExpirationPolicyWorker do
         subject
       end
 
-      context 'with load balancing enabled' do
-        before do
-          allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
-        end
-
+      context 'with load balancing enabled', :db_load_balancing 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