Linear version of Namespace#self_and_descendants

Replace the recursive version of Namespace#self_and_descendants with a
faster linear version that utilizes the namespaces.traversal_ids column.

app/models/namespaces/traversal/linear.rb

@@ -38,15 +38,31 @@ module Namespaces
     module Linear
       extend ActiveSupport::Concern
 
+      UnboundedSearch = Class.new(StandardError)
+
       included do
         after_create :sync_traversal_ids, if: -> { sync_traversal_ids? }
         after_update :sync_traversal_ids, if: -> { sync_traversal_ids? && saved_change_to_parent_id? }
+
+        scope :traversal_ids_contains, ->(ids) { where("traversal_ids @> (?)", ids) }
       end
 
       def sync_traversal_ids?
         Feature.enabled?(:sync_traversal_ids, root_ancestor, default_enabled: :yaml)
       end
 
+      def use_traversal_ids?
+        Feature.enabled?(:use_traversal_ids, root_ancestor, default_enabled: :yaml)
+      end
+
+      def self_and_descendants
+        if use_traversal_ids?
+          lineage(self)
+        else
+          super
+        end
+      end
+
       private
 
       # Update the traversal_ids for the full hierarchy.
@@ -58,6 +74,38 @@ module Namespaces
 
         Namespace::TraversalHierarchy.for_namespace(root_ancestor).sync_traversal_ids!
       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
+        self.class.unscope(where: :type)
+      end
+
+      # Search this namespace's lineage. Bound inclusively by top node.
+      def lineage(top)
+        raise UnboundedSearch.new('Must bound search by a top') unless top
+
+        without_sti_condition
+          .traversal_ids_contains(latest_traversal_ids(top))
+      end
+
+      # traversal_ids are a cached value.
+      #
+      # The traversal_ids value in a loaded object can become stale when compared
+      # to the database value. For example, if you load a hierarchy and then move
+      # a group, any previously loaded descendant objects will have out of date
+      # traversal_ids.
+      #
+      # To solve this problem, we never depend on the object's traversal_ids
+      # value. We always query the database first with a sub-select for the
+      # latest traversal_ids.
+      #
+      # Note that ActiveRecord will cache query results. You can avoid this by
+      # using `Model.uncached { ... }`
+      def latest_traversal_ids(namespace = self)
+        without_sti_condition.where('id = (?)', namespace)
+                .select('traversal_ids as latest_traversal_ids')
+      end
     end
   end
 end

changelogs/unreleased/301142-replace-namespace-self_and_descendants-with-linear-version.yml

+---
+title: Linear version of Namespace#self_and_descendants
+merge_request: 56296
+author:
+type: performance

config/feature_flags/development/use_traversal_ids.yml

+---
+name: use_traversal_ids
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/56296
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/321948
+milestone: '13.11'
+type: development
+group: group::access
+default_enabled: false

db/migrate/20210311001204_change_traversal_ids_to_bigint.rb

-# frozen_string_literal: true
-
-class ChangeTraversalIdsToBigint < ActiveRecord::Migration[6.0]
-  include Gitlab::Database::MigrationHelpers
-
-  DOWNTIME = false
-
-  disable_ddl_transaction!
-
-  def up
-    change_column_type_concurrently :namespaces, :traversal_ids, 'bigint[]'
-  end
-
-  def down
-    undo_change_column_type_concurrently :namespaces, :traversal_ids
-  end
-end

db/schema_migrations/20210311001204

-55b0e53343a5a053097d88842e1b8ebce7a61d168ba14f6866b5800d3b28a4fa
\ No newline at end of file

db/structure.sql

@@ -150,15 +150,6 @@ $$;
 
 COMMENT ON FUNCTION table_sync_function_2be879775d() IS 'Partitioning migration: table sync for audit_events table';
 
