Store the load balancer proxy in ActiveRecord

Currently we are storing this in a class instance variable in
Gitlab::Database::LoadBalancing. This class is auto-loaded which
means it is also unloaded when a Ruby change is detected.

This moves this to a class instance variable in ActiveRecord::Base

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

@@ -54,6 +54,7 @@ RSpec.describe Gitlab::Checks::MatchingMergeRequest do
       let(:all_caught_up) { true }
 
       before do
+        allow(ActiveRecord::Base).to receive(:load_balancing_proxy)
         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)
 

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
@@ -395,9 +393,9 @@ RSpec.describe Gitlab::Database::LoadBalancing do
         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))
+        allow(ActiveRecord::Base).to receive(:load_balancing_proxy).and_return(
+          ::Gitlab::Database::LoadBalancing::ConnectionProxy.new(hosts)
+        )
 
         original_db_config = Gitlab::Database.main.config
         modified_db_config = original_db_config.merge(load_balancing: { hosts: hosts })

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
@@ -296,6 +294,7 @@ RSpec.describe Gitlab::SidekiqLogging::StructuredLogger do
       context 'when load balancing is enabled' do
         before do
           allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
+          allow(ActiveRecord::Base).to receive(:load_balancing_proxy)
         end
 
         let(:db_config_name) { ::Gitlab::Database.db_config_name(ApplicationRecord.connection) }

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

@@ -347,6 +347,7 @@ RSpec.describe Ci::Build do
     it 'sticks the build if the status changed' do
       job = create(:ci_build, :pending)
 
+      allow(ActiveRecord::Base).to receive(:load_balancing_proxy)
       allow(Gitlab::Database::LoadBalancing).to receive(:enable?)
         .and_return(true)
 

spec/models/project_feature_usage_spec.rb

@@ -128,8 +128,6 @@ 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) }
 
@@ -138,9 +136,10 @@ RSpec.describe ProjectFeatureUsage, type: :model do
       context 'database load balancing is configured' 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)
+          allow(ActiveRecord::Base).to receive(:load_balancing_proxy=)
+          proxy = ::Gitlab::Database::LoadBalancing.configure_proxy
+          allow(ActiveRecord::Base).to receive(:connection).and_return(proxy)
+
           ::Gitlab::Database::LoadBalancing::Session.clear_session
         end
 

spec/services/ci/register_job_service_spec.rb

@@ -19,12 +19,12 @@ module Ci
 
         before do
           project.update!(shared_runners_enabled: false)
+
+          allow(ActiveRecord::Base).to receive(:load_balancing_proxy)
+          allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
         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 +32,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

@@ -86,8 +86,6 @@ RSpec.describe Users::ActivityService do
   end
 
   context 'with DB Load Balancing', :request_store, :redis, :clean_gitlab_redis_shared_state do
-    include_context 'clear DB Load Balancing configuration'
-
     let(:user) { create(:user, last_activity_on: last_activity_on) }
 
     context 'when last activity is in the past' do
@@ -96,9 +94,9 @@ RSpec.describe Users::ActivityService do
       context 'database load balancing is configured' 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)
+          allow(ActiveRecord::Base).to receive(:load_balancing_proxy=)
+          proxy = ::Gitlab::Database::LoadBalancing.configure_proxy
+          allow(ActiveRecord::Base).to receive(:connection).and_return(proxy)
         end
 
         let(:service) do

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

@@ -47,6 +47,7 @@ RSpec.describe AuthorizedProjectUpdate::UserRefreshFromReplicaWorker do
     context 'with load balancing enabled' do
       before do
         allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
+        allow(ActiveRecord::Base).to receive(:load_balancing_proxy)
       end
 
       it 'reads from the replica database' do

spec/workers/container_expiration_policy_worker_spec.rb

@@ -159,6 +159,7 @@ RSpec.describe ContainerExpirationPolicyWorker do
       context 'with load balancing enabled' do
         before do
           allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
+          allow(ActiveRecord::Base).to receive(:load_balancing_proxy)
         end
 
         it 'reads the counts from the replica' do