Replace traversal_ids lock mechanism

Lock all roots when moving namespace
Lock on root ancestor instead of the results of the CTE

app/models/namespace/traversal_hierarchy.rb

@@ -34,20 +34,23 @@ class Namespace
       sql = """
             UPDATE namespaces
             SET traversal_ids = cte.traversal_ids
-            FROM (#{recursive_traversal_ids(lock: true)}) as cte
+            FROM (#{recursive_traversal_ids}) as cte
             WHERE namespaces.id = cte.id
               AND namespaces.traversal_ids <> cte.traversal_ids
             """
-      Namespace.connection.exec_query(sql)
+      Namespace.transaction do
+        @root.lock!
+        Namespace.connection.exec_query(sql)
+      end
     rescue ActiveRecord::Deadlocked
       db_deadlock_counter.increment(source: 'Namespace#sync_traversal_ids!')
       raise
     end
 
     # Identify all incorrect traversal_ids in the current namespace hierarchy.
-    def incorrect_traversal_ids(lock: false)
+    def incorrect_traversal_ids
       Namespace
-        .joins("INNER JOIN (#{recursive_traversal_ids(lock: lock)}) as cte ON namespaces.id = cte.id")
+        .joins("INNER JOIN (#{recursive_traversal_ids}) as cte ON namespaces.id = cte.id")
         .where('namespaces.traversal_ids <> cte.traversal_ids')
     end
 
@@ -58,13 +61,10 @@ class Namespace
     #
     # Note that the traversal_ids represent a calculated traversal path for the
     # namespace and not the value stored within the traversal_ids attribute.
-    #
-    # Optionally locked with FOR UPDATE to ensure isolation between concurrent
-    # updates of the heirarchy.
-    def recursive_traversal_ids(lock: false)
+    def recursive_traversal_ids
       root_id = Integer(@root.id)
 
