Preload group ancestor to decrease N+1 issue

When checking permissions for a group, it always tries to load
group ancestor (to check SSO). If we check permissions for related
groups (for which we are sure that all of them have same root), we
just preset this root group before checking permission.

ee/app/finders/epics_finder.rb

@@ -78,7 +78,9 @@ class EpicsFinder < IssuableFinder
                # The `group` method takes care of checking permissions
                [group]
              else
-               groups_user_can_read_epics(related_groups)
+               # `same_root` should be set only if we are sure that all groups
+               # in related_groups have the same ancestor root group
+               ::Group.groups_user_can_read_epics(related_groups, current_user, same_root: true)
              end
 
     Epic.where(group: groups)
@@ -112,16 +114,6 @@ class EpicsFinder < IssuableFinder
     end
   end
 
-  # rubocop: disable CodeReuse/ActiveRecord
-  def groups_user_can_read_epics(groups)
-    groups = Gitlab::GroupPlansPreloader.new.preload(groups)
-
-    DeclarativePolicy.user_scope do
-      groups.select { |g| Ability.allowed?(current_user, :read_epic, g) }
-    end
-  end
-  # rubocop: enable CodeReuse/ActiveRecord
-
   # rubocop: disable CodeReuse/ActiveRecord
   def by_timeframe(items)
     return items unless params[:start_date] && params[:end_date]

ee/app/models/ee/epic.rb

@@ -195,21 +195,6 @@ module EE
         ::Gitlab::ObjectHierarchy.new(self.where(parent_id: nil)).max_descendants_depth
       end
 
-      def groups_user_can_read_epics(epics, user)
-        groups = if ::Feature.enabled?(:optimized_groups_user_can_read_epics_method)
-                   epics_query = epics.select(:group_id)
-                   ::Group.joins("INNER JOIN (#{epics_query.to_sql}) as epics on epics.group_id = namespaces.id")
-                 else
-                   ::Group.where(id: epics.select(:group_id))
-                 end
-
-        groups = ::Gitlab::GroupPlansPreloader.new.preload(groups)
-
-        DeclarativePolicy.user_scope do
-          groups.select { |g| Ability.allowed?(user, :read_epic, g) }
-        end
-      end
-
       def related_issues(ids:, preload: nil)
         ::Issue.select('issues.*, epic_issues.id as epic_issue_id, epic_issues.relative_position, epic_issues.epic_id as epic_id')
           .joins(:epic_issue)

ee/app/models/ee/group.rb

@@ -72,6 +72,15 @@ module EE
         ).where.not(file_template_project_id: nil)
       end
 
+      scope :for_epics, ->(epics) do
+        if ::Feature.enabled?(:optimized_groups_user_can_read_epics_method)
+          epics_query = epics.select(:group_id)
+          joins("INNER JOIN (#{epics_query.to_sql}) as epics on epics.group_id = namespaces.id")
+        else
+          where(id: epics.select(:group_id))
+        end
+      end
+
       state_machine :ldap_sync_status, namespace: :ldap_sync, initial: :ready do
         state :ready
         state :started
@@ -114,6 +123,29 @@ module EE
       end
     end
 
+    class_methods do
+      def groups_user_can_read_epics(groups, user, same_root: false)
+        groups = ::Gitlab::GroupPlansPreloader.new.preload(groups)
+
+        # if we are sure that all groups have the same root group, we can
+        # preset root_ancestor for all of them to avoid an additional SQL query
+        # done for each group permission check:
+        # https://gitlab.com/gitlab-org/gitlab/issues/11539
+        preset_root_ancestor_for(groups) if same_root && ::Feature.enabled?(:preset_group_root)
+
+        DeclarativePolicy.user_scope do
+          groups.select { |group| Ability.allowed?(user, :read_epic, group) }
+        end
+      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
       return unless ip_restrictions.present?
 

ee/app/models/ee/namespace.rb

@@ -25,6 +25,8 @@ module EE
     prepended do
       include EachBatch
 
+      attr_writer :root_ancestor
+
       belongs_to :plan
 
       has_one :namespace_statistics

ee/app/services/epics/descendant_count_service.rb

@@ -48,7 +48,8 @@ module Epics
     def accessible_epics
       strong_memoize(:epics) do
         epics = epic.base_and_descendants
-        groups = Epic.groups_user_can_read_epics(epics, current_user)
+        epic_groups = Group.for_epics(epics)
+        groups = Group.groups_user_can_read_epics(epic_groups, current_user, same_root: true)
         epics.in_selected_groups(groups)
       end
     end

ee/spec/models/epic_spec.rb

@@ -101,64 +101,6 @@ describe Epic do
     end
   end
 
