Merge branch 'load-balancer-connection-arguments' into 'master'

Refactor LoadBalancing::LoadBalancer to support scoping it to different models

See merge request gitlab-org/gitlab!67057

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) }