1 files changed, 157 insertions, 164 deletions

diff --git a/spec/lib/gitlab/database/load_balancing/sticking_spec.rb b/spec/lib/gitlab/database/load_balancing/sticking_spec.rb
index cf52e59db3a..8ceda52ee85 100644
--- a/spec/lib/gitlab/database/load_balancing/sticking_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/sticking_spec.rb
@@ -3,55 +3,82 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
+  let(:sticking) do
+    described_class.new(ActiveRecord::Base.connection.load_balancer)
+  end
+
   after do
     Gitlab::Database::LoadBalancing::Session.clear_session
   end
 
-  describe '.stick_if_necessary' do
-    context 'when sticking is disabled' do
-      it 'does not perform any sticking' do
-        expect(described_class).not_to receive(:stick)
+  describe '#stick_or_unstick_request' do
+    it 'sticks or unsticks a single object and updates the Rack environment' do
+      expect(sticking)
+        .to receive(:unstick_or_continue_sticking)
+        .with(:user, 42)
 
-        described_class.stick_if_necessary(:user, 42)
-      end
+      env = {}
+
+      sticking.stick_or_unstick_request(env, :user, 42)
+
+      expect(env[Gitlab::Database::LoadBalancing::RackMiddleware::STICK_OBJECT].to_a)
+        .to eq([[ActiveRecord::Base, :user, 42]])
     end
 
-    context 'when sticking is enabled' do
-      before do
-        allow(Gitlab::Database::LoadBalancing).to receive(:enable?)
-          .and_return(true)
-      end
+    it 'sticks or unsticks multiple objects and updates the Rack environment' do
+      expect(sticking)
+        .to receive(:unstick_or_continue_sticking)
+        .with(:user, 42)
+        .ordered
 
-      it 'does not stick if no write was performed' do
-        allow(Gitlab::Database::LoadBalancing::Session.current)
-          .to receive(:performed_write?)
-          .and_return(false)
+      expect(sticking)
+        .to receive(:unstick_or_continue_sticking)
+        .with(:runner, '123456789')
+        .ordered
 
-        expect(described_class).not_to receive(:stick)
+      env = {}
 
-        described_class.stick_if_necessary(:user, 42)
-      end
+      sticking.stick_or_unstick_request(env, :user, 42)
+      sticking.stick_or_unstick_request(env, :runner, '123456789')
 
-      it 'sticks to the primary if a write was performed' do
-        allow(Gitlab::Database::LoadBalancing::Session.current)
-          .to receive(:performed_write?)
-          .and_return(true)
+      expect(env[Gitlab::Database::LoadBalancing::RackMiddleware::STICK_OBJECT].to_a).to eq([
+        [ActiveRecord::Base, :user, 42],
+        [ActiveRecord::Base, :runner, '123456789']
+      ])
+    end
+  end
 
-        expect(described_class).to receive(:stick).with(:user, 42)
+  describe '#stick_if_necessary' do
+    it 'does not stick if no write was performed' do
+      allow(Gitlab::Database::LoadBalancing::Session.current)
+        .to receive(:performed_write?)
+        .and_return(false)
 
-        described_class.stick_if_necessary(:user, 42)
-      end
+      expect(sticking).not_to receive(:stick)
+
+      sticking.stick_if_necessary(:user, 42)
+    end
+
+    it 'sticks to the primary if a write was performed' do
+      allow(Gitlab::Database::LoadBalancing::Session.current)
+        .to receive(:performed_write?)
+        .and_return(true)
+
+      expect(sticking)
+        .to receive(:stick)
+        .with(:user, 42)
+
+      sticking.stick_if_necessary(:user, 42)
     end
   end
 
-  describe '.all_caught_up?' do
-    let(:lb) { double(:lb) }
+  describe '#all_caught_up?' do
+    let(:lb) { ActiveRecord::Base.connection.load_balancer }
     let(:last_write_location) { 'foo' }
 
     before do
-      allow(described_class).to receive(:load_balancer).and_return(lb)
-
-      allow(described_class).to receive(:last_write_location_for)
+      allow(sticking)
+        .to receive(:last_write_location_for)
         .with(:user, 42)
         .and_return(last_write_location)
     end
