The `use_model_load_balancing` results in a wrong sticking context used

This solves a subtle bug with `sticking` when `use_model_load_balancing`
is used.

The problem was that for `Web/Sidekiq`:

1. If prior request would have `use_model_load_balancing=false` it would
   see load balancers to be `main/main`. This would result in sticking context
   to include only `main`.

2. If next request would evaluate `use_model_load_balancing=true` it would
   see load balancers for `main/ci`, but would only for wal locations see `main=lag`.
   As a result we would not check replication lag against `ci`.

3. This result in a situation that if `ci` is way behind `main`, the 2.
   would not understand the state of `ci`. Thus would consider `ci` to be up-to date
   and stick to `main`.

This commit fixes that:

1. If `ci:` is configured we always store the `ci:` in replication lag. This makes
   us the `wal_locations` (and `RackMiddleware` sticking context) to always include
   `main+ci` in all cases.

2. This results would behave (as for number of requests): we are reading primary / replica
   location from all databases.

3. Since we always have `ci` even if `use_model_load_balancing=false` we
   can properly evaluate state of replicas.

config/initializers/load_balancing.rb

@@ -20,7 +20,7 @@ Gitlab::Database::LoadBalancing.base_models.each do |model|
   Gitlab::Cluster::LifecycleEvents.on_before_fork do
     # When forking, we don't want to wait until the connections aren't in use
     # any more, as this could delay the boot cycle.
-    model.connection.load_balancer.disconnect!(timeout: 0)
+    model.load_balancer.disconnect!(timeout: 0)
   end
 
   # Service discovery only needs to run in the worker processes, as the main one

lib/gitlab/database/load_balancing.rb

@@ -26,7 +26,7 @@ module Gitlab
         return to_enum(__method__) unless block_given?
 
         base_models.each do |model|
-          yield model.connection.load_balancer
+          yield model.load_balancer
         end
       end
 

lib/gitlab/database/load_balancing/connection_proxy.rb

@@ -13,7 +13,15 @@ module Gitlab
         WriteInsideReadOnlyTransactionError = Class.new(StandardError)
         READ_ONLY_TRANSACTION_KEY = :load_balacing_read_only_transaction
 
-        attr_reader :load_balancer
+        # The load balancer is intentionally not exposed since the returned instance
+        # might be different `model.connection.load_balancer` vs `model.load_balancer`
+        #
+        # The used `model.connection` is dependent on `use_model_load_balancing`.
+        # See more in: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/73949.
+        #
+        # Always use `model.load_balancer` or `model.sticking`.
+        #
+        # attr_reader :load_balancer
 
         # These methods perform writes after which we need to stick to the
         # primary.

lib/gitlab/database/load_balancing/setup.rb

@@ -40,6 +40,7 @@ module Gitlab
         def setup_connection_proxy
           # We just use a simple `class_attribute` here so we don't need to
           # inject any modules and/or expose unnecessary methods.
+          setup_class_attribute(:load_balancer, load_balancer)
           setup_class_attribute(:connection, ConnectionProxy.new(load_balancer))
           setup_class_attribute(:sticking, Sticking.new(load_balancer))
         end
@@ -107,10 +108,6 @@ module Gitlab
           def connection
             use_model_load_balancing? ? super : ActiveRecord::Base.connection
           end
-
-          def sticking
-            use_model_load_balancing? ? super : ActiveRecord::Base.sticking
-          end
           # rubocop:enable Database/MultipleDatabases
         end
       end

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

@@ -3,12 +3,9 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
-  let(:proxy) do
-    config = Gitlab::Database::LoadBalancing::Configuration
-      .new(ActiveRecord::Base)
-
-    described_class.new(Gitlab::Database::LoadBalancing::LoadBalancer.new(config))
-  end
+  let(:config) { Gitlab::Database::LoadBalancing::Configuration.new(ActiveRecord::Base) }
+  let(:load_balancer) { Gitlab::Database::LoadBalancing::LoadBalancer.new(config) }
+  let(:proxy) { described_class.new(load_balancer) }
 
   describe '#select' do
     it 'performs a read' do
@@ -143,9 +140,9 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
 
       context 'with a read query' do
         it 'runs the transaction and any nested queries on the replica' do
