Merge branch 'cablett-recursive-epic-tree-counts-weights' into 'master'

Bulk lazy loader for epic aggregates

See merge request gitlab-org/gitlab!23406

app/graphql/gitlab_schema.rb

@@ -141,3 +141,5 @@ class GitlabSchema < GraphQL::Schema
     end
   end
 end
+
+GitlabSchema.prepend_if_ee('EE::GitlabSchema')

doc/api/graphql/reference/gitlab_schema.graphql

@@ -1948,8 +1948,8 @@ type Epic implements Noteable {
   descendantCounts: EpicDescendantCount
 
   """
-  Total weight of open and closed descendant epic's issues. Available only when
-  feature flag `unfiltered_epic_aggregates` is enabled.
+  Total weight of open and closed issues in the epic and its descendants.
+  Available only when feature flag `unfiltered_epic_aggregates` is enabled.
   """
   descendantWeightSum: EpicDescendantWeights
 

doc/api/graphql/reference/index.md

@@ -317,7 +317,7 @@ Represents an epic.
 | `closedAt` | Time | Timestamp of the epic's closure |
 | `createdAt` | Time | Timestamp of the epic's creation |
 | `descendantCounts` | EpicDescendantCount | Number of open and closed descendant epics and issues |
-| `descendantWeightSum` | EpicDescendantWeights | Total weight of open and closed descendant epic's issues. Available only when feature flag `unfiltered_epic_aggregates` is enabled. |
+| `descendantWeightSum` | EpicDescendantWeights | Total weight of open and closed issues in the epic and its descendants. Available only when feature flag `unfiltered_epic_aggregates` is enabled. |
 | `description` | String | Description of the epic |
 | `downvotes` | Int! | Number of downvotes the epic has received |
 | `dueDate` | Time | Due date of the epic |

ee/app/graphql/ee/gitlab_schema.rb

+# frozen_string_literal: true
+
+module EE
+  module GitlabSchema
+    extend ActiveSupport::Concern
+
+    prepended do
+      lazy_resolve ::Gitlab::Graphql::Aggregations::Epics::LazyEpicAggregate, :epic_aggregate
+    end
+  end
+end

ee/app/graphql/types/epic_type.rb

@@ -2,6 +2,8 @@
 
 module Types
   class EpicType < BaseObject
+    include ::Gitlab::Graphql::Aggregations::Epics::Constants
+
     graphql_name 'Epic'
     description 'Represents an epic.'
 
@@ -121,21 +123,20 @@ module Types
           resolver: Resolvers::EpicIssuesResolver
 
     field :descendant_counts, Types::EpicDescendantCountType, null: true, complexity: 10,
-          description: 'Number of open and closed descendant epics and issues',
-          resolve: -> (epic, args, ctx) do
-            Epics::DescendantCountService.new(epic, ctx[:current_user])
-          end
+      description: 'Number of open and closed descendant epics and issues',
+      resolve: -> (epic, args, ctx) do
+        if Feature.enabled?(:unfiltered_epic_aggregates)
+          Gitlab::Graphql::Aggregations::Epics::LazyEpicAggregate.new(ctx, epic.id, COUNT)
+        else
+          Epics::DescendantCountService.new(epic, ctx[:current_user])
+        end
+      end
 
     field :descendant_weight_sum, Types::EpicDescendantWeightSumType, null: true, complexity: 10,
-          description: "Total weight of open and closed descendant epic's issues",
-          feature_flag: :unfiltered_epic_aggregates
-
-    def descendant_weight_sum
-      OpenStruct.new(
-        # We shouldn't stop the whole query, so returning -1 for a semi-noisy error
-        opened_issues: -1,
-        closed_issues: -1
-      )
-    end
+      description: "Total weight of open and closed issues in the epic and its descendants",
+      feature_flag: :unfiltered_epic_aggregates,
+      resolve: -> (epic, args, ctx) do
+        Gitlab::Graphql::Aggregations::Epics::LazyEpicAggregate.new(ctx, epic.id, WEIGHT_SUM)
+      end
   end
 end

ee/lib/gitlab/graphql/aggregations/epics/constants.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Graphql
+    module Aggregations
+      module Epics
+        module Constants
+          ISSUE_TYPE = :issue
+          EPIC_TYPE = :epic
+
+          CLOSED_ISSUE_STATE = Issue.available_states[:closed].freeze
+          OPENED_ISSUE_STATE = Issue.available_states[:opened].freeze
+
+          CLOSED_EPIC_STATE = Epic.available_states[:closed].freeze
+          OPENED_EPIC_STATE = Epic.available_states[:opened].freeze
+
+          COUNT = :count
+          WEIGHT_SUM = :weight_sum
+        end
+      end
+    end
+  end
+end

ee/lib/gitlab/graphql/aggregations/epics/epic_node.rb

+# frozen_string_literal: true
+
+# This class represents an Epic's aggregate information (added up counts) about its child epics and direct issues
+
+module Gitlab
+  module Graphql
+    module Aggregations
+      module Epics
+        class EpicNode
+          include ::Gitlab::Graphql::Aggregations::Epics::Constants
+          include Gitlab::Utils::StrongMemoize
+
+          attr_reader :epic_id, :epic_state_id, :epic_info_flat_list, :parent_id,
+                      :count_aggregate, :weight_sum_aggregate
+
+          attr_accessor :children, :calculated_count_totals, :calculated_weight_sum_totals
+
+          def initialize(epic_id, flat_info_list)
+            # epic aggregate records from the DB loader look like the following:
+            # { 1 => [{iid: 1, parent_id: nil, epic_state_id: 1, issues_count: 1, issues_weight_sum: 2, issues_state_id: 2}] ... }
+            # They include the sum of each epic's direct issues, grouped by status,
+            # so in order to get a sum of the entire tree, we have to add that up recursively
+            @epic_id = epic_id
+            @epic_info_flat_list = flat_info_list
+            @children = []
+            @sums = {}
+
+            set_epic_attributes(flat_info_list.first) # there will always be one
+          end
+
+          def aggregate_count
+            strong_memoize(:count_aggregate) do
+              OpenStruct.new({
+                opened_issues: sum_objects(COUNT, OPENED_ISSUE_STATE, ISSUE_TYPE),
+                closed_issues: sum_objects(COUNT, CLOSED_ISSUE_STATE, ISSUE_TYPE),
+                opened_epics: sum_objects(COUNT, OPENED_EPIC_STATE, EPIC_TYPE),
+                closed_epics: sum_objects(COUNT, CLOSED_EPIC_STATE, EPIC_TYPE)
+             })
+            end
+          end
+
+          def aggregate_weight_sum
+            strong_memoize(:weight_sum_aggregate) do
+              OpenStruct.new({
+                opened_issues: sum_objects(WEIGHT_SUM, OPENED_ISSUE_STATE, ISSUE_TYPE),
+                closed_issues: sum_objects(WEIGHT_SUM, CLOSED_ISSUE_STATE, ISSUE_TYPE)
+               })
+            end
+          end
+
+          def to_s
+            {
+                epic_id: @epic_id,
+                parent_id: @parent_id,
+                children: children,
+                object_id: object_id
+            }.to_s
+          end
+
+          def sum_objects(facet, state, type)
+            key = [facet, state, type]
+            return @sums[key] if @sums[key]
+
+            direct_sum = value_from_records(*key)
+            sum_from_children = children.inject(0) do |total, child|
+              total += child.sum_objects(*key)
+              total
+            end
+
+            @sums[key] = direct_sum + sum_from_children
+          end
+
+          private
+
+          def set_epic_attributes(record)
+            @epic_state_id = record[:epic_state_id]
+            @parent_id = record[:parent_id]
+          end
+
+          def value_from_records(facet, state, type)
+            # DB records look like:
+            # {iid: 1, epic_state_id: 1, issues_count: 1, issues_weight_sum: 2, parent_id: nil, issues_state_id: 2}
+            if type == EPIC_TYPE
+              # can only be COUNT
+              children.select { |node| node.epic_state_id == state }.count
+            else
+              matching_record = epic_info_flat_list.find do |record|
+                record[:issues_state_id] == state
+              end || {}
+
+              matching_record.fetch("issues_#{facet}".to_sym, 0)
+            end
+          end
+        end
+      end
+    end
+  end
+end

ee/lib/gitlab/graphql/aggregations/epics/lazy_epic_aggregate.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Graphql
+    module Aggregations
+      module Epics
+        class LazyEpicAggregate
+          include ::Gitlab::Graphql::Aggregations::Epics::Constants
+
+          attr_reader :facet, :epic_id, :lazy_state
+
+          PERMITTED_FACETS = [COUNT, WEIGHT_SUM].freeze
+
+          # Because facets "count" and "weight_sum" share the same db query, but have a different graphql type object,
+          # we can separate them and serve only the fields which are requested by the GraphQL query
+          def initialize(query_ctx, epic_id, aggregate_facet)
+            @epic_id = epic_id
+
+            error = validate_facet(aggregate_facet)
+            if error
+              raise ArgumentError.new("#{error}. Please specify either #{COUNT} or #{WEIGHT_SUM}")
+            end
+
+            @facet = aggregate_facet.to_sym
+
+            # Initialize the loading state for this query,
+            # or get the previously-initiated state
+            @lazy_state = query_ctx[:lazy_epic_aggregate] ||= {
+                pending_ids: Set.new,
+                tree: {}
+            }
+            # Register this ID to be loaded later:
+            @lazy_state[:pending_ids] << epic_id
+          end
+
+          # Return the loaded record, hitting the database if needed
+          def epic_aggregate
+            # Check if the record was already loaded:
+            # load from tree by epic
+            unless tree[@epic_id]
+              load_records_into_tree
+            end
+
+            aggregate_object(tree[@epic_id])
+          end
+
+          private
+
+          def validate_facet(aggregate_facet)
+            unless aggregate_facet.present?
+              return "No aggregate facet provided."
+            end
+
+            unless PERMITTED_FACETS.include?(aggregate_facet.to_sym)
+              return "Invalid aggregate facet #{aggregate_facet} provided."
+            end
+          end
+
+          def tree
+            @lazy_state[:tree]
+          end
+
+          def load_records_into_tree
+            # The record hasn't been loaded yet, so
+            # hit the database with all pending IDs
+            pending_ids = @lazy_state[:pending_ids].to_a
+
+            # Fire off the db query and get the results (grouped by epic_id and facet)
+            raw_epic_aggregates = Gitlab::Graphql::Loaders::BulkEpicAggregateLoader.new(epic_ids: pending_ids).execute
+            create_epic_nodes(raw_epic_aggregates)
+            @lazy_state[:pending_ids].clear
+          end
+
+          def create_epic_nodes(aggregate_records)
+            aggregate_records.each do |epic_id, aggregates|
+              next if aggregates.blank?
+
+              tree[epic_id] = EpicNode.new(epic_id, aggregates)
+            end
+
+            relate_parents_and_children
+          end
+
+          def relate_parents_and_children
+            tree.each do |_, node|
+              parent = tree[node.parent_id]
+              next if parent.nil?
+
+              parent.children << node
+            end
+          end
+
+          def aggregate_object(node)
+            if @facet == COUNT
+              node.aggregate_count
+            else
+              node.aggregate_weight_sum
+            end
+          end
+        end
+      end
+    end
+  end
+end

ee/lib/gitlab/graphql/loaders/bulk_epic_aggregate_loader.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Graphql
+    module Loaders
+      class BulkEpicAggregateLoader
+        include ::Gitlab::Graphql::Aggregations::Epics::Constants
+
+        MAXIMUM_LOADABLE = 100_001
+
+        attr_reader :target_epic_ids, :results
+
+        # This class retrieves each epic and its child epics recursively
+        # It allows us to recreate the epic tree structure in POROs
+        def initialize(epic_ids:)
+          @results = {}
+          @target_epic_ids = epic_ids
+        end
+
+        # rubocop: disable CodeReuse/ActiveRecord
+        def execute
+          return {} unless target_epic_ids
+
+          # We do a left outer join in order to capture epics with no issues
+          # This is so we can aggregate the epic counts for every epic
+          raw_results = ::Gitlab::ObjectHierarchy.new(Epic.where(id: target_epic_ids)).base_and_descendants
+            .left_joins(epic_issues: :issue)
+            .group("issues.state_id", "epics.id", "epics.iid", "epics.parent_id", "epics.state_id")
+            .select("epics.id, epics.iid, epics.parent_id, epics.state_id AS epic_state_id, issues.state_id AS issues_state_id, COUNT(issues) AS issues_count, SUM(COALESCE(issues.weight, 0)) AS issues_weight_sum")
+            .limit(MAXIMUM_LOADABLE)
+
+          raw_results = raw_results.map { |record| record.attributes.with_indifferent_access }
+
+          raise ArgumentError.new("There are too many records to load. Please select fewer epics or contact your administrator.") if raw_results.count == MAXIMUM_LOADABLE
+
+          @results = raw_results.group_by { |record| record[:id] }
+        end
+        # rubocop: enable CodeReuse/ActiveRecord
+      end
+    end
+  end
+end

ee/spec/lib/gitlab/graphql/aggregations/epics/epic_node_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Gitlab::Graphql::Aggregations::Epics::EpicNode do
+  include_context 'includes EpicAggregate constants'
+
+  let(:epic_id) { 34 }
+  let(:epic_iid) { 5 }
+
+  describe '#initialize' do
+    let(:fake_data) do
+      [
+          { iid: epic_iid, epic_state_id: epic_state_id, issues_count: 1, issues_weight_sum: 2, parent_id: parent_id, issues_state_id: OPENED_ISSUE_STATE },
+          { iid: epic_iid, epic_state_id: epic_state_id, issues_count: 2, issues_weight_sum: 2, parent_id: parent_id, issues_state_id: CLOSED_ISSUE_STATE }
+      ]
+    end
+
+    shared_examples 'setting attributes based on the first record' do |attributes|
+      let(:parent_id) { attributes[:parent_id] }
+      let(:epic_state_id) { attributes[:epic_state_id] }
+
+      it 'sets epic attributes based on the first record' do
+        new_node = described_class.new(epic_id, fake_data)
+
+        expect(new_node.parent_id).to eq parent_id
+        expect(new_node.epic_state_id).to eq epic_state_id
+      end
+    end
+
+    it_behaves_like 'setting attributes based on the first record', { epic_state_id: OPENED_EPIC_STATE, parent_id: nil }
+    it_behaves_like 'setting attributes based on the first record', { epic_state_id: CLOSED_EPIC_STATE, parent_id: 2 }
+  end
+
+  describe 'recursive totals' do
+    subject { described_class.new(epic_id, [{ parent_id: nil, epic_state_id: CLOSED_EPIC_STATE }]) }
+
+    before do
+      allow(subject).to receive(:epic_info_flat_list).and_return(flat_info)
+    end
+
+    context 'an epic with no child epics' do
+      context 'with no child issues', :aggregate_results do
+        let(:flat_info) { [] }
+
+        it 'has the correct aggregates', :aggregate_failures do
+          expect(subject).to have_aggregate(ISSUE_TYPE, COUNT, OPENED_ISSUE_STATE, 0)
+          expect(subject).to have_aggregate(ISSUE_TYPE, COUNT, CLOSED_ISSUE_STATE, 0)
+          expect(subject).to have_aggregate(ISSUE_TYPE, WEIGHT_SUM, OPENED_ISSUE_STATE, 0)
+          expect(subject).to have_aggregate(ISSUE_TYPE, WEIGHT_SUM, CLOSED_ISSUE_STATE, 0)
+
+          expect(subject).to have_aggregate(EPIC_TYPE, COUNT, OPENED_EPIC_STATE, 0)
+          expect(subject).to have_aggregate(EPIC_TYPE, COUNT, CLOSED_EPIC_STATE, 0)
+        end
+      end
+
+      context 'with an issue with 0 weight', :aggregate_results do
+        let(:flat_info) do
+          [
+              record_for(epic_id: epic_id, parent_id: nil, epic_state_id: CLOSED_EPIC_STATE, issues_state_id: OPENED_ISSUE_STATE, issues_count: 1, issues_weight_sum: 0)
+          ]
+        end
+
+        it 'has the correct aggregates', :aggregate_failures do
+          expect(subject).to have_aggregate(ISSUE_TYPE, COUNT, OPENED_ISSUE_STATE, 1)
+          expect(subject).to have_aggregate(ISSUE_TYPE, COUNT, CLOSED_ISSUE_STATE, 0)
+          expect(subject).to have_aggregate(ISSUE_TYPE, WEIGHT_SUM, OPENED_ISSUE_STATE, 0)
+          expect(subject).to have_aggregate(ISSUE_TYPE, WEIGHT_SUM, CLOSED_ISSUE_STATE, 0)
+
+          expect(subject).to have_aggregate(EPIC_TYPE, COUNT, OPENED_EPIC_STATE, 0)
+          expect(subject).to have_aggregate(EPIC_TYPE, COUNT, CLOSED_EPIC_STATE, 0)
+        end
+      end
+
+      context 'with an issue with nonzero weight' do
+        let(:flat_info) do
+          [
+              record_for(epic_id: epic_id, parent_id: nil, epic_state_id: CLOSED_EPIC_STATE, issues_state_id: OPENED_ISSUE_STATE, issues_count: 1, issues_weight_sum: 2)
+          ]
+        end
+
+        it 'has the correct aggregates', :aggregate_failures do
+          expect(subject).to have_aggregate(ISSUE_TYPE, COUNT, OPENED_ISSUE_STATE, 1)
+          expect(subject).to have_aggregate(ISSUE_TYPE, COUNT, CLOSED_ISSUE_STATE, 0)
+          expect(subject).to have_aggregate(ISSUE_TYPE, WEIGHT_SUM, OPENED_ISSUE_STATE, 2)
+          expect(subject).to have_aggregate(ISSUE_TYPE, WEIGHT_SUM, CLOSED_ISSUE_STATE, 0)
+
+          expect(subject).to have_aggregate(EPIC_TYPE, COUNT, OPENED_EPIC_STATE, 0)
+          expect(subject).to have_aggregate(EPIC_TYPE, COUNT, CLOSED_EPIC_STATE, 0)
+        end
+      end
+    end
+
+    context 'an epic with child epics' do
+      let(:child_epic_id) { 45 }
+      let(:child_epic_node) { described_class.new(child_epic_id, child_flat_info) }
+      let(:flat_info) do
+        [
+            record_for(epic_id: epic_id, parent_id: nil, epic_state_id: OPENED_EPIC_STATE, issues_state_id: OPENED_ISSUE_STATE, issues_count: 0, issues_weight_sum: 0)
+        ]
+      end
+
+      before do
+        subject.children << child_epic_node
+      end
+
+      context 'with a child that has issues of nonzero weight' do
+        let(:child_flat_info) do
+          [
+            record_for(epic_id: epic_id, parent_id: nil, epic_state_id: OPENED_EPIC_STATE, issues_state_id: OPENED_ISSUE_STATE, issues_count: 1, issues_weight_sum: 2)
+          ]
+        end
+
+        it 'has the correct aggregates', :aggregate_failures do
+          expect(subject).to have_aggregate(ISSUE_TYPE, COUNT, OPENED_ISSUE_STATE, 1)
+          expect(subject).to have_aggregate(ISSUE_TYPE, COUNT, CLOSED_ISSUE_STATE, 0)
+          expect(subject).to have_aggregate(ISSUE_TYPE, WEIGHT_SUM, OPENED_ISSUE_STATE, 2)
+          expect(subject).to have_aggregate(ISSUE_TYPE, WEIGHT_SUM, CLOSED_ISSUE_STATE, 0)
+          expect(subject).to have_aggregate(EPIC_TYPE, COUNT, OPENED_EPIC_STATE, 1)
+          expect(subject).to have_aggregate(EPIC_TYPE, COUNT, CLOSED_EPIC_STATE, 0)
+        end
+      end
+    end
+  end
+
+  def record_for(epic_id:, parent_id:, epic_state_id:, issues_state_id:, issues_count:, issues_weight_sum:)
+    {
+      epic_id: epic_id,
+      issues_count: issues_count,
+      issues_weight_sum: issues_weight_sum,
+      parent_id: parent_id,
+      issues_state_id: issues_state_id,
+      epic_state_id: epic_state_id
+    }
+  end
+end

ee/spec/lib/gitlab/graphql/aggregations/epics/lazy_epic_aggregate_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Gitlab::Graphql::Aggregations::Epics::LazyEpicAggregate do
+  include_context 'includes EpicAggregate constants'
+
+  let(:query_ctx) do
+    {}
+  end
+  let(:epic_id) { 37 }
+  let(:epic_iid) { 18 }
+  let(:child_epic_id) { 38 }
+
+  describe '#initialize' do
+    it 'requires either :weight_sum or :count as a facet', :aggregate_failures do
+      expect { described_class.new(query_ctx, epic_id, :nonsense) }.to raise_error(ArgumentError, /Invalid aggregate facet/)
+      expect { described_class.new(query_ctx, epic_id, nil) }.to raise_error(ArgumentError, /No aggregate facet/)
+      expect { described_class.new(query_ctx, epic_id, "") }.to raise_error(ArgumentError, /No aggregate facet/)
+    end
+
+    context 'with valid facets :weight_sum or :count' do
+      specify 'as a symbol', :aggregate_failures do
+        [WEIGHT_SUM, COUNT].each do |valid_facet|
+          expect { described_class.new(query_ctx, epic_id, valid_facet) }.not_to raise_error
+        end
+      end
+
+      specify 'as a string', :aggregate_failures do
+        %w(weight_sum count).each do |valid_facet|
+          expect { described_class.new(query_ctx, epic_id, valid_facet) }.not_to raise_error
+        end
+      end
+    end
+
+    it 'adds the epic_id to lazy state' do
+      described_class.new(query_ctx, epic_id, COUNT)
+
+      expect(query_ctx[:lazy_epic_aggregate][:pending_ids]).to match [epic_id]
+    end
+  end
+
+  describe '#epic_aggregate' do
+    let(:single_record) do
+      { iid: 6, issues_count: 4, issues_weight_sum: 9, parent_id: nil, issues_state_id: OPENED_ISSUE_STATE, epic_state_id: OPENED_EPIC_STATE }
+    end
+    let(:epic_info_node) { Gitlab::Graphql::Aggregations::Epics::EpicNode.new(epic_id, [single_record] ) }
+
+    subject { described_class.new(query_ctx, epic_id, COUNT) }
+
+    before do
+      subject.instance_variable_set(:@lazy_state, fake_state)
+    end
+
+    context 'if the record has already been loaded' do
+      let(:fake_state) do
+        { pending_ids: Set.new, tree: { epic_id => epic_info_node } }
+      end
+
+      it 'does not make the query again' do
+        expect(epic_info_node).to receive(:aggregate_count)
+        expect(Gitlab::Graphql::Loaders::BulkEpicAggregateLoader).not_to receive(:new)
+
+        subject.epic_aggregate
+      end
+    end
+
+    context 'if the record has not been loaded' do
+      let(:other_epic_id) { 39 }
+      let(:fake_state) do
+        { pending_ids: Set.new([epic_id, child_epic_id]), tree: {} }
+      end
+      let(:fake_data) do
+        {
+            epic_id => [{ epic_state_id: OPENED_EPIC_STATE, issues_count: 2, issues_weight_sum: 5, parent_id: nil, issues_state_id: OPENED_ISSUE_STATE }],
+            child_epic_id => [{ epic_state_id: CLOSED_EPIC_STATE, issues_count: 4, issues_weight_sum: 17, parent_id: epic_id, issues_state_id: CLOSED_ISSUE_STATE }],
+            other_epic_id => [{ epic_state_id: OPENED_EPIC_STATE, issues_count: 0, issues_weight_sum: 0, parent_id: nil, issues_state_id: nil }] # represents an epic with no parent and no issues
+        }
+      end
+
+      before do
+        expect_next_instance_of(Gitlab::Graphql::Loaders::BulkEpicAggregateLoader) do |loader|
+          expect(loader).to receive(:execute).and_return(fake_data)
+        end
+      end
+
+      it 'clears the pending IDs' do
+        subject.epic_aggregate
+
+        lazy_state = subject.instance_variable_get(:@lazy_state)
+
+        expect(lazy_state[:pending_ids]).to be_empty
+      end
+
+      it 'creates the parent-child associations', :aggregate_failures do
+        subject.epic_aggregate
+
+        expect(tree[child_epic_id].parent_id).to eq epic_id
+        expect(tree[epic_id].children.map(&:epic_id)).to match_array([child_epic_id])
+      end
+
+      context 'for a parent-child relationship' do
+        it 'assembles recursive sums for the parent', :aggregate_failures do
+          subject.epic_aggregate
+
+          expect(tree[epic_id]).to have_aggregate(ISSUE_TYPE, COUNT, OPENED_ISSUE_STATE, 2)
+          expect(tree[epic_id]).to have_aggregate(ISSUE_TYPE, COUNT, CLOSED_ISSUE_STATE, 4)
+          expect(tree[epic_id]).to have_aggregate(ISSUE_TYPE, WEIGHT_SUM, OPENED_ISSUE_STATE, 5)
+          expect(tree[epic_id]).to have_aggregate(ISSUE_TYPE, WEIGHT_SUM, CLOSED_ISSUE_STATE, 17)
+          expect(tree[epic_id]).to have_aggregate(EPIC_TYPE, COUNT, CLOSED_EPIC_STATE, 1)
+        end
+      end
+
+      context 'for a standalone epic with no issues' do
+        it 'assembles recursive sums', :aggregate_failures do
+          subject.epic_aggregate
+
+          expect(tree[other_epic_id]).to have_aggregate(ISSUE_TYPE, COUNT, OPENED_ISSUE_STATE, 0)
+          expect(tree[other_epic_id]).to have_aggregate(ISSUE_TYPE, COUNT, CLOSED_ISSUE_STATE, 0)
+          expect(tree[other_epic_id]).to have_aggregate(ISSUE_TYPE, WEIGHT_SUM, OPENED_ISSUE_STATE, 0)
+          expect(tree[other_epic_id]).to have_aggregate(ISSUE_TYPE, WEIGHT_SUM, CLOSED_ISSUE_STATE, 0)
+          expect(tree[other_epic_id]).to have_aggregate(EPIC_TYPE, COUNT, CLOSED_EPIC_STATE, 0)
+        end
+      end
+    end
+  end
+
+  def tree
+    lazy_state = subject.instance_variable_get(:@lazy_state)
+    lazy_state[:tree]
+  end
+end

ee/spec/lib/gitlab/graphql/loaders/bulk_epic_aggregate_loader_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Gitlab::Graphql::Loaders::BulkEpicAggregateLoader do
+  include_context 'includes EpicAggregate constants'
+
+  let_it_be(:group) { create(:group, :public) }
+  let_it_be(:subgroup) { create(:group, :private, parent: group)}
+
+  let_it_be(:project) { create(:project, namespace: group) }
+  let_it_be(:subproject) { create(:project, namespace: subgroup) }
+
+  let_it_be(:parent_epic) { create(:epic, group: group, title: 'parent epic') }
+  let_it_be(:epic_with_issues) { create(:epic, group: subgroup, parent: parent_epic, state: :opened, title: 'epic with issues') }
+
+  # closed, no issues
+  let_it_be(:epic_without_issues) { create(:epic, group: subgroup, parent: parent_epic, state: :closed, title: 'epic without issues') }
+
+  # open, public
+  let_it_be(:issue1) { create(:issue, project: project, weight: 1, state: :opened) }
+  let_it_be(:issue2) { create(:issue, project: project, weight: 1, state: :opened) }
+  # closed
+  let_it_be(:issue3) { create(:issue, project: project, weight: 1, state: :closed) }
+  let_it_be(:issue4) { create(:issue, project: project, weight: 1, state: :closed) }
+  # confidential
+  let_it_be(:issue5) { create(:issue, project: project, weight: 1, confidential: true, state: :opened) }
+  let_it_be(:issue6) { create(:issue, project: project, weight: 1, confidential: true, state: :opened) }
+  # in private project, private subgroup
+  let_it_be(:issue7) { create(:issue, project: subproject, weight: 1, state: :opened) }
+  let_it_be(:issue8) { create(:issue, project: subproject, weight: 1, state: :opened) }
+  # private project, confidential, private subgroup
+  let_it_be(:issue9) { create(:issue, project: subproject, weight: 1, confidential: true, state: :opened) }
+  let_it_be(:issue10) { create(:issue, project: subproject, weight: 1, confidential: true, state: :opened) }
+  # nil weight doesn't break it
+  let_it_be(:issue11) { create(:issue, project: project, weight: 0, state: :opened) }
+  let_it_be(:issue12) { create(:issue, project: project, weight: nil, state: :opened) }
+
+  let_it_be(:epic_issue1) { create(:epic_issue, epic: parent_epic, issue: issue1) }
+  let_it_be(:epic_issue2) { create(:epic_issue, epic: epic_with_issues, issue: issue2) }
+  let_it_be(:epic_issue3) { create(:epic_issue, epic: parent_epic, issue: issue3) }
+  let_it_be(:epic_issue4) { create(:epic_issue, epic: epic_with_issues, issue: issue4) }
+  let_it_be(:epic_issue5) { create(:epic_issue, epic: parent_epic, issue: issue5) }
+  let_it_be(:epic_issue6) { create(:epic_issue, epic: epic_with_issues, issue: issue6) }
+  let_it_be(:epic_issue7) { create(:epic_issue, epic: parent_epic, issue: issue7) }
+  let_it_be(:epic_issue8) { create(:epic_issue, epic: epic_with_issues, issue: issue8) }
+  let_it_be(:epic_issue9) { create(:epic_issue, epic: parent_epic, issue: issue9) }
+  let_it_be(:epic_issue10) { create(:epic_issue, epic: epic_with_issues, issue: issue10) }
+  let_it_be(:epic_issue11) { create(:epic_issue, epic: parent_epic, issue: issue11) }
+  let_it_be(:epic_issue12) { create(:epic_issue, epic: epic_with_issues, issue: issue12) }
+
+  subject { described_class.new(epic_ids: target_ids) }
+
+  before do
+    stub_licensed_features(epics: true)
+  end
+
+  context 'when epic ids with issues is provided' do
+    let(:target_ids) { parent_epic.id }
+
+    it 'sums all the weights, even confidential, or in private groups' do
+      expected_result = {
+        parent_epic.id => [
+            result_for(parent_epic, issues_state: OPENED_ISSUE_STATE, issues_count: 5, issues_weight_sum: 4),
+            result_for(parent_epic, issues_state: CLOSED_ISSUE_STATE, issues_count: 1, issues_weight_sum: 1)
+          ],
+        epic_with_issues.id => [
+          result_for(epic_with_issues, issues_state: OPENED_ISSUE_STATE, issues_count: 5, issues_weight_sum: 4),
+          result_for(epic_with_issues, issues_state: CLOSED_ISSUE_STATE, issues_count: 1, issues_weight_sum: 1)
+        ],
+        epic_without_issues.id => [
+          result_for(epic_without_issues, issues_state: nil, issues_count: 0, issues_weight_sum: 0)
+        ]
+      }
+
+      result = subject.execute
+
+      expected_result.each do |epic_id, records|
+        expect(result[epic_id]).to match_array records
+      end
+    end
+
+    it 'contains results for all epics, even if they do not have issues' do
+      result = subject.execute
+
+      # epic_without_issues is included, even if it has none
+      expect(result.keys).to match_array([parent_epic.id, epic_with_issues.id, epic_without_issues.id])
+    end
+
+    it 'errors when the number of retrieved records exceeds the maximum' do
+      stub_const("Gitlab::Graphql::Loaders::BulkEpicAggregateLoader::MAXIMUM_LOADABLE", 1)
+
+      expect { subject.execute }.to raise_error(ArgumentError, /too many records/)
+    end
+
+    context 'testing for a single database query' do
+      it 'does not repeat database queries for subepics' do
+        recorder = ActiveRecord::QueryRecorder.new { described_class.new(epic_ids: epic_with_issues.id).execute }
+
+        # this one has sub-epics, but there should still only be one query
+        expect { described_class.new(epic_ids: [parent_epic.id, epic_with_issues.id]).execute }.not_to exceed_query_limit(recorder)
+      end
+
+      it 'avoids N+1' do
+        recorder = ActiveRecord::QueryRecorder.new { described_class.new(epic_ids: epic_with_issues.id).execute }
+
+        expect { described_class.new(epic_ids: [epic_with_issues.id, parent_epic.id]).execute }.not_to exceed_query_limit(recorder)
+      end
+    end
+  end
+
+  context 'when an epic without issues is provided' do
+    let(:target_ids) { epic_without_issues.id }
+
+    it 'returns a placeholder' do
+      expected_result = [
+        result_for(epic_without_issues, issues_state: nil, issues_count: 0, issues_weight_sum: 0)
+      ]
+
+      actual_result = subject.execute
+
+      expect(actual_result[epic_without_issues.id]).to match_array(expected_result)
+    end
+  end
+
+  context 'when no epic ids are provided' do
+    [nil, [], ""].each do |empty_arg|
+      let(:target_ids) { empty_arg }
+
+      it 'returns an empty set' do
+        expect(subject.execute).to eq({})
+      end
+    end
+  end
+
+  def result_for(epic, issues_state:, issues_count:, issues_weight_sum:)
+    {
+      id: epic.id,
+      iid: epic.iid,
+      issues_count: issues_count,
+      issues_weight_sum: issues_weight_sum,
+      parent_id: epic.parent_id,
+      issues_state_id: issues_state,
+      epic_state_id: Epic.available_states[epic.state_id]
+    }.stringify_keys
+  end
+end

ee/spec/requests/api/graphql/group/epic/epic_aggregate_query_spec.rb

@@ -15,18 +15,18 @@ describe 'Epic aggregates (count and weight)' do
 
   let(:epic_aggregates_query) do
     <<~QUERY
-        nodes {
-          descendantWeightSum {
-            openedIssues
-            closedIssues
-          }
-          descendantCounts {
-            openedEpics
-            closedEpics
-            openedIssues
-            closedIssues
-          }
-        }
+    nodes {
+      descendantWeightSum {
+        openedIssues
+        closedIssues
+      }
+      descendantCounts {
+        openedEpics
+        closedEpics
+        openedIssues
+        closedIssues
+      }
+    }
     QUERY
   end
 
@@ -34,68 +34,126 @@ describe 'Epic aggregates (count and weight)' do
     stub_licensed_features(epics: true)
   end
 
-  context 'with feature flag enabled' do
+  context 'count and weight totals' do
+    subject { graphql_data.dig('group', 'epics', 'nodes') }
+
+    let_it_be(:subgroup) { create(:group, :private, parent: group)}
+    let_it_be(:subsubgroup) { create(:group, :private, parent: subgroup)}
+
+    let_it_be(:project) { create(:project, namespace: group) }
+
+    let_it_be(:epic_with_issues) { create(:epic, id: 2, group: subgroup, parent: parent_epic, title: 'epic with issues') }
+    let_it_be(:epic_without_issues) { create(:epic, :closed, id: 3, group: subgroup, parent: parent_epic, title: 'epic without issues') }
+    let_it_be(:closed_epic) { create(:epic, :closed, id: 4, group: subgroup, parent: parent_epic, title: 'closed epic') }
+
+    let_it_be(:issue1) { create(:issue, project: project, weight: 5, state: :opened) }
+    let_it_be(:issue2) { create(:issue, project: project, weight: 7, state: :closed) }
+
+    let_it_be(:epic_issue1) { create(:epic_issue, epic: epic_with_issues, issue: issue1) }
+    let_it_be(:epic_issue2) { create(:epic_issue, epic: epic_with_issues, issue: issue2) }
+
     before do
-      stub_feature_flags(unfiltered_epic_aggregates: true)
+      group.add_developer(current_user)
+      post_graphql(query, current_user: current_user)
     end
 
-    it 'returns a placeholder with -1 weights and does not error' do
-      post_graphql(query, current_user: current_user)
+    shared_examples 'counts properly' do
+      it_behaves_like 'a working graphql query'
 
-      actual_result = graphql_data.dig('group', 'epics', 'nodes').first
-      expected_result = {
-        "descendantWeightSum" => {
-          "openedIssues" => -1,
-          "closedIssues" => -1
+      it 'returns the epic counts' do
+        epic_count_result = {
+          "openedEpics" => 1,
+          "closedEpics" => 2
         }
-      }
 
-      expect(actual_result).to include expected_result
-    end
-  end
+        is_expected.to include(
+          a_hash_including('descendantCounts' => a_hash_including(epic_count_result))
+        )
+      end
 
-  context 'with feature flag disabled' do
-    before do
-      stub_feature_flags(unfiltered_epic_aggregates: false)
+      it 'returns the issue counts' do
+        issue_count_result = {
+          "openedIssues" => 1,
+          "closedIssues" => 1
+        }
+
+        is_expected.to include(
+          a_hash_including('descendantCounts' => a_hash_including(issue_count_result))
+        )
+      end
     end
 
-    context 'when requesting counts' do
-      let(:epic_aggregates_query) do
-        <<~QUERY
-        nodes {
-          descendantCounts {
-            openedEpics
-            closedEpics
-            openedIssues
-            closedIssues
-          }
-        }
-        QUERY
+    context 'with feature flag enabled' do
+      before do
+        stub_feature_flags(unfiltered_epic_aggregates: true)
       end
 
-      it 'uses the DescendantCountService' do
-        expect(Epics::DescendantCountService).to receive(:new)
+      it 'uses the LazyEpicAggregate service' do
+        # one for count, one for weight_sum, even though the share the same tree state as part of the context
+        expect(Gitlab::Graphql::Aggregations::Epics::LazyEpicAggregate).to receive(:new).twice
 
         post_graphql(query, current_user: current_user)
       end
+
+      it_behaves_like 'counts properly'
+
+      it 'returns the weights' do
+        descendant_weight_result = {
+            "openedIssues" => 5,
+            "closedIssues" => 7
+        }
+
+        is_expected.to include(
+          a_hash_including('descendantWeightSum' => a_hash_including(descendant_weight_result))
+        )
+      end
     end
 
-    context 'when requesting weights' do
-      let(:epic_aggregates_query) do
-        <<~QUERY
-        nodes {
-          descendantWeightSum {
-            openedIssues
-            closedIssues
+    context 'with feature flag disabled' do
+      before do
+        stub_feature_flags(unfiltered_epic_aggregates: false)
+      end
+
+      context 'when requesting counts' do
+        let(:epic_aggregates_query) do
+          <<~QUERY
+          nodes {
+            descendantCounts {
+              openedEpics
+              closedEpics
+              openedIssues
+              closedIssues
+            }
           }
-        }
-        QUERY
+          QUERY
+        end
+
+        it 'uses the DescendantCountService' do
+          expect(Epics::DescendantCountService).to receive(:new)
+
+          post_graphql(query, current_user: current_user)
+        end
+
+        it_behaves_like 'counts properly'
       end
 
-      it 'returns an error' do
-        post_graphql(query, current_user: current_user)
+      context 'when requesting weights' do
+        let(:epic_aggregates_query) do
+          <<~QUERY
+          nodes {
+            descendantWeightSum {
+              openedIssues
+              closedIssues
+            }
+          }
+          QUERY
+        end
+
+        it 'returns an error' do
+          post_graphql(query, current_user: current_user)
 
-        expect_graphql_errors_to_include /Field 'descendantWeightSum' doesn't exist on type 'Epic/
+          expect_graphql_errors_to_include /Field 'descendantWeightSum' doesn't exist on type 'Epic/
+        end
       end
     end
   end

ee/spec/support/matchers/ee/epic_aggregate_matchers.rb

+# frozen_string_literal: true
+
+RSpec::Matchers.define :have_aggregate do |type, facet, state, expected_value|
+  match do |epic_node_result|
+    aggregate_object = epic_node_result.public_send(:"aggregate_#{facet}")
+    expect(aggregate_object.public_send(method_name(type, state))).to eq expected_value
+  end
+
+  failure_message do |epic_node_result|
+    aggregate_object = epic_node_result.public_send(:"aggregate_#{facet}")
+    aggregate_method = method_name(type, state)
+    "Epic node with id #{epic_node_result.epic_id} called #{aggregate_method} on aggregate object. Value was expected to be #{expected_value} but was #{aggregate_object.send(aggregate_method)}."
+  end
+
+  def method_name(type, state)
+    if type == ISSUE_TYPE
+      return :opened_issues if state == OPENED_ISSUE_STATE
+
+      :closed_issues
+    elsif type == EPIC_TYPE
+      return :opened_epics if state == OPENED_EPIC_STATE
+
+      :closed_epics
+    end
+  end
+end

ee/spec/support/shared_contexts/epic_aggregate_constants.rb

+# frozen_string_literal: true
+
+shared_context 'includes EpicAggregate constants' do
+  EPIC_TYPE = Gitlab::Graphql::Aggregations::Epics::Constants::EPIC_TYPE
+  ISSUE_TYPE = Gitlab::Graphql::Aggregations::Epics::Constants::ISSUE_TYPE
+
+  OPENED_EPIC_STATE = Gitlab::Graphql::Aggregations::Epics::Constants::OPENED_EPIC_STATE
+  CLOSED_EPIC_STATE = Gitlab::Graphql::Aggregations::Epics::Constants::CLOSED_EPIC_STATE
+  OPENED_ISSUE_STATE = Gitlab::Graphql::Aggregations::Epics::Constants::OPENED_ISSUE_STATE
+  CLOSED_ISSUE_STATE = Gitlab::Graphql::Aggregations::Epics::Constants::CLOSED_ISSUE_STATE
+
+  WEIGHT_SUM = Gitlab::Graphql::Aggregations::Epics::Constants::WEIGHT_SUM
+  COUNT = Gitlab::Graphql::Aggregations::Epics::Constants::COUNT
+end