Merge branch 'lb-reuses-similar-connections' into 'master'

Make `LoadBalancer` to re-use very similar connection

See merge request gitlab-org/gitlab!72698

lib/gitlab/database/load_balancing/configuration.rb

@@ -7,7 +7,7 @@ module Gitlab
       class Configuration
         attr_accessor :hosts, :max_replication_difference,
                       :max_replication_lag_time, :replica_check_interval,
-                      :service_discovery, :model
+                      :service_discovery
 
         # Creates a configuration object for the given ActiveRecord model.
         def self.for_model(model)
@@ -41,6 +41,8 @@ module Gitlab
             end
           end
 
+          config.reuse_primary_connection!
+
           config
         end
 
@@ -59,6 +61,24 @@ module Gitlab
             disconnect_timeout: 120,
             use_tcp: false
           }
+
+          # Temporary model for GITLAB_LOAD_BALANCING_REUSE_PRIMARY_
+          # To be removed with FF
+          @primary_model = nil
+        end
+
+        def db_config_name
+          @model.connection_db_config.name.to_sym
+        end
+
+        # With connection re-use the primary connection can be overwritten
+        # to be used from different model
+        def primary_connection_specification_name
+          (@primary_model || @model).connection_specification_name
+        end
+
+        def replica_db_config
+          @model.connection_db_config
         end
 
         def pool_size
@@ -86,6 +106,30 @@ module Gitlab
         def service_discovery_enabled?
           service_discovery[:record].present?
         end
+
+        # TODO: This is temporary code to allow re-use of primary connection
+        # if the two connections are pointing to the same host. This is needed
+        # to properly support transaction visibility.
+        #
+        # This behavior is required to support [Phase 3](https://gitlab.com/groups/gitlab-org/-/epics/6160#progress).
+        # This method is meant to be removed as soon as it is finished.
+        #
+        # The remapping is done as-is:
+        #   export GITLAB_LOAD_BALANCING_REUSE_PRIMARY_<name-of-connection>=<new-name-of-connection>
+        #
+        # Ex.:
+        #   export GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci=main
+        #
+        def reuse_primary_connection!
+          new_connection = ENV["GITLAB_LOAD_BALANCING_REUSE_PRIMARY_#{db_config_name}"]
+          return unless new_connection.present?
+
+          @primary_model = Gitlab::Database.database_base_models[new_connection.to_sym]
+
+          unless @primary_model
+            raise "Invalid value for 'GITLAB_LOAD_BALANCING_REUSE_PRIMARY_#{db_config_name}=#{new_connection}'"
+          end
+        end
       end
     end
   end

lib/gitlab/database/load_balancing/load_balancer.rb

@@ -12,7 +12,7 @@ module Gitlab
 
         REPLICA_SUFFIX = '_replica'
 
-        attr_reader :name, :host_list, :configuration
+        attr_reader :host_list, :configuration
 
         # configuration - An instance of `LoadBalancing::Configuration` that
         #                 contains the configuration details (such as the hosts)
@@ -26,8 +26,10 @@ module Gitlab
             else
               HostList.new(configuration.hosts.map { |addr| Host.new(addr, self) })
             end
+        end
 
-          @name = @configuration.model.connection_db_config.name.to_sym
+        def name
+          @configuration.db_config_name
         end
 
         def primary_only?
@@ -227,7 +229,7 @@ module Gitlab
         # host - An optional host name to use instead of the default one.
         # port - An optional port to connect to.
         def create_replica_connection_pool(pool_size, host = nil, port = nil)
-          db_config = pool.db_config
+          db_config = @configuration.replica_db_config
 
           env_config = db_config.configuration_hash.dup
           env_config[:pool] = pool_size
@@ -252,7 +254,7 @@ module Gitlab
         # leverage that.
         def pool
           ActiveRecord::Base.connection_handler.retrieve_connection_pool(
-            @configuration.model.connection_specification_name,
+            @configuration.primary_connection_specification_name,
             role: ActiveRecord::Base.writing_role,
             shard: ActiveRecord::Base.default_shard
           ) || raise(::ActiveRecord::ConnectionNotEstablished)

lib/gitlab/database/load_balancing/rack_middleware.rb

@@ -38,8 +38,8 @@ module Gitlab
         def unstick_or_continue_sticking(env)
           namespaces_and_ids = sticking_namespaces(env)
 
-          namespaces_and_ids.each do |(model, namespace, id)|
-            model.sticking.unstick_or_continue_sticking(namespace, id)
+          namespaces_and_ids.each do |(sticking, namespace, id)|
+            sticking.unstick_or_continue_sticking(namespace, id)
           end
         end
 
