Preload Group plans in EpicsFinder

When displaying a Group details page we check if epics are enabled for
the group and all of its descendants. For this set we will eventually
run at least one SQL query per Group to get the plans said Group has
access to. In case of a hierarchy with 20 groups this means at least 20
SQL queries.

In this commit we introduce Gitlab::GroupPlansPreloader. This is a
custom class that takes an ActiveRecord::Relation of Group objects and
efficiently preloads the plans for these groups. This class only needs
three SQL queries to do its work:

1. A SQL query to get the groups and all of their ancestors.
2. A SQL query to get all the plans all of these groups have access to.
3. A SQL query to get the groups to assign the plans to, then return to
   the caller.

Using this data the class can then determine which groups have access to
what plans.

Fixes https://gitlab.com/gitlab-org/gitlab-ee/issues/6206

ee/app/finders/epics_finder.rb

@@ -52,6 +52,8 @@ class EpicsFinder < IssuableFinder
   private
 
   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

ee/app/models/ee/namespace.rb

@@ -155,6 +155,19 @@ module EE
       actual_plan&.pipeline_size_limit.to_i
     end
 
+    def memoized_plans=(plans)
+      @plans = plans # rubocop: disable Gitlab/ModuleWithInstanceVariables
+    end
+
+    def plans
+      @plans ||=
+        if parent_id
+          Plan.where(id: self_and_ancestors.with_plan.reorder(nil).select(:plan_id))
+        else
+          Array(plan)
+        end
+    end
+
     private
 
     def validate_plan_name
@@ -180,14 +193,5 @@ module EE
         globally_available
       end
     end
-
-    def plans
-      @plans ||=
-        if parent_id
-          Plan.where(id: self_and_ancestors.with_plan.reorder(nil).select(:plan_id))
-        else
-          Array(plan)
-        end
-    end
   end
 end

ee/changelogs/unreleased/ee-groups-controller-show-performance.yml

+---
+title: Preload Group plans in EpicsFinder
+merge_request:
+author:
+type: performance

ee/lib/gitlab/group_plans_preloader.rb

+# frozen_string_literal: true
+
+module Gitlab
+  # Preloading of Plans for one or more groups.
+  #
+  # This class can be used to efficiently preload the plans of a given list of
+  # groups, including any plans the groups may have access to based on their
+  # parent groups.
+  class GroupPlansPreloader
+    # Preloads all the plans for the given Groups.
+    #
+    # groups - An ActiveRecord::Relation returning a set of Group instances.
+    #
+    # Returns an Array containing all the Groups, including their preloaded
+    # plans.
+    def preload(groups)
+      groups_and_ancestors = groups_and_ancestors_for(groups)
+
+      # A Hash mapping group IDs to their corresponding Group instances.
+      groups_map = groups_and_ancestors.each_with_object({}) do |group, hash|
+        hash[group.id] = group
+      end
+
+      all_plan_ids = Set.new
+
+      # A Hash that for every group ID maps _all_ the plan IDs this group has
+      # access to.
+      plans_map = groups_and_ancestors
+        .each_with_object(Hash.new { |h, k| h[k] = [] }) do |group, hash|
+          current = group
+
+          while current
+
+            if (plan_id = current.plan_id)
+              hash[group.id] << plan_id
+              all_plan_ids << plan_id
+            end
+
+            current = groups_map[current.parent_id]
+          end
+        end
+
+      # Grab all the plans for all the Groups, using only a single query.
+      plans = Plan
+        .where(id: all_plan_ids.to_a)
+        .each_with_object({}) do |plan, hash|
+          hash[plan.id] = plan
+        end
+
+      # Assign all the plans to the groups that have access to them.
+      groups.each do |group|
+        group.memoized_plans = plans_map[group.id].map { |id| plans[id] }
+      end
+    end
+
+    # Returns an ActiveRecord::Relation that includes the given groups, and all
+    # their (recursive) ancestors.
+    def groups_and_ancestors_for(groups)
+      Gitlab::GroupHierarchy
+        .new(groups)
+        .base_and_ancestors
+        .select(:id, :parent_id, :plan_id)
+    end
+  end
+end

ee/spec/finders/epics_finder_spec.rb

@@ -53,6 +53,12 @@ describe EpicsFinder do
           expect(epics).to contain_exactly(epic1, epic2, epic3)
         end
 
+        it 'does not execute more than 7 SQL queries' do
+          amount = ActiveRecord::QueryRecorder.new { epics.to_a }.count
+
+          expect(amount).to be <= 7
+        end
+
         context 'by created_at' do
           it 'returns all epics created before the given date' do
             expect(epics(created_before: 2.days.ago)).to contain_exactly(epic1, epic2)
@@ -97,6 +103,24 @@ describe EpicsFinder do
           it 'returns all epics that belong to the given group and its subgroups' do
             expect(epics).to contain_exactly(epic1, epic2, epic3, subepic1, subepic2)
           end
+
+          it 'does not execute more than 9 SQL queries' do
+            amount = ActiveRecord::QueryRecorder.new { epics.to_a }.count
+
+            expect(amount).to be <= 9
+          end
+
+          it 'does not execute more than 11 SQL queries when checking namespace plans' do
+            allow(Gitlab::CurrentSettings)
+              .to receive(:should_check_namespace_plan?)
+              .and_return(true)
+
+            group.update(plan: create(:gold_plan))
+
+            amount = ActiveRecord::QueryRecorder.new { epics.to_a }.count
+
+            expect(amount).to be <= 10
+          end
         end
 
         context 'by timeframe' do

ee/spec/lib/gitlab/group_plans_preloader_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Gitlab::GroupPlansPreloader do
+  describe '#preload' do
+    let!(:plan1) { create(:free_plan, name: 'plan-1') }
+    let!(:plan2) { create(:free_plan, name: 'plan-2') }
+
+    let(:preloaded_groups) do
+      # We don't use the factory objects here because they might have the plan
+      # loaded already (as we specify the plan when creating them).
+      described_class.new.preload(Group.order(id: :asc))
+    end
+
+    before do
+      group1 = create(:group, name: 'group-1', plan_id: plan1.id)
+
+      create(:group, name: 'group-2', plan_id: plan2.id)
+      create(:group, name: 'group-3', parent: group1)
+    end
+
+    it 'only executes three SQL queries to preload the data' do
+      amount = ActiveRecord::QueryRecorder
+        .new { preloaded_groups }
+        .count
+
+      # One query to get the groups and their ancestors, one query to get their
+      # plans, and one query to _just_ get the groups.
+      expect(amount).to eq(3)
+    end
+
+    it 'associates the correct plans with the correct groups' do
+      expect(preloaded_groups[0].plans).to eq([plan1])
+      expect(preloaded_groups[1].plans).to eq([plan2])
+      expect(preloaded_groups[2].plans).to eq([plan1])
+    end
+
+    it 'does not execute any queries for preloaded plans' do
+      preloaded_groups
+
+      amount = ActiveRecord::QueryRecorder
+        .new { preloaded_groups.each(&:plans) }
+        .count
+
+      expect(amount).to be_zero
+    end
+  end
+end