-      sql = <<~SQL
+      <<~SQL
       WITH RECURSIVE cte(id, traversal_ids, cycle) AS (
         VALUES(#{root_id}, ARRAY[#{root_id}], false)
       UNION ALL
@@ -74,10 +74,6 @@ class Namespace
       )
       SELECT id, traversal_ids FROM cte
       SQL
-
-      sql += ' FOR UPDATE' if lock
-
-      sql
     end
 
     # This is essentially Namespace#root_ancestor which will soon be rewritten

app/models/namespaces/traversal/linear.rb

@@ -41,6 +41,7 @@ module Namespaces
       UnboundedSearch = Class.new(StandardError)
 
       included do
+        before_update :lock_both_roots, if: -> { sync_traversal_ids? && parent_id_changed? }
         after_create :sync_traversal_ids, if: -> { sync_traversal_ids? }
         after_update :sync_traversal_ids, if: -> { sync_traversal_ids? && saved_change_to_parent_id? }
 
@@ -77,6 +78,23 @@ module Namespaces
         Namespace::TraversalHierarchy.for_namespace(root_ancestor).sync_traversal_ids!
       end
 
+      # Lock the root of the hierarchy we just left, and lock the root of the hierarchy
+      # we just joined. In most cases the two hierarchies will be the same.
+      def lock_both_roots
+        parent_ids = [
+          parent_id_was || self.id,
+          parent_id || self.id
+        ].compact
+
+        roots = Gitlab::ObjectHierarchy
+          .new(Namespace.where(id: parent_ids))
+          .base_and_ancestors
+          .reorder(nil)
+          .where(parent_id: nil)
+
+        Namespace.lock.select(:id).where(id: roots).order(id: :asc).load
+      end
+
       # Make sure we drop the STI `type = 'Group'` condition for better performance.
       # Logically equivalent so long as hierarchies remain homogeneous.
       def without_sti_condition

ee/spec/lib/ee/gitlab/usage_data_non_sql_metrics_spec.rb

@@ -15,7 +15,7 @@ RSpec.describe Gitlab::UsageDataNonSqlMetrics do
         described_class.uncached_data
       end
 
-      expect(recorder.count).to eq(50)
+      expect(recorder.count).to eq(52)
     end
   end
 end

spec/models/group_spec.rb

@@ -395,18 +395,94 @@ RSpec.describe Group do
       end
     end
 
-    context 'assigning a new parent' do
-      let!(:old_parent) { create(:group) }
-      let!(:new_parent) { create(:group) }
+    context 'assign a new parent' do
       let!(:group) { create(:group, parent: old_parent) }
+      let(:recorded_queries) { ActiveRecord::QueryRecorder.new }
+
+      subject do
+        recorded_queries.record do
+          group.update(parent: new_parent)
+        end
+      end
 
       before do
-        group.update(parent: new_parent)
+        subject
         reload_models(old_parent, new_parent, group)
       end
 
-      it 'updates traversal_ids' do
-        expect(group.traversal_ids).to eq [new_parent.id, group.id]
+      context 'within the same hierarchy' do
+        let!(:root) { create(:group).reload }
+        let!(:old_parent) { create(:group, parent: root) }
+        let!(:new_parent) { create(:group, parent: root) }
+
+        it 'updates traversal_ids' do
+          expect(group.traversal_ids).to eq [root.id, new_parent.id, group.id]
+        end
+
+        it_behaves_like 'hierarchy with traversal_ids'
+        it_behaves_like 'locked row' do
+          let(:row) { root }
+        end
+      end
+
+      context 'to another hierarchy' do
+        let!(:old_parent) { create(:group) }
+        let!(:new_parent) { create(:group) }
+        let!(:group) { create(:group, parent: old_parent) }
+
+        it 'updates traversal_ids' do
+          expect(group.traversal_ids).to eq [new_parent.id, group.id]
+        end
+
+        it_behaves_like 'locked rows' do
+          let(:rows) { [old_parent, new_parent] }
+        end
+
+        context 'old hierarchy' do
+          let(:root) { old_parent.root_ancestor }
+
+          it_behaves_like 'hierarchy with traversal_ids'
+        end
+
+        context 'new hierarchy' do
+          let(:root) { new_parent.root_ancestor }
+
+          it_behaves_like 'hierarchy with traversal_ids'
+        end
+      end
+
+      context 'from being a root ancestor' do
+        let!(:old_parent) { nil }
+        let!(:new_parent) { create(:group) }
+
+        it 'updates traversal_ids' do
+          expect(group.traversal_ids).to eq [new_parent.id, group.id]
+        end
+
+        it_behaves_like 'locked rows' do
+          let(:rows) { [group, new_parent] }
+        end
+
+        it_behaves_like 'hierarchy with traversal_ids' do
+          let(:root) { new_parent }
+        end
+      end
+
+      context 'to being a root ancestor' do
+        let!(:old_parent) { create(:group) }
+        let!(:new_parent) { nil }
+
+        it 'updates traversal_ids' do
+          expect(group.traversal_ids).to eq [group.id]
+        end
+
+        it_behaves_like 'locked rows' do
+          let(:rows) { [old_parent, group] }
+        end
+
+        it_behaves_like 'hierarchy with traversal_ids' do
+          let(:root) { group }
+        end
       end
     end
 

spec/models/namespace/traversal_hierarchy_spec.rb

@@ -43,16 +43,6 @@ RSpec.describe Namespace::TraversalHierarchy, type: :model do
     end
   end
 
-  shared_examples 'locked update query' do
-    it 'locks query with FOR UPDATE' do
-      qr = ActiveRecord::QueryRecorder.new do
-        subject
-      end
-      expect(qr.count).to eq 1
-      expect(qr.log.first).to match /FOR UPDATE/
-    end
-  end
-
   describe '#incorrect_traversal_ids' do
     let!(:hierarchy) { described_class.new(root) }
 
@@ -63,12 +53,6 @@ RSpec.describe Namespace::TraversalHierarchy, type: :model do
     end
 
     it { is_expected.to match_array Namespace.all }
-
-    context 'when lock is true' do
-      subject { hierarchy.incorrect_traversal_ids(lock: true).load }
-
-      it_behaves_like 'locked update query'
-    end
   end
 
   describe '#sync_traversal_ids!' do
@@ -79,14 +63,18 @@ RSpec.describe Namespace::TraversalHierarchy, type: :model do
     it { expect(hierarchy.incorrect_traversal_ids).to be_empty }
 
     it_behaves_like 'hierarchy with traversal_ids'
-    it_behaves_like 'locked update query'
+    it_behaves_like 'locked row' do
+      let(:recorded_queries) { ActiveRecord::QueryRecorder.new }
+      let(:row) { root }
 
-    context 'when deadlocked' do
       before do
-        connection_double = double(:connection)
+        recorded_queries.record { subject }
+      end
+    end
 
-        allow(Namespace).to receive(:connection).and_return(connection_double)
-        allow(connection_double).to receive(:exec_query) { raise ActiveRecord::Deadlocked }
+    context 'when deadlocked' do
+      before do
+        allow(root).to receive(:lock!) { raise ActiveRecord::Deadlocked }
       end
 
       it { expect { subject }.to raise_error(ActiveRecord::Deadlocked) }

spec/support/helpers/query_recorder.rb

@@ -13,6 +13,10 @@ module ActiveRecord
       @skip_cached = skip_cached
       @query_recorder_debug = ENV['QUERY_RECORDER_DEBUG'] || query_recorder_debug
       @log_file = log_file
+      record(&block) if block_given?
+    end
+
+    def record(&block)
       # force replacement of bind parameters to give tests the ability to check for ids
       ActiveRecord::Base.connection.unprepared_statement do
         ActiveSupport::Notifications.subscribed(method(:callback), 'sql.active_record', &block)

spec/support/helpers/reload_helpers.rb

@@ -2,7 +2,7 @@
 
 module ReloadHelpers
   def reload_models(*models)
-    models.map(&:reload)
+    models.compact.map(&:reload)
   end
 
   def subject_and_reload(*models)

spec/support/shared_examples/row_lock_shared_examples.rb

+# frozen_string_literal: true
+
+# Ensure that a SQL command to lock this row(s) was requested.
+# Ensure a transaction also occurred.
+# Be careful! This form of spec is not foolproof, but better than nothing.
+
+RSpec.shared_examples 'locked row' do
+  it "has locked row" do
+    table_name = row.class.table_name
+    ids_regex = /SELECT.*FROM.*#{table_name}.*"#{table_name}"."id" = #{row.id}.+FOR UPDATE/m
+
+    expect(recorded_queries.log).to include a_string_matching 'SAVEPOINT'
+    expect(recorded_queries.log).to include a_string_matching ids_regex
+  end
+end
+
+RSpec.shared_examples 'locked rows' do
+  it "has locked rows" do
+    table_name = rows.first.class.table_name
+
+    row_ids = rows.map(&:id).join(', ')
+    ids_regex = /SELECT.+FROM.+"#{table_name}".+"#{table_name}"."id" IN \(#{row_ids}\).+FOR UPDATE/m
+
+    expect(recorded_queries.log).to include a_string_matching 'SAVEPOINT'
+    expect(recorded_queries.log).to include a_string_matching ids_regex
+  end
+end