-  shared_examples '.groups_user_can_read_epics examples' do
-    let_it_be(:private_group) { create(:group, :private) }
-    let_it_be(:epic) { create(:epic, group: private_group) }
-
-    subject do
-      epics = described_class.where(id: epic.id)
-      described_class.groups_user_can_read_epics(epics, user)
-    end
-
-    it 'does not return inaccessible groups' do
-      expect(subject).to be_empty
-    end
-
-    context 'with authorized user' do
-      before do
-        private_group.add_developer(user)
-      end
-
-      context 'with epics enabled' do
-        before do
-          stub_licensed_features(epics: true)
-        end
-
-        it 'returns epic groups user can access' do
-          expect(subject).to eq [private_group]
-        end
-      end
-
-      context 'with epics are disabled' do
-        before do
-          stub_licensed_features(epics: false)
-        end
-
-        it 'returns an empty list' do
-          expect(subject).to be_empty
-        end
-      end
-    end
-  end
-
-  describe '.groups_user_can_read_epics' do
-    context 'with `optimized_groups_user_can_read_epics_method` feature flag enabled' do
-      before do
-        stub_feature_flags(optimized_groups_user_can_read_epics_method: true)
-      end
-
-      include_examples '.groups_user_can_read_epics examples'
-    end
-
-    context 'with `optimized_groups_user_can_read_epics_method` feature flag disabled' do
-      before do
-        stub_feature_flags(optimized_groups_user_can_read_epics_method: false)
-      end
-
-      include_examples '.groups_user_can_read_epics examples'
-    end
-  end
-
   describe '#valid_parent?' do
     context 'basic checks' do
       let(:epic) { build(:epic, group: group) }

ee/spec/models/group_spec.rb

@@ -72,6 +72,36 @@ describe Group do
                                         group_not_marked_for_deletion)
       end
     end
+
+    describe '.for_epics' do
+      let_it_be(:epic1) { create(:epic) }
+      let_it_be(:epic2) { create(:epic) }
+
+      shared_examples '.for_epics examples' do
+        it 'returns groups only for selected epics' do
+          epics = ::Epic.where(id: epic1)
+          expect(described_class.for_epics(epics)).to contain_exactly(epic1.group)
+        end
+      end
+
+      context 'with `optimized_groups_user_can_read_epics_method` feature flag' do
+        before do
+          stub_feature_flags(optimized_groups_user_can_read_epics_method: flag_state)
+        end
+
+        context 'enabled' do
+          let(:flag_state) { true }
+
+          include_examples '.for_epics examples'
+        end
+
+        context 'disabled' do
+          let(:flag_state) { false }
+
+          include_examples '.for_epics examples'
+        end
+      end
+    end
   end
 
   describe 'validations' do
@@ -167,6 +197,107 @@ describe Group do
     end
   end
 
+  describe '.groups_user_can_read_epics' do
+    let_it_be(:user) { create(:user) }
+    let_it_be(:private_group) { create(:group, :private) }
+
+    subject do
+      groups = described_class.where(id: private_group.id)
+      described_class.groups_user_can_read_epics(groups, user)
+    end
+
+    it 'does not return inaccessible groups' do
+      expect(subject).to be_empty
+    end
+
+    context 'with authorized user' do
+      before do
+        private_group.add_developer(user)
+      end
+
+      context 'with epics enabled' do
+        before do
+          stub_licensed_features(epics: true)
+        end
+
+        it 'returns epic groups user can access' do
+          expect(subject).to eq [private_group]
+        end
+      end
+
+      context 'with epics disabled' do
+        before do
+          stub_licensed_features(epics: false)
+        end
+
+        it 'returns an empty list' do
+          expect(subject).to be_empty
+        end
+      end
+    end
+
+    context 'getting group root ancestor' do
+      let_it_be(:subgroup1) { create(:group, :private, parent: private_group) }
+      let_it_be(:subgroup2) { create(:group, :private, parent: subgroup1) }
+
+      shared_examples 'group root ancestor' do
+        it 'does not exceed SQL queries count' do
+          groups = described_class.where(id: subgroup1)
+          control_count = ActiveRecord::QueryRecorder.new do
+            described_class.groups_user_can_read_epics(groups, user, params)
+          end.count
+
+          groups = described_class.where(id: [subgroup1, subgroup2])
+          expect { described_class.groups_user_can_read_epics(groups, user, params) }
+            .not_to exceed_query_limit(control_count + extra_query_count)
+        end
+      end
+
+      context 'with `preset_group_root` feature flag disabled' do
+        before do
+          stub_feature_flags(preset_group_root: false)
+        end
+
+        it_behaves_like 'group root ancestor' do
+          let(:params) { {} }
+          let(:extra_query_count) { 6 }
+        end
+      end
+
+      context 'with `preset_group_root` feature flag enabled' do
+        before do
+          stub_feature_flags(preset_group_root: true)
+        end
+
+        context 'when same_root is false' do
+          let(:params) { { same_root: false } }
+
+          # extra 6 queries:
+          # * getting root_ancestor
+          # * getting root ancestor's saml_provider
+          # * check if group has projects
+          # * max_member_access_for_user_from_shared_groups
+          # * max_member_access_for_user
+          # * self_and_ancestors_ids
+          it_behaves_like 'group root ancestor' do
+            let(:extra_query_count) { 6 }
+          end
+        end
+
+        context 'when same_root is true' do
+          let(:params) { { same_root: true } }
+
+          # avoids 2 queries from the list above:
+          # * getting root ancestor
+          # * getting root ancestor's saml_provider
+          it_behaves_like 'group root ancestor' do
+            let(:extra_query_count) { 4 }
+          end
+        end
+      end
+    end
+  end
+
   describe '#vulnerable_projects' do
     it "fetches the group's projects that have vulnerabilities" do
       vulnerable_project = create(:project, namespace: group)