Preset root ancestor when finding labels

Because we search for permissioned labels which all have the same
top-level group, we can preset root_ancestor for them before permission
check. This optimization saves us extra query to fetch root ancestor
(called from group policy).

app/finders/labels_finder.rb

@@ -100,6 +100,10 @@ class LabelsFinder < UnionFinder
     strong_memoize(:group_ids) do
       groups = groups_to_include(group)
 
+      # Because we are sure that all groups are in the same hierarchy tree
+      # we can preset root group for all of them to optimize permission checks
+      Group.preset_root_ancestor_for(groups) if Feature.enabled?(:preset_root_ancestor_for_labels, group)
+
       groups_user_can_read_labels(groups).map(&:id)
     end
   end

app/models/group.rb

@@ -169,6 +169,15 @@ class Group < Namespace
       where('NOT EXISTS (?)', services)
     end
 
+    # This method can be used only if all groups have the same top-level
+    # group
+    def preset_root_ancestor_for(groups)
+      return groups if groups.size < 2
+
+      root = groups.first.root_ancestor
+      groups.drop(1).each { |group| group.root_ancestor = root }
+    end
+
     private
 
     def public_to_user_arel(user)

app/models/namespace.rb

@@ -103,6 +103,10 @@ class Namespace < ApplicationRecord
       )
   end
 
+  # Make sure that the name is same as strong_memoize name in root_ancestor
+  # method
+  attr_writer :root_ancestor
+
   class << self
     def by_path(path)
       find_by('lower(path) = :value', value: path.downcase)
@@ -315,6 +319,8 @@ class Namespace < ApplicationRecord
   def root_ancestor
     return self if persisted? && parent_id.nil?
 
+    # Make sure that strong_memoize name is in sync with root_ancestor's
+    # attr_writer name
     strong_memoize(:root_ancestor) do
       self_and_ancestors.reorder(nil).find_by(parent_id: nil)
     end

config/feature_flags/development/preset_root_ancestor_for_labels.yml

+---
+name: preset_root_ancestor_for_labels
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/52236
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/299521
+milestone: '13.9'
+type: development
+group: group::plan
+default_enabled: false

ee/app/models/ee/group.rb

@@ -163,13 +163,6 @@ module EE
       def groups_user_can_read_epics(groups, user, same_root: false)
         groups_user_can(groups, user, :read_epic, same_root: same_root)
       end
-
-      def preset_root_ancestor_for(groups)
-        return groups if groups.size < 2
-
-        root = groups.first.root_ancestor
-        groups.drop(1).each { |group| group.root_ancestor = root }
-      end
     end
 
     def ip_restriction_ranges

ee/app/models/ee/namespace.rb

@@ -23,8 +23,6 @@ module EE
     prepended do
       include EachBatch
 
-      attr_writer :root_ancestor
-
       has_one :namespace_statistics
       has_one :namespace_limit, inverse_of: :namespace
       has_one :gitlab_subscription

ee/spec/requests/groups/labels_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe 'view audit events' do
+  describe 'GET /groups/:group/-/audit_events' do
+    let_it_be(:group) { create(:group, :public) }
+    let_it_be(:group1) { create(:group, parent: group) }
+
+    it 'returns 200 response' do
+      send_request
+
+      expect(response).to have_gitlab_http_status(:ok)
+    end
+
+    it 'avoids N+1 DB queries', :request_store do
+      control = ActiveRecord::QueryRecorder.new(skip_cached: false) { send_request }
+
+      create_list(:group, 3, parent: group)
+
+      expect { send_request }.not_to exceed_all_query_limit(control)
+    end
+
+    context 'when preset_root_ancestor_for_labels flag is disabled' do
+      before do
+        stub_feature_flags(preset_root_ancestor_for_labels: false)
+      end
+
+      it 'does additional N+1 DB queries', :request_store do
+        control = ActiveRecord::QueryRecorder.new(skip_cached: false) { send_request }
+
+        create_list(:group, 3, parent: group)
+
+        expect { send_request }.to exceed_all_query_limit(control)
+      end
+    end
+
+    def send_request
+      get group_labels_path(group, include_descendant_groups: true, format: :json)
+    end
+  end
+end

spec/models/group_spec.rb

@@ -1492,6 +1492,28 @@ RSpec.describe Group do
     end
   end
 
+  describe '.preset_root_ancestor_for' do
+    let_it_be(:rootgroup, reload: true) { create(:group) }
+    let_it_be(:subgroup, reload: true) { create(:group, parent: rootgroup) }
+    let_it_be(:subgroup2, reload: true) { create(:group, parent: subgroup) }
+
+    it 'does noting for single group' do
+      expect(subgroup).not_to receive(:self_and_ancestors)
+
+      described_class.preset_root_ancestor_for([subgroup])
+    end
+
+    it 'sets the same root_ancestor for multiple groups' do
+      expect(subgroup).not_to receive(:self_and_ancestors)
+      expect(subgroup2).not_to receive(:self_and_ancestors)
+
+      described_class.preset_root_ancestor_for([rootgroup, subgroup, subgroup2])
+
+      expect(subgroup.root_ancestor).to eq(rootgroup)
+      expect(subgroup2.root_ancestor).to eq(rootgroup)
+    end
+  end
+
   def subject_and_reload(*models)
     subject
     models.map(&:reload)