Use `GITLAB_LOAD_BALANCING_REUSE_PRIMARY_*` to allow primary re-use

The purpose of this feature is to allow Phase 3 of decomposition rollout
as described in https://gitlab.com/groups/gitlab-org/-/epics/6160#progress.

The idea is that we configure a separate DB (CI) with replicas, but due
to transaction visibility we need to share primary connection to `main`
database.

This adds a temporary feature enabled via `FF` to allow this type
of behavior change.

lib/gitlab/database/load_balancing/configuration.rb

@@ -5,8 +5,6 @@ module Gitlab
     module LoadBalancing
       # Configuration settings for a single LoadBalancer instance.
       class Configuration
-        attr_reader :connection_specification_name, :connection_db_config
-
         attr_accessor :hosts, :max_replication_difference,
                       :max_replication_lag_time, :replica_check_interval,
                       :service_discovery
@@ -43,12 +41,12 @@ module Gitlab
             end
           end
 
+          config.reuse_primary_connection!
+
           config
         end
 
         def initialize(model, hosts = [])
-          @connection_specification_name = model.connection_specification_name
-          @connection_db_config = model.connection_db_config
           @max_replication_difference = 8.megabytes
           @max_replication_lag_time = 60.0
           @replica_check_interval = 60.0
@@ -63,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
@@ -90,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.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.connection_specification_name,
+            @configuration.primary_connection_specification_name,
             role: ActiveRecord::Base.writing_role,
             shard: ActiveRecord::Base.default_shard
           ) || raise(::ActiveRecord::ConnectionNotEstablished)

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

@@ -3,20 +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,
-      connection_specification_name: ActiveRecord::Base.connection_specification_name
-    )
-  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)
 
@@ -108,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
@@ -183,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

@@ -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/setup_spec.rb

@@ -23,10 +23,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do
         env_name: 'test',
         name: 'main'
       )
-      model = double(:model,
-        connection_db_config: config,
-        connection_specification_name: ActiveRecord::Base.connection_specification_name
-      )
+      model = double(:model, connection_db_config: config)
 
       expect(ActiveRecord::DatabaseConfigurations::HashConfig)
         .to receive(:new)

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