-          expect(proxy.load_balancer).to receive(:read)
+          expect(load_balancer).to receive(:read)
             .twice.and_yield(replica)
-          expect(proxy.load_balancer).not_to receive(:read_write)
+          expect(load_balancer).not_to receive(:read_write)
           expect(session).not_to receive(:write!)
 
           proxy.transaction { proxy.select('true') }
@@ -154,8 +151,8 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
 
       context 'with a write query' do
         it 'raises an exception' do
-          allow(proxy.load_balancer).to receive(:read).and_yield(replica)
-          allow(proxy.load_balancer).to receive(:read_write).and_yield(replica)
+          allow(load_balancer).to receive(:read).and_yield(replica)
+          allow(load_balancer).to receive(:read_write).and_yield(replica)
 
           expect do
             proxy.transaction { proxy.insert('something') }
@@ -178,9 +175,9 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
 
       context 'with a read query' do
         it 'runs the transaction and any nested queries on the primary and stick to it' do
-          expect(proxy.load_balancer).to receive(:read_write)
+          expect(load_balancer).to receive(:read_write)
             .twice.and_yield(primary)
-          expect(proxy.load_balancer).not_to receive(:read)
+          expect(load_balancer).not_to receive(:read)
           expect(session).to receive(:write!)
 
           proxy.transaction { proxy.select('true') }
@@ -189,9 +186,9 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
 
       context 'with a write query' do
         it 'runs the transaction and any nested queries on the primary and stick to it' do
-          expect(proxy.load_balancer).to receive(:read_write)
+          expect(load_balancer).to receive(:read_write)
             .twice.and_yield(primary)
-          expect(proxy.load_balancer).not_to receive(:read)
+          expect(load_balancer).not_to receive(:read)
           expect(session).to receive(:write!).twice
 
           proxy.transaction { proxy.insert('something') }
@@ -209,7 +206,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
     end
 
     it 'properly forwards keyword arguments' do
-      allow(proxy.load_balancer).to receive(:read_write)
+      allow(load_balancer).to receive(:read_write)
 
       expect(proxy).to receive(:write_using_load_balancer).and_call_original
 
@@ -234,7 +231,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
       end
 
       it 'properly forwards keyword arguments' do
-        allow(proxy.load_balancer).to receive(:read)
+        allow(load_balancer).to receive(:read)
 
         expect(proxy).to receive(:read_using_load_balancer).and_call_original
 
@@ -259,7 +256,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
         allow(session).to receive(:use_replicas_for_read_queries?).and_return(false)
 
         expect(connection).to receive(:foo).with('foo')
-        expect(proxy.load_balancer).to receive(:read).and_yield(connection)
+        expect(load_balancer).to receive(:read).and_yield(connection)
 
         proxy.read_using_load_balancer(:foo, 'foo')
       end
@@ -271,7 +268,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
         allow(session).to receive(:use_replicas_for_read_queries?).and_return(true)
 
         expect(connection).to receive(:foo).with('foo')
-        expect(proxy.load_balancer).to receive(:read).and_yield(connection)
+        expect(load_balancer).to receive(:read).and_yield(connection)
 
         proxy.read_using_load_balancer(:foo, 'foo')
       end
@@ -283,7 +280,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
         allow(session).to receive(:use_replicas_for_read_queries?).and_return(true)
 
         expect(connection).to receive(:foo).with('foo')
-        expect(proxy.load_balancer).to receive(:read).and_yield(connection)
+        expect(load_balancer).to receive(:read).and_yield(connection)
 
         proxy.read_using_load_balancer(:foo, 'foo')
       end
@@ -296,7 +293,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
 
         expect(connection).to receive(:foo).with('foo')
 
-        expect(proxy.load_balancer).to receive(:read_write)
+        expect(load_balancer).to receive(:read_write)
           .and_yield(connection)
 
         proxy.read_using_load_balancer(:foo, 'foo')
@@ -314,7 +311,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
     end
 
     it 'uses but does not stick to the primary' do
-      expect(proxy.load_balancer).to receive(:read_write).and_yield(connection)
+      expect(load_balancer).to receive(:read_write).and_yield(connection)
       expect(connection).to receive(:foo).with('foo')
       expect(session).not_to receive(:write!)
 

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