@@ -47,8 +47,8 @@ module Gitlab
         def stick_if_necessary(env)
           namespaces_and_ids = sticking_namespaces(env)
 
-          namespaces_and_ids.each do |model, namespace, id|
-            model.sticking.stick_if_necessary(namespace, id)
+          namespaces_and_ids.each do |sticking, namespace, id|
+            sticking.stick_if_necessary(namespace, id)
           end
         end
 
@@ -74,7 +74,7 @@ module Gitlab
             # models that support load balancing. In the future (if we
             # determined this to be OK) we may be able to relax this.
             ::Gitlab::Database::LoadBalancing.base_models.map do |model|
-              [model, :user, warden.user.id]
+              [model.sticking, :user, warden.user.id]
             end
           elsif env[STICK_OBJECT].present?
             env[STICK_OBJECT].to_a

lib/gitlab/database/load_balancing/sticking.rb

@@ -12,7 +12,6 @@ module Gitlab
 
         def initialize(load_balancer)
           @load_balancer = load_balancer
-          @model = load_balancer.configuration.model
         end
 
         # Unsticks or continues sticking the current request.
@@ -28,7 +27,7 @@ module Gitlab
           unstick_or_continue_sticking(namespace, id)
 
           env[RackMiddleware::STICK_OBJECT] ||= Set.new
-          env[RackMiddleware::STICK_OBJECT] << [@model, namespace, id]
+          env[RackMiddleware::STICK_OBJECT] << [self, namespace, id]
         end
 
         # Sticks to the primary if a write was performed.

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

@@ -3,17 +3,12 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::LoadBalancing::Configuration do
-  let(:model) do
-    config = ActiveRecord::DatabaseConfigurations::HashConfig
-      .new('main', 'test', configuration_hash)
-
-    double(:model, connection_db_config: config)
-  end
+  let(:configuration_hash) { {} }
+  let(:db_config) { ActiveRecord::DatabaseConfigurations::HashConfig.new('test', 'ci', configuration_hash) }
+  let(:model) { double(:model, connection_db_config: db_config) }
 
   describe '.for_model' do
     context 'when load balancing is not configured' do
-      let(:configuration_hash) { {} }
-
       it 'uses the default settings' do
         config = described_class.for_model(model)
 
@@ -105,6 +100,14 @@ RSpec.describe Gitlab::Database::LoadBalancing::Configuration do
         expect(config.pool_size).to eq(4)
       end
     end
+
+    it 'calls reuse_primary_connection!' do
+      expect_next_instance_of(described_class) do |subject|
+        expect(subject).to receive(:reuse_primary_connection!).and_call_original
+      end
+
+      described_class.for_model(model)
+    end
   end
 
   describe '#load_balancing_enabled?' do
@@ -180,4 +183,58 @@ RSpec.describe Gitlab::Database::LoadBalancing::Configuration do
       end
     end
   end
+
+  describe '#db_config_name' do
+    let(:config) { described_class.new(model) }
+
+    subject { config.db_config_name }
+
+    it 'returns connection name as symbol' do
+      is_expected.to eq(:ci)
+    end
+  end
+
+  describe '#replica_db_config' do
+    let(:model) { double(:model, connection_db_config: db_config, connection_specification_name: 'Ci::CiDatabaseRecord') }
+    let(:config) { described_class.for_model(model) }
+
+    it 'returns exactly db_config' do
+      expect(config.replica_db_config).to eq(db_config)
+    end
+
+    context 'when GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci=main' do
+      it 'does not change replica_db_config' do
+        stub_env('GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci', 'main')
+
+        expect(config.replica_db_config).to eq(db_config)
+      end
+    end
+  end
+
+  describe 'reuse_primary_connection!' do
+    let(:model) { double(:model, connection_db_config: db_config, connection_specification_name: 'Ci::CiDatabaseRecord') }
+    let(:config) { described_class.for_model(model) }
+
+    context 'when GITLAB_LOAD_BALANCING_REUSE_PRIMARY_* not configured' do
+      it 'the primary connection uses default specification' do
+        expect(config.primary_connection_specification_name).to eq('Ci::CiDatabaseRecord')
+      end
+    end
+
+    context 'when GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci=main' do
+      it 'the primary connection uses main connection' do
+        stub_env('GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci', 'main')
+
+        expect(config.primary_connection_specification_name).to eq('ActiveRecord::Base')
+      end
+    end
+
+    context 'when GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci=unknown' do
+      it 'raises exception' do
+        stub_env('GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci', 'unknown')
+
+        expect { config.reuse_primary_connection! }.to raise_error /Invalid value for/
+      end
+    end
+  end
 end

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

