Merge branch 'sh-fix-sidekiq-mark-write-location' into 'master'

Fix intermittent merge errors due to replication lag

See merge request gitlab-org/gitlab!44045

ee/lib/gitlab/database/load_balancing.rb

@@ -84,8 +84,19 @@ module Gitlab
 
       # Returns true if load balancing is to be enabled.
       def self.enable?
-        return false unless ::License.feature_available?(:db_load_balancing)
         return false if program_name == 'rake' || Gitlab::Runtime.sidekiq?
+        return false unless self.configured?
+
+        true
+      end
+
+      # Returns true if load balancing has been configured. Since
+      # Sidekiq does not currently use load balancing, we
+      # may want Web application servers to detect replication lag by
+      # posting the write location of the database if load balancing is
+      # configured.
+      def self.configured?
+        return false unless ::License.feature_available?(:db_load_balancing)
 
         hosts.any? || service_discovery_enabled?
       end

ee/lib/gitlab/database/load_balancing/load_balancer.rb

@@ -105,18 +105,14 @@ module Gitlab
 
         # Returns the transaction write location of the primary.
         def primary_write_location
-          read_write do |connection|
-            row = connection
-              .select_all("SELECT pg_current_wal_insert_lsn()::text AS location")
-              .first
+          location = read_write do |connection|
+            ::Gitlab::Database.get_write_location(connection)
+          end
+
+          return location if location
 
-            if row
-              row['location']
-            else
           raise 'Failed to determine the write location of the primary database'
         end
-          end
-        end
 
         # Returns true if all hosts have caught up to the given transaction
         # write location.

ee/lib/gitlab/database/load_balancing/sticking.rb

@@ -49,18 +49,24 @@ module Gitlab
         def self.stick(namespace, id)
           return unless LoadBalancing.enable?
 
-          mark_primary_write_location_helper(namespace, id)
+          mark_primary_write_location(namespace, id)
           Session.current.use_primary!
         end
 
         def self.mark_primary_write_location(namespace, id)
-          return unless LoadBalancing.enable?
+          return unless LoadBalancing.configured?
 
-          mark_primary_write_location_helper(namespace, id)
+          # Load balancing could be enabled for the Web application server,
+          # but it's not activated for Sidekiq. We should update Redis with
+          # the write location just in case load balancing is being used.
+          location =
+            if LoadBalancing.enable?
+              load_balancer.primary_write_location
+            else
+              Gitlab::Database.get_write_location(ActiveRecord::Base.connection)
             end
 
-        def self.mark_primary_write_location_helper(namespace, id)
-          location = load_balancer.primary_write_location
+          return if location.blank?
 
           set_write_location_for(namespace, id, location)
         end

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

@@ -140,6 +140,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
       it 'sticks an entity to the primary' do
         allow(Gitlab::Database::LoadBalancing).to receive(:enable?)
           .and_return(true)
+        allow(Gitlab::Database::LoadBalancing).to receive(:configured?).and_return(true)
 
         lb = double(:lb, primary_write_location: 'foo')
 
@@ -157,30 +158,60 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
   end
 
   describe '.mark_primary_write_location' do
-    context 'when load balancing is disabled' do
+    context 'when enabled' do
       before do
-        allow(Gitlab::Database::LoadBalancing).to receive(:enable?)
-          .and_return(false)
+        allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
+        allow(Gitlab::Database::LoadBalancing).to receive(:configured?).and_return(true)
       end
 
-      it 'does not attempt to write' do
-        expect(described_class).not_to receive(:set_write_location_for)
+      it 'updates the write location with the load balancer' do
+        lb = double(:lb, primary_write_location: 'foo')
+
+        allow(described_class).to receive(:load_balancer).and_return(lb)
+
+        expect(described_class).to receive(:set_write_location_for)
+          .with(:user, 42, 'foo')
 
         described_class.mark_primary_write_location(:user, 42)
       end
     end
 
-    context 'when load balancing is enabled' do
-      it 'updates the write location' do
-        allow(Gitlab::Database::LoadBalancing).to receive(:enable?)
-          .and_return(true)
+    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
 
-        lb = double(:lb, primary_write_location: 'foo')
+      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)
 
-        allow(described_class).to receive(:load_balancer).and_return(lb)
+        described_class.mark_primary_write_location(:user, 42)
+      end
 
-        expect(described_class).to receive(:set_write_location_for)
-          .with(:user, 42, 'foo')
+      context 'when write location is nil' do
+        before do
+          allow(Gitlab::Database).to receive(:get_write_location).and_return(nil)
+        end
+
+        it 'does not update the write location' do
+          expect(described_class).not_to receive(:set_write_location_for)
+
+          described_class.mark_primary_write_location(:user, 42)
+        end
+      end
+    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)
+      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

ee/spec/lib/gitlab/database/load_balancing_spec.rb

@@ -199,6 +199,49 @@ RSpec.describe Gitlab::Database::LoadBalancing do
     end
   end
 
+  describe '.configured?' do
+    let!(:license) { create(:license, plan: ::License::PREMIUM_PLAN) }
+
+    it 'returns true when Sidekiq is being used' do
+      allow(described_class).to receive(:hosts).and_return(%w(foo))
+      allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(true)
+      expect(described_class).not_to receive(:program_name)
+
+      expect(described_class.configured?).to eq(true)
+    end
+
+    it 'returns true when service discovery is enabled in Sidekiq' do
+      allow(described_class).to receive(:hosts).and_return([])
+      allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(true)
+
+      allow(described_class)
+        .to receive(:service_discovery_enabled?)
+        .and_return(true)
+
+      expect(described_class.configured?).to eq(true)
+    end
+
+    it 'returns false when neither service discovery nor hosts are configured' do
+      allow(described_class).to receive(:hosts).and_return([])
+
+      allow(described_class)
+        .to receive(:service_discovery_enabled?)
+        .and_return(false)
+
+      expect(described_class.configured?).to eq(false)
+    end
+
+    context 'without a license' do
+      before do
+        License.destroy_all # rubocop: disable Cop/DestroyAll
+      end
+
+      it 'is not configured' do
+        expect(described_class.configured?).to eq(false)
+      end
+    end
+  end
+
   describe '.program_name' do
     it 'returns a String' do
       expect(described_class.program_name).to be_an_instance_of(String)

lib/gitlab/database.rb

@@ -250,6 +250,14 @@ module Gitlab
       false
     end
 
+    def self.get_write_location(ar_connection)
+      row = ar_connection
+        .select_all("SELECT pg_current_wal_insert_lsn()::text AS location")
+        .first
+
+      row['location'] if row
+    end
+
     private_class_method :database_version
 
     def self.add_post_migrate_path_to_rails(force: false)

spec/lib/gitlab/database_spec.rb

@@ -352,6 +352,20 @@ RSpec.describe Gitlab::Database do
     end
   end
 
+  describe '.get_write_location' do
+    it 'returns a string' do
+      connection = ActiveRecord::Base.connection
+
+      expect(described_class.get_write_location(connection)).to be_a(String)
+    end
+
+    it 'returns nil if there are no results' do
+      connection = double(select_all: [])
+
+      expect(described_class.get_write_location(connection)).to be_nil
+    end
+  end
+
   describe '#true_value' do
     it 'returns correct value' do
       expect(described_class.true_value).to eq "'t'"