@@ -60,13 +87,9 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
       let(:last_write_location) { nil }
 
       it 'returns true' do
-        allow(described_class).to receive(:last_write_location_for)
-          .with(:user, 42)
-          .and_return(nil)
-
         expect(lb).not_to receive(:select_up_to_date_host)
 
-        expect(described_class.all_caught_up?(:user, 42)).to eq(true)
+        expect(sticking.all_caught_up?(:user, 42)).to eq(true)
       end
     end
 
@@ -76,9 +99,11 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
       end
 
       it 'returns true, and unsticks' do
-        expect(described_class).to receive(:unstick).with(:user, 42)
+        expect(sticking)
+          .to receive(:unstick)
+          .with(:user, 42)
 
-        expect(described_class.all_caught_up?(:user, 42)).to eq(true)
+        expect(sticking.all_caught_up?(:user, 42)).to eq(true)
       end
 
       it 'notifies with the proper event payload' do
@@ -87,7 +112,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
           .with('caught_up_replica_pick.load_balancing', { result: true })
           .and_call_original
 
-        described_class.all_caught_up?(:user, 42)
+        sticking.all_caught_up?(:user, 42)
       end
     end
 
@@ -97,7 +122,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
       end
 
       it 'returns false' do
-        expect(described_class.all_caught_up?(:user, 42)).to eq(false)
+        expect(sticking.all_caught_up?(:user, 42)).to eq(false)
       end
 
       it 'notifies with the proper event payload' do
@@ -106,42 +131,43 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
           .with('caught_up_replica_pick.load_balancing', { result: false })
           .and_call_original
 
-        described_class.all_caught_up?(:user, 42)
+        sticking.all_caught_up?(:user, 42)
       end
     end
   end
 
-  describe '.unstick_or_continue_sticking' do
-    let(:lb) { double(:lb) }
-
-    before do
-      allow(described_class).to receive(:load_balancer).and_return(lb)
-    end
+  describe '#unstick_or_continue_sticking' do
+    let(:lb) { ActiveRecord::Base.connection.load_balancer }
 
     it 'simply returns if no write location could be found' do
-      allow(described_class).to receive(:last_write_location_for)
+      allow(sticking)
+        .to receive(:last_write_location_for)
         .with(:user, 42)
         .and_return(nil)
 
       expect(lb).not_to receive(:select_up_to_date_host)
 
-      described_class.unstick_or_continue_sticking(:user, 42)
+      sticking.unstick_or_continue_sticking(:user, 42)
     end
 
     it 'unsticks if all secondaries have caught up' do
-      allow(described_class).to receive(:last_write_location_for)
+      allow(sticking)
+        .to receive(:last_write_location_for)
         .with(:user, 42)
         .and_return('foo')
 
       allow(lb).to receive(:select_up_to_date_host).with('foo').and_return(true)
 
-      expect(described_class).to receive(:unstick).with(:user, 42)
+      expect(sticking)
+        .to receive(:unstick)
+        .with(:user, 42)
 
-      described_class.unstick_or_continue_sticking(:user, 42)
+      sticking.unstick_or_continue_sticking(:user, 42)
     end
 
     it 'continues using the primary if the secondaries have not yet caught up' do
-      allow(described_class).to receive(:last_write_location_for)
+      allow(sticking)
+        .to receive(:last_write_location_for)
         .with(:user, 42)
         .and_return('foo')
 
@@ -150,184 +176,151 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
       expect(Gitlab::Database::LoadBalancing::Session.current)
         .to receive(:use_primary!)
 
-      described_class.unstick_or_continue_sticking(:user, 42)
+      sticking.unstick_or_continue_sticking(:user, 42)
     end
   end
 
   RSpec.shared_examples 'sticking' do
-    context 'when sticking is disabled' do
-      it 'does not perform any sticking', :aggregate_failures do
-        expect(described_class).not_to receive(:set_write_location_for)
-        expect(Gitlab::Database::LoadBalancing::Session.current).not_to receive(:use_primary!)
-
-        described_class.bulk_stick(:user, ids)
-      end
+    before do
+      allow(ActiveRecord::Base.connection.load_balancer)
+        .to receive(:primary_write_location)
+        .and_return('foo')
     end
 