@@ -4,10 +4,11 @@ require 'spec_helper'
 
 RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
   let(:conflict_error) { Class.new(RuntimeError) }
-  let(:db_host) { ActiveRecord::Base.connection_pool.db_config.host }
+  let(:model) { ActiveRecord::Base }
+  let(:db_host) { model.connection_pool.db_config.host }
   let(:config) do
     Gitlab::Database::LoadBalancing::Configuration
-      .new(ActiveRecord::Base, [db_host, db_host])
+      .new(model, [db_host, db_host])
   end
 
   let(:lb) { described_class.new(config) }
@@ -452,4 +453,51 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
       expect(lb.send(:get_write_location, double(select_all: []))).to be_nil
     end
   end
+
+  describe 'primary connection re-use', :reestablished_active_record_base do
+    let(:model) { Ci::CiDatabaseRecord }
+
+    before do
+      # fake additional Database
+      model.establish_connection(
+        ActiveRecord::DatabaseConfigurations::HashConfig.new(Rails.env, 'ci', ActiveRecord::Base.connection_db_config.configuration_hash)
+      )
+    end
+
+    describe '#read' do
+      it 'returns ci replica connection' do
+        expect { |b| lb.read(&b) }.to yield_with_args do |args|
+          expect(args.pool.db_config.name).to eq('ci_replica')
+        end
+      end
+
+      context 'when GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci=main' do
+        it 'returns ci replica connection' do
+          stub_env('GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci', 'main')
+
+          expect { |b| lb.read(&b) }.to yield_with_args do |args|
+            expect(args.pool.db_config.name).to eq('ci_replica')
+          end
+        end
+      end
+    end
+
+    describe '#read_write' do
+      it 'returns Ci::CiDatabaseRecord connection' do
+        expect { |b| lb.read_write(&b) }.to yield_with_args do |args|
+          expect(args.pool.db_config.name).to eq('ci')
+        end
+      end
+
+      context 'when GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci=main' do
+        it 'returns ActiveRecord::Base connection' do
+          stub_env('GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci', 'main')
+
+          expect { |b| lb.read_write(&b) }.to yield_with_args do |args|
+            expect(args.pool.db_config.name).to eq('main')
+          end
+        end
+      end
+    end
+  end
 end

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

@@ -6,12 +6,12 @@ RSpec.describe Gitlab::Database::LoadBalancing::RackMiddleware, :redis do
   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([[ActiveRecord::Base, :user, 42]]) }
+  let(:single_sticking_object) { Set.new([[ActiveRecord::Base.sticking, :user, 42]]) }
   let(:multiple_sticking_objects) do
     Set.new([
-      [ActiveRecord::Base, :user, 42],
-      [ActiveRecord::Base, :runner, '123456789'],
-      [ActiveRecord::Base, :runner, '1234']
+      [ActiveRecord::Base.sticking, :user, 42],
+      [ActiveRecord::Base.sticking, :runner, '123456789'],
+      [ActiveRecord::Base.sticking, :runner, '1234']
     ])
   end
 
@@ -162,7 +162,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::RackMiddleware, :redis 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]
+          [model.sticking, :user, 42]
         end
 
         expect(middleware.sticking_namespaces(env)).to eq(ids)
@@ -181,9 +181,9 @@ RSpec.describe Gitlab::Database::LoadBalancing::RackMiddleware, :redis do
         env = { described_class::STICK_OBJECT => multiple_sticking_objects }
 
         expect(middleware.sticking_namespaces(env)).to eq([
-          [ActiveRecord::Base, :user, 42],
-          [ActiveRecord::Base, :runner, '123456789'],
-          [ActiveRecord::Base, :runner, '1234']
+          [ActiveRecord::Base.sticking, :user, 42],
+          [ActiveRecord::Base.sticking, :runner, '123456789'],
+          [ActiveRecord::Base.sticking, :runner, '1234']
         ])
       end
     end

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

@@ -22,7 +22,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
       sticking.stick_or_unstick_request(env, :user, 42)
 
       expect(env[Gitlab::Database::LoadBalancing::RackMiddleware::STICK_OBJECT].to_a)
-        .to eq([[ActiveRecord::Base, :user, 42]])
+        .to eq([[sticking, :user, 42]])
     end
 
     it 'sticks or unsticks multiple objects and updates the Rack environment' do
@@ -42,8 +42,8 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
       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']
+        [sticking, :user, 42],
+        [sticking, :runner, '123456789']
       ])
     end
   end