@@ -61,7 +61,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do
 
       setup.setup_connection_proxy
 
-      expect(model.connection.load_balancer).to eq(lb)
+      expect(model.load_balancer).to eq(lb)
       expect(model.sticking)
         .to be_an_instance_of(Gitlab::Database::LoadBalancing::Sticking)
     end
@@ -265,17 +265,19 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do
         stub_env('GITLAB_USE_MODEL_LOAD_BALANCING', env_GITLAB_USE_MODEL_LOAD_BALANCING)
         stub_env('GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci', env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci)
         stub_feature_flags(use_model_load_balancing: ff_use_model_load_balancing)
-      end
 
-      it 'results match expectations' do
-        result = models.transform_values do |model|
         # Make load balancer to force init with a dedicated replicas connections
+        models.each do |_, model|
           described_class.new(model).tap do |subject|
             subject.configuration.hosts = [subject.configuration.replica_db_config.host]
             subject.setup
           end
+        end
+      end
 
-          load_balancer = model.connection.load_balancer
+      it 'results match expectations' do
+        result = models.transform_values do |model|
+          load_balancer = model.connection.instance_variable_get(:@load_balancer)
 
           {
             read: load_balancer.read { |connection| connection.pool.db_config.name },
@@ -285,6 +287,13 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do
 
         expect(result).to eq(expectations)
       end
+
+      it 'does return load_balancer assigned to a given connection' do
+        models.each do |name, model|
+          expect(model.load_balancer.name).to eq(name)
+          expect(model.sticking.instance_variable_get(:@load_balancer)).to eq(model.load_balancer)
+        end
+      end
     end
   end
 end

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

@@ -64,7 +64,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware, :clean_
       let(:wal_locations) { { Gitlab::Database::MAIN_DATABASE_NAME.to_sym => location } }
 
       it 'does not stick to the primary', :aggregate_failures do
-        expect(ActiveRecord::Base.connection.load_balancer)
+        expect(ActiveRecord::Base.load_balancer)
           .to receive(:select_up_to_date_host)
           .with(location)
           .and_return(true)
@@ -107,7 +107,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware, :clean_
         let(:job) { { 'job_id' => 'a180b47c-3fd6-41b8-81e9-34da61c3400e', 'dedup_wal_locations' => wal_locations } }
 
         before do
-          allow(ActiveRecord::Base.connection.load_balancer)
+          allow(ActiveRecord::Base.load_balancer)
             .to receive(:select_up_to_date_host)
             .with(wal_locations[:main])
             .and_return(true)
@@ -120,7 +120,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware, :clean_
         let(:job) { { 'job_id' => 'a180b47c-3fd6-41b8-81e9-34da61c3400e', 'database_write_location' => '0/D525E3A8' } }
 
         before do
-          allow(ActiveRecord::Base.connection.load_balancer)
+          allow(ActiveRecord::Base.load_balancer)
             .to receive(:select_up_to_date_host)
             .with('0/D525E3A8')
             .and_return(true)

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

@@ -4,7 +4,7 @@ require 'spec_helper'
 
 RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
   let(:sticking) do
-    described_class.new(ActiveRecord::Base.connection.load_balancer)
+    described_class.new(ActiveRecord::Base.load_balancer)
   end
 
   after do
@@ -73,7 +73,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
   end
 
   describe '#all_caught_up?' do
-    let(:lb) { ActiveRecord::Base.connection.load_balancer }
+    let(:lb) { ActiveRecord::Base.load_balancer }
     let(:last_write_location) { 'foo' }
 
     before do
@@ -137,7 +137,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
   end
 
   describe '#unstick_or_continue_sticking' do
-    let(:lb) { ActiveRecord::Base.connection.load_balancer }
+    let(:lb) { ActiveRecord::Base.load_balancer }
 
     it 'simply returns if no write location could be found' do
       allow(sticking)
@@ -182,13 +182,13 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
 
   RSpec.shared_examples 'sticking' do
     before do
-      allow(ActiveRecord::Base.connection.load_balancer)
+      allow(ActiveRecord::Base.load_balancer)
         .to receive(:primary_write_location)
         .and_return('foo')
     end
 
     it 'sticks an entity to the primary', :aggregate_failures do
-      allow(ActiveRecord::Base.connection.load_balancer)
+      allow(ActiveRecord::Base.load_balancer)
         .to receive(:primary_only?)
         .and_return(false)
 
@@ -227,11 +227,11 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
 
   describe '#mark_primary_write_location' do
     it 'updates the write location with the load balancer' do
-      allow(ActiveRecord::Base.connection.load_balancer)
+      allow(ActiveRecord::Base.load_balancer)
         .to receive(:primary_write_location)
         .and_return('foo')
 
-      allow(ActiveRecord::Base.connection.load_balancer)
+      allow(ActiveRecord::Base.load_balancer)
         .to receive(:primary_only?)
         .and_return(false)
 
@@ -291,7 +291,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
   end
 
   describe '#select_caught_up_replicas' do
-    let(:lb) { ActiveRecord::Base.connection.load_balancer }
+    let(:lb) { ActiveRecord::Base.load_balancer }
 
     context 'with no write location' do
       before do

spec/lib/gitlab/database/load_balancing_spec.rb

@@ -76,7 +76,7 @@ RSpec.describe Gitlab::Database::LoadBalancing do
 
       context 'when a read connection is used' do
         it 'returns :replica' do
-          proxy.load_balancer.read do |connection|
+          load_balancer.read do |connection|
             expect(described_class.db_role_for_connection(connection)).to eq(:replica)
           end
         end
@@ -84,7 +84,7 @@ RSpec.describe Gitlab::Database::LoadBalancing do
 
       context 'when a read_write connection is used' do
         it 'returns :primary' do
-          proxy.load_balancer.read_write do |connection|
+          load_balancer.read_write do |connection|
             expect(described_class.db_role_for_connection(connection)).to eq(:primary)
           end
         end

spec/lib/gitlab/database/query_analyzer_spec.rb

@@ -64,7 +64,7 @@ RSpec.describe Gitlab::Database::QueryAnalyzer do
     end
 
     def process_sql(sql)
-      ApplicationRecord.connection.load_balancer.read_write do |connection|
+      ApplicationRecord.load_balancer.read_write do |connection|
         described_class.new.send(:process_sql, sql, connection)
       end
     end

spec/lib/gitlab/database_spec.rb

@@ -205,7 +205,7 @@ RSpec.describe Gitlab::Database do
 
     context 'when replicas are configured', :database_replica do
       it 'returns the name for a replica' do
-        replica = ActiveRecord::Base.connection.load_balancer.host
+        replica = ActiveRecord::Base.load_balancer.host
 
         expect(described_class.db_config_name(replica)).to eq('main_replica')
       end

spec/support/database_load_balancing.rb

@@ -2,17 +2,17 @@
 
 RSpec.configure do |config|
   config.around(:each, :database_replica) do |example|
-    old_proxies = []
+    old_proxies = {}
 
     Gitlab::Database::LoadBalancing.base_models.each do |model|
+      old_proxies[model] = [model.load_balancer, model.connection, model.sticking]
+
       config = Gitlab::Database::LoadBalancing::Configuration
         .new(model, [model.connection_db_config.configuration_hash[:host]])
-      lb = Gitlab::Database::LoadBalancing::LoadBalancer.new(config)
-
-      old_proxies << [model, model.connection]
 
-      model.connection =
-        Gitlab::Database::LoadBalancing::ConnectionProxy.new(lb)
+      model.load_balancer = Gitlab::Database::LoadBalancing::LoadBalancer.new(config)
+      model.sticking = Gitlab::Database::LoadBalancing::Sticking.new(model.load_balancer)
+      model.connection = Gitlab::Database::LoadBalancing::ConnectionProxy.new(model.load_balancer)
     end
 
     Gitlab::Database::LoadBalancing::Session.clear_session
@@ -23,8 +23,8 @@ RSpec.configure do |config|
     Gitlab::Database::LoadBalancing::Session.clear_session
     redis_shared_state_cleanup!
 
-    old_proxies.each do |(model, proxy)|
-      model.connection = proxy
+    old_proxies.each do |model, proxy|
+      model.load_balancer, model.connection, model.sticking = proxy
     end
   end
 end