-    context 'when sticking is enabled' do
-      before do
-        allow(Gitlab::Database::LoadBalancing).to receive(:configured?).and_return(true)
-
-        lb = double(:lb, primary_write_location: 'foo')
+    it 'sticks an entity to the primary', :aggregate_failures do
+      allow(ActiveRecord::Base.connection.load_balancer)
+        .to receive(:primary_only?)
+        .and_return(false)
 
-        allow(described_class).to receive(:load_balancer).and_return(lb)
+      ids.each do |id|
+        expect(sticking)
+          .to receive(:set_write_location_for)
+          .with(:user, id, 'foo')
       end
 
-      it 'sticks an entity to the primary', :aggregate_failures do
-        ids.each do |id|
-          expect(described_class).to receive(:set_write_location_for)
-                                       .with(:user, id, 'foo')
-        end
+      expect(Gitlab::Database::LoadBalancing::Session.current)
+        .to receive(:use_primary!)
 
-        expect(Gitlab::Database::LoadBalancing::Session.current)
-          .to receive(:use_primary!)
+      subject
+    end
 
-        subject
-      end
+    it 'does not update the write location when no replicas are used' do
+      expect(sticking).not_to receive(:set_write_location_for)
+
+      subject
     end
   end
 
-  describe '.stick' do
+  describe '#stick' do
     it_behaves_like 'sticking' do
       let(:ids) { [42] }
-      subject { described_class.stick(:user, ids.first) }
+      subject { sticking.stick(:user, ids.first) }
     end
   end
 
-  describe '.bulk_stick' do
+  describe '#bulk_stick' do
     it_behaves_like 'sticking' do
       let(:ids) { [42, 43] }
-      subject { described_class.bulk_stick(:user, ids) }
+      subject { sticking.bulk_stick(:user, ids) }
     end
   end
 
-  describe '.mark_primary_write_location' do
-    context 'when enabled' do
-      before do
-        allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
-        allow(Gitlab::Database::LoadBalancing).to receive(:configured?).and_return(true)
-      end
-
-      it 'updates the write location with the load balancer' do
-        lb = double(:lb, primary_write_location: 'foo')
+  describe '#mark_primary_write_location' do
+    it 'updates the write location with the load balancer' do
+      allow(ActiveRecord::Base.connection.load_balancer)
+        .to receive(:primary_write_location)
+        .and_return('foo')
 
-        allow(described_class).to receive(:load_balancer).and_return(lb)
+      allow(ActiveRecord::Base.connection.load_balancer)
+        .to receive(:primary_only?)
+        .and_return(false)
 
-        expect(described_class).to receive(:set_write_location_for)
-          .with(:user, 42, 'foo')
+      expect(sticking)
+        .to receive(:set_write_location_for)
+        .with(:user, 42, 'foo')
 
-        described_class.mark_primary_write_location(:user, 42)
-      end
+      sticking.mark_primary_write_location(:user, 42)
     end
 
-    context 'when load balancing is configured but not enabled' do
-      before do
-        allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(false)
-        allow(Gitlab::Database::LoadBalancing).to receive(:configured?).and_return(true)
-      end
-
-      it 'updates the write location with the main ActiveRecord connection' do
-        allow(described_class).to receive(:load_balancer).and_return(nil)
-        expect(ActiveRecord::Base).to receive(:connection).and_call_original
-        expect(described_class).to receive(:set_write_location_for)
-          .with(:user, 42, anything)
+    it 'does nothing when no replicas are used' do
+      expect(sticking).not_to receive(:set_write_location_for)
 
-        described_class.mark_primary_write_location(:user, 42)
-      end
-
-      context 'when write location is nil' do
-        before do
-          allow(Gitlab::Database.main).to receive(:get_write_location).and_return(nil)
-        end
+      sticking.mark_primary_write_location(:user, 42)
+    end
+  end
 
-        it 'does not update the write location' do
-          expect(described_class).not_to receive(:set_write_location_for)
+  describe '#unstick' do
+    it 'removes the sticking data from Redis' do
+      sticking.set_write_location_for(:user, 4, 'foo')
+      sticking.unstick(:user, 4)
 
-          described_class.mark_primary_write_location(:user, 42)
-        end
-      end
+      expect(sticking.last_write_location_for(:user, 4)).to be_nil
     end
 