-CREATE FUNCTION trigger_26c50fd02f49() RETURNS trigger
-    LANGUAGE plpgsql
-    AS $$
-BEGIN
-  NEW."traversal_ids_for_type_change" := NEW."traversal_ids";
-  RETURN NEW;
-END;
-$$;
-
 CREATE TABLE audit_events (
     id bigint NOT NULL,
     author_id integer NOT NULL,
@@ -24563,8 +24554,6 @@ CREATE TRIGGER table_sync_trigger_b99eb6998c AFTER INSERT OR DELETE OR UPDATE ON
 
 CREATE TRIGGER table_sync_trigger_ee39a25f9d AFTER INSERT OR DELETE OR UPDATE ON audit_events FOR EACH ROW EXECUTE PROCEDURE table_sync_function_2be879775d();
 
-CREATE TRIGGER trigger_26c50fd02f49 BEFORE INSERT OR UPDATE ON namespaces FOR EACH ROW EXECUTE PROCEDURE trigger_26c50fd02f49();
-
 CREATE TRIGGER trigger_has_external_issue_tracker_on_delete AFTER DELETE ON services FOR EACH ROW WHEN ((((old.category)::text = 'issue_tracker'::text) AND (old.active = true) AND (old.project_id IS NOT NULL))) EXECUTE PROCEDURE set_has_external_issue_tracker();
 
 CREATE TRIGGER trigger_has_external_issue_tracker_on_insert AFTER INSERT ON services FOR EACH ROW WHEN ((((new.category)::text = 'issue_tracker'::text) AND (new.active = true) AND (new.project_id IS NOT NULL))) EXECUTE PROCEDURE set_has_external_issue_tracker();

spec/lib/gitlab/object_hierarchy_spec.rb

@@ -195,9 +195,24 @@ RSpec.describe Gitlab::ObjectHierarchy do
     it_behaves_like 'Gitlab::ObjectHierarchy test cases'
 
     it 'calls DISTINCT' do
-      expect(parent.self_and_descendants.to_sql).to include("DISTINCT")
       expect(child2.self_and_ancestors.to_sql).to include("DISTINCT")
     end
+
+    context 'when use_traversal_ids feature flag is enabled' do
+      it 'does not call DISTINCT' do
+        expect(parent.self_and_descendants.to_sql).not_to include("DISTINCT")
+      end
+    end
+
+    context 'when use_traversal_ids feature flag is disabled' do
+      before do
+        stub_feature_flags(use_traversal_ids: false)
+      end
+
+      it 'calls DISTINCT' do
+        expect(parent.self_and_descendants.to_sql).to include("DISTINCT")
+      end
+    end
   end
 
   context 'when the use_distinct_for_all_object_hierarchy feature flag is enabled' do
@@ -209,10 +224,24 @@ RSpec.describe Gitlab::ObjectHierarchy do
     it_behaves_like 'Gitlab::ObjectHierarchy test cases'
 
     it 'calls DISTINCT' do
-      expect(parent.self_and_descendants.to_sql).to include("DISTINCT")
       expect(child2.self_and_ancestors.to_sql).to include("DISTINCT")
     end
 
+    context 'when use_traversal_ids feature flag is enabled' do
+      it 'does not call DISTINCT' do
+        expect(parent.self_and_descendants.to_sql).not_to include("DISTINCT")
+      end
+    end
+
+    context 'when use_traversal_ids feature flag is disabled' do
+      before do
+        stub_feature_flags(use_traversal_ids: false)
+      end
+
+      it 'calls DISTINCT' do
+        expect(parent.self_and_descendants.to_sql).to include("DISTINCT")
+      end
+
       context 'when the skip_ordering option is set' do
         let(:options) { { skip_ordering: true } }
 
@@ -226,6 +255,7 @@ RSpec.describe Gitlab::ObjectHierarchy do
         end
       end
     end
+  end
 
   context 'when the use_distinct_in_object_hierarchy feature flag is disabled' do
     before do

spec/models/namespace/traversal_hierarchy_spec.rb

@@ -62,10 +62,6 @@ RSpec.describe Namespace::TraversalHierarchy, type: :model do
       Namespace.update_all(traversal_ids: [])
     end
 
-    before do
-      Namespace.update_all(traversal_ids: [])
-    end
-
     it { is_expected.to match_array Namespace.all }
 
     context 'when lock is true' do

spec/models/namespace_spec.rb

@@ -873,7 +873,51 @@ RSpec.describe Namespace do
     end
   end
 
-  it_behaves_like 'recursive namespace traversal'
+  describe '#use_traversal_ids?' do
+    let_it_be(:namespace) { build(:namespace) }
+
+    subject { namespace.use_traversal_ids? }
+
+    context 'when use_traversal_ids feature flag is true' do
+      before do
+        stub_feature_flags(use_traversal_ids: true)
+      end
+
+      it { is_expected.to eq true }
+    end
+
+    context 'when use_traversal_ids feature flag is false' do
+      before do
+        stub_feature_flags(use_traversal_ids: false)
+      end
+
+      it { is_expected.to eq false }
+    end
+  end
+
+  context 'when use_traversal_ids feature flag is true' do
+    it_behaves_like 'namespace traversal'
+
+    describe '#self_and_descendants' do
+      subject { namespace.self_and_descendants }
+
+      it { expect(subject.to_sql).to include 'traversal_ids @>' }
+    end
+  end
+
+  context 'when use_traversal_ids feature flag is false' do
+    before do
+      stub_feature_flags(use_traversal_ids: false)
+    end
+
+    it_behaves_like 'namespace traversal'
+
+    describe '#self_and_descendants' do
+      subject { namespace.self_and_descendants }
+
+      it { expect(subject.to_sql).not_to include 'traversal_ids @>' }
+    end
+  end
 
   describe '#users_with_descendants' do
     let(:user_a) { create(:user) }

spec/support/shared_examples/namespaces/recursive_traversal_examples.rb → spec/support/shared_examples/namespaces/namespace_traversal_examples.rb

 # frozen_string_literal: true
 
-RSpec.shared_examples 'recursive namespace traversal' do
+RSpec.shared_examples 'namespace traversal' do
   describe '#self_and_hierarchy' do
     let!(:group) { create(:group, path: 'git_lab') }
     let!(:nested_group) { create(:group, parent: group) }