Prepare LoadBalancer for multiple databases

This refactors Gitlab::Database::LoadBalancing::LoadBalancer so it can
support the use of multiple databases (each with their own load
balancer) in the future.

Instead of directly depending on ActiveRecord::Base, LoadBalancer now
takes a ActiveRecord model. It will then scope its work to that model,
including the data that it caches in RequestStore. This ensures two
different LoadBalancer instances won't conflict with each other.

The new setup of retrieving connections is such that we can ensure we
always have the right connection (e.g. a writable connection when we
need one).

See https://gitlab.com/gitlab-org/gitlab/-/issues/331776 for more
information.
Co-Authored-By: Thong Kuah <tkuah@gitlab.com>

lib/gitlab/database/load_balancing/load_balancer.rb

@@ -7,19 +7,17 @@ module Gitlab
       #
       # Each host in the load balancer uses the same credentials as the primary
       # database.
-      #
-      # This class *requires* that `ActiveRecord::Base.retrieve_connection`
-      # always returns a connection to the primary.
       class LoadBalancer
         CACHE_KEY = :gitlab_load_balancer_host
 
         attr_reader :host_list
 
         # hosts - The hostnames/addresses of the additional databases.
-        def initialize(hosts = [])
+        def initialize(hosts = [], model = ActiveRecord::Base)
           @host_list = HostList.new(hosts.map { |addr| Host.new(addr, self) })
           @connection_db_roles = {}.compare_by_identity
           @connection_db_roles_count = {}.compare_by_identity
+          @model = model
         end
 
         # Yields a connection that can be used for reads.
@@ -94,7 +92,7 @@ module Gitlab
           # Instead of immediately grinding to a halt we'll retry the operation
           # a few times.
           retry_with_backoff do
-            connection = ActiveRecord::Base.retrieve_connection
+            connection = pool.connection
             track_connection_role(connection, ROLE_PRIMARY)
 
             yield connection
@@ -109,7 +107,7 @@ module Gitlab
         def db_role_for_connection(connection)
           return @connection_db_roles[connection] if @connection_db_roles[connection]
           return ROLE_REPLICA if @host_list.manage_pool?(connection.pool)
-          return ROLE_PRIMARY if connection.pool == ActiveRecord::Base.connection_pool
+          return ROLE_PRIMARY if connection.pool == pool
         end
 
         # Returns a host to use for queries.
@@ -117,21 +115,21 @@ module Gitlab
         # Hosts are scoped per thread so that multiple threads don't
         # accidentally re-use the same host + connection.
         def host
-          RequestStore[CACHE_KEY] ||= @host_list.next
+          request_cache[CACHE_KEY] ||= @host_list.next
         end
 
         # Releases the host and connection for the current thread.
         def release_host
-          if host = RequestStore[CACHE_KEY]
+          if host = request_cache[CACHE_KEY]
             host.disable_query_cache!
             host.release_connection
           end
 
-          RequestStore.delete(CACHE_KEY)
+          request_cache.delete(CACHE_KEY)
         end
 
         def release_primary_connection
-          ActiveRecord::Base.connection_pool.release_connection
+          pool.release_connection
         end
 
         # Returns the transaction write location of the primary.
@@ -152,7 +150,7 @@ module Gitlab
 
           return false unless host
 
-          RequestStore[CACHE_KEY] = host
+          request_cache[CACHE_KEY] = host
 
           true
         end
@@ -209,6 +207,17 @@ module Gitlab
 
         private
 
+        # ActiveRecord::ConnectionAdapters::ConnectionHandler handles fetching,
+        # and caching for connections pools for each "connection", so we
+        # leverage that.
+        def pool
+          ActiveRecord::Base.connection_handler.retrieve_connection_pool(
+            @model.connection_specification_name,
+            role: ActiveRecord::Base.writing_role,
+            shard: ActiveRecord::Base.default_shard
+          )
+        end
+
         def ensure_caching!
           host.enable_query_cache! unless host.query_cache_enabled
         end
@@ -228,6 +237,11 @@ module Gitlab
             @connection_db_roles_count.delete(connection)
           end
         end
+
+        def request_cache
+          base = RequestStore[:gitlab_load_balancer] ||= {}
+          base[pool] ||= {}
+        end
       end
     end
   end

spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb

@@ -7,6 +7,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
   let(:conflict_error) { Class.new(RuntimeError) }
 
   let(:lb) { described_class.new(%w(localhost localhost)) }
+  let(:request_cache) { lb.send(:request_cache) }
 
   before do
     allow(Gitlab::Database.main).to receive(:create_connection_pool)
@@ -123,8 +124,9 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
 
   describe '#read_write' do
     it 'yields a connection for a write' do
-      expect { |b| lb.read_write(&b) }
-        .to yield_with_args(ActiveRecord::Base.retrieve_connection)
+      connection = ActiveRecord::Base.connection_pool.connection
+
+      expect { |b| lb.read_write(&b) }.to yield_with_args(connection)
     end
 
     it 'uses a retry with exponential backoffs' do
@@ -260,13 +262,24 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
     end
 
     it 'stores the host in a thread-local variable' do
-      RequestStore.delete(described_class::CACHE_KEY)
+      request_cache.delete(described_class::CACHE_KEY)
 
       expect(lb.host_list).to receive(:next).once.and_call_original
 
       lb.host
       lb.host
     end
+
+    it 'does not create conflicts with other load balancers when caching hosts' do
+      lb1 = described_class.new(%w(localhost localhost), ActiveRecord::Base)
+      lb2 = described_class.new(%w(localhost localhost), Ci::CiDatabaseRecord)
+
+      host1 = lb1.host
+      host2 = lb2.host
+
+      expect(lb1.send(:request_cache)[described_class::CACHE_KEY]).to eq(host1)
+      expect(lb2.send(:request_cache)[described_class::CACHE_KEY]).to eq(host2)
+    end
   end
 
   describe '#release_host' do
@@ -277,7 +290,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
 
       lb.release_host
 
-      expect(RequestStore[described_class::CACHE_KEY]).to be_nil
+      expect(request_cache[described_class::CACHE_KEY]).to be_nil
     end
   end
 
@@ -415,7 +428,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
   describe '#select_up_to_date_host' do
     let(:location) { 'AB/12345'}
     let(:hosts) { lb.host_list.hosts }
-    let(:set_host) { RequestStore[described_class::CACHE_KEY] }
+    let(:set_host) { request_cache[described_class::CACHE_KEY] }
 
     subject { lb.select_up_to_date_host(location) }