-    context 'when load balancing is disabled' do
-      before do
-        allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(false)
-        allow(Gitlab::Database::LoadBalancing).to receive(:configured?).and_return(false)
+    it 'removes the old key' do
+      Gitlab::Redis::SharedState.with do |redis|
+        redis.set(sticking.send(:old_redis_key_for, :user, 4), 'foo', ex: 30)
       end
 
-      it 'updates the write location with the main ActiveRecord connection' do
-        expect(described_class).not_to receive(:set_write_location_for)
-
-        described_class.mark_primary_write_location(:user, 42)
-      end
+      sticking.unstick(:user, 4)
+      expect(sticking.last_write_location_for(:user, 4)).to be_nil
     end
   end
 
-  describe '.unstick' do
-    it 'removes the sticking data from Redis' do
-      described_class.set_write_location_for(:user, 4, 'foo')
-      described_class.unstick(:user, 4)
+  describe '#last_write_location_for' do
+    it 'returns the last WAL write location for a user' do
+      sticking.set_write_location_for(:user, 4, 'foo')
 
-      expect(described_class.last_write_location_for(:user, 4)).to be_nil
+      expect(sticking.last_write_location_for(:user, 4)).to eq('foo')
     end
-  end
 
-  describe '.last_write_location_for' do
-    it 'returns the last WAL write location for a user' do
-      described_class.set_write_location_for(:user, 4, 'foo')
+    it 'falls back to reading the old key' do
+      Gitlab::Redis::SharedState.with do |redis|
+        redis.set(sticking.send(:old_redis_key_for, :user, 4), 'foo', ex: 30)
+      end
 
-      expect(described_class.last_write_location_for(:user, 4)).to eq('foo')
+      expect(sticking.last_write_location_for(:user, 4)).to eq('foo')
     end
   end
 
-  describe '.redis_key_for' do
+  describe '#redis_key_for' do
     it 'returns a String' do
-      expect(described_class.redis_key_for(:user, 42))
-        .to eq('database-load-balancing/write-location/user/42')
+      expect(sticking.redis_key_for(:user, 42))
+        .to eq('database-load-balancing/write-location/main/user/42')
     end
   end
 
-  describe '.load_balancer' do
-    it 'returns a the load balancer' do
-      proxy = double(:proxy)
-
-      expect(Gitlab::Database::LoadBalancing).to receive(:proxy)
-        .and_return(proxy)
-
-      expect(proxy).to receive(:load_balancer)
-
-      described_class.load_balancer
-    end
-  end
-
-  describe '.select_caught_up_replicas' do
-    let(:lb) { double(:lb) }
-
-    before do
-      allow(described_class).to receive(:load_balancer).and_return(lb)
-    end
+  describe '#select_caught_up_replicas' do
+    let(:lb) { ActiveRecord::Base.connection.load_balancer }
 
     context 'with no write location' do
       before do
-        allow(described_class).to receive(:last_write_location_for)
-          .with(:project, 42).and_return(nil)
+        allow(sticking)
+          .to receive(:last_write_location_for)
+          .with(:project, 42)
+          .and_return(nil)
       end
 
       it 'returns false and does not try to find caught up hosts' do
         expect(lb).not_to receive(:select_up_to_date_host)
-        expect(described_class.select_caught_up_replicas(:project, 42)).to be false
+        expect(sticking.select_caught_up_replicas(:project, 42)).to be false
       end
     end
 
     context 'with write location' do
       before do
-        allow(described_class).to receive(:last_write_location_for)
-          .with(:project, 42).and_return('foo')
+        allow(sticking)
+          .to receive(:last_write_location_for)
+          .with(:project, 42)
+          .and_return('foo')
       end
 
       it 'returns true, selects hosts, and unsticks if any secondary has caught up' do
         expect(lb).to receive(:select_up_to_date_host).and_return(true)
-        expect(described_class).to receive(:unstick).with(:project, 42)
-        expect(described_class.select_caught_up_replicas(:project, 42)).to be true
+        expect(sticking)
+          .to receive(:unstick)
+          .with(:project, 42)
+        expect(sticking.select_caught_up_replicas(:project, 42)).to be true
       end
     end
   end