Add lookahead optimisations for merge requests

This uses lookahead to preload associations we know we are going to
need.

To ensure this works, the keyset connection is updated to support
pre-loaded associations This enables results sets that are already in
memory to be efficiently handled without making extra unnecessary
queries.

app/graphql/resolvers/concerns/looks_ahead.rb

+# frozen_string_literal: true
+
+module LooksAhead
+  extend ActiveSupport::Concern
+
+  FEATURE_FLAG = :graphql_lookahead_support
+
+  included do
+    attr_accessor :lookahead
+  end
+
+  def resolve(**args)
+    self.lookahead = args.delete(:lookahead)
+
+    resolve_with_lookahead(**args)
+  end
+
+  def apply_lookahead(query)
+    return query unless Feature.enabled?(FEATURE_FLAG)
+
+    selection = node_selection
+
+    includes = preloads.each.flat_map do |name, requirements|
+      selection&.selects?(name) ? requirements : []
+    end
+    preloads = (unconditional_includes + includes).uniq
+
+    return query if preloads.empty?
+
+    query.preload(*preloads) # rubocop: disable CodeReuse/ActiveRecord
+  end
+
+  private
+
+  def unconditional_includes
+    []
+  end
+
+  def preloads
+    {}
+  end
+
+  def node_selection
+    return unless lookahead
+
+    if lookahead.selects?(:nodes)
+      lookahead.selection(:nodes)
+    elsif lookahead.selects?(:edges)
+      lookahead.selection(:edges).selection(:nodes)
+    end
+  end
+end

app/graphql/resolvers/concerns/resolves_merge_requests.rb

@@ -4,12 +4,13 @@
 # that `MergeRequestsFinder` can handle, so you may need to use aliasing.
 module ResolvesMergeRequests
   extend ActiveSupport::Concern
+  include LooksAhead
 
   included do
     type Types::MergeRequestType, null: true
   end
 
-  def resolve(**args)
+  def resolve_with_lookahead(**args)
     args[:iids] = Array.wrap(args[:iids]) if args[:iids]
     args.compact!
 
@@ -18,7 +19,7 @@ module ResolvesMergeRequests
     else
       args[:project_id] ||= project
 
-      MergeRequestsFinder.new(current_user, args).execute
+      apply_lookahead(MergeRequestsFinder.new(current_user, args).execute)
     end.then(&(single? ? :first : :itself))
   end
 
@@ -41,10 +42,26 @@ module ResolvesMergeRequests
   # rubocop: disable CodeReuse/ActiveRecord
   def batch_load(iid)
     BatchLoader::GraphQL.for(iid.to_s).batch(key: project) do |iids, loader, args|
-      args[:key].merge_requests.where(iid: iids).each do |mr|
+      query = args[:key].merge_requests.where(iid: iids)
+
+      apply_lookahead(query).each do |mr|
         loader.call(mr.iid.to_s, mr)
       end
     end
   end
   # rubocop: enable CodeReuse/ActiveRecord
+
+  def unconditional_includes
+    [:target_project]
+  end
+
+  def preloads
+    {
+      assignees: [:assignees],
+      labels: [:labels],
+      author: [:author],
+      milestone: [:milestone],
+      head_pipeline: [:merge_request_diff, { head_pipeline: [:merge_request] }]
+    }
+  end
 end

app/graphql/types/project_type.rb

@@ -125,6 +125,7 @@ module Types
           Types::MergeRequestType.connection_type,
           null: true,
           description: 'Merge requests of the project',
+          extras: [:lookahead],
           resolver: Resolvers::MergeRequestsResolver
 
     field :merge_request,

changelogs/unreleased/ajk-gql-lookahead.yml

+---
+title: Add GraphQL lookahead support
+merge_request: 32373
+author:
+type: performance

doc/development/api_graphql_styleguide.md

@@ -644,6 +644,59 @@ abstractions Resolvers should not be considered re-usable, finders are to be
 preferred), remember to call the `ready?` method and check the boolean flag
 before calling `resolve`! An example can be seen in our [`GraphQLHelpers`](https://gitlab.com/gitlab-org/gitlab/-/blob/2d395f32d2efbb713f7bc861f96147a2a67e92f2/spec/support/helpers/graphql_helpers.rb#L20-27).
 
+### Look-Ahead
+
+The full query is known in advance during execution, which means we can make use
+of [lookahead](https://graphql-ruby.org/queries/lookahead.html) to optimize our
+queries, and batch load associations we know we will need. Consider adding
+lookahead support in your resolvers to avoid `N+1` performance issues.
+
+To enable support for common lookahead use-cases (pre-loading associations when
+child fields are requested), you can
+include [`LooksAhead`](https://gitlab.com/gitlab-org/gitlab/-/blob/master/app/graphql/resolvers/concerns/looks_ahead.rb). For example:
+
+```ruby
+# Assuming a model `MyThing` with attributes `[child_attribute, other_attribute, nested]`,
+# where nested has an attribute named `included_attribute`.
+class MyThingResolver < BaseResolver
+  include LooksAhead
+
+  # Rather than defining `resolve(**args)`, we implement: `resolve_with_lookahead(**args)`
+  def resolve_with_lookahead(**args)
+    apply_lookahead(MyThingFinder.new(current_user).execute)
+  end
+
+  # We list things that should always be preloaded:
+  # For example, if child_attribute is always needed (during authorization
+  # perhaps), then we can include it here.
+  def unconditional_includes
+    [:child_attribute]
+  end
+
+  # We list things that should be included if a certain field is selected:
+  def preloads
+    {
+        field_one: [:other_attribute],
+        field_two: [{ nested: [:included_attribute] }]
+    }
+  end
+end
+```
+
+The final thing that is needed is that every field that uses this resolver needs
+to advertise the need for lookahead:
+
+```ruby
+  # in ParentType
+  field :my_things, MyThingType.connection_type, null: true,
+        extras: [:lookahead], # Necessary
+        resolver: MyThingResolver,
+        description: 'My things'
+```
+
+For an example of real world use, please
+see [`ResolvesMergeRequests`](https://gitlab.com/gitlab-org/gitlab/-/blob/master/app/graphql/resolvers/concerns/resolves_merge_requests.rb).
+
 ## Mutations
 
 Mutations are used to change any stored values, or to trigger

lib/gitlab/graphql/pagination/keyset/connection.rb

@@ -82,7 +82,7 @@ module Gitlab
           def sliced_nodes
             @sliced_nodes ||=
               begin
-                OrderInfo.validate_ordering(ordered_items, order_list)
+                OrderInfo.validate_ordering(ordered_items, order_list) unless loaded?(ordered_items)
 
                 sliced = ordered_items
                 sliced = slice_nodes(sliced, before, :before) if before.present?
@@ -113,16 +113,14 @@ module Gitlab
                 # grab one more than we need
                 paginated_nodes = sliced_nodes.last(limit_value + 1)
 
-                if paginated_nodes.count > limit_value
-                  # there is an extra node, so there is a previous page
-                  @has_previous_page = true
-                  paginated_nodes = paginated_nodes.last(limit_value)
-                end
+                # there is an extra node, so there is a previous page
+                @has_previous_page = paginated_nodes.count > limit_value
+                @has_previous_page ? paginated_nodes.last(limit_value) : paginated_nodes
+              elsif loaded?(sliced_nodes)
+                sliced_nodes.take(limit_value) # rubocop: disable CodeReuse/ActiveRecord
               else
-                paginated_nodes = sliced_nodes.limit(limit_value) # rubocop: disable CodeReuse/ActiveRecord
+                sliced_nodes.limit(limit_value) # rubocop: disable CodeReuse/ActiveRecord
               end
-
-              paginated_nodes
             end
           end
 
@@ -141,6 +139,15 @@ module Gitlab
             @limit_value ||= [first, last, max_page_size].compact.min
           end
 
+          def loaded?(items)
+            case items
+            when Array
+              true
+            else
+              items.loaded?
+            end
+          end
+
           def ordered_items
             strong_memoize(:ordered_items) do
               unless items.primary_key.present?
@@ -149,6 +156,16 @@ module Gitlab
 
               list = OrderInfo.build_order_list(items)
 
+              if loaded?(items)
+                @order_list = list.presence || [items.primary_key]
+
+                # already sorted, or trivially sorted
+                next items if list.present? || items.size <= 1
+
+                pkey = items.primary_key.to_sym
+                next items.sort_by { |item| item[pkey] }.reverse
+              end
+
               # ensure there is a primary key ordering
               if list&.last&.attribute_name != items.primary_key
                 items.order(arel_table[items.primary_key].desc) # rubocop: disable CodeReuse/ActiveRecord

spec/factories/merge_requests.rb

@@ -109,6 +109,17 @@ FactoryBot.define do
       end
     end
 
+    trait :with_head_pipeline do
+      after(:build) do |merge_request|
+        merge_request.head_pipeline = build(
+          :ci_pipeline,
+          :running,
+          project: merge_request.source_project,
+          ref: merge_request.source_branch,
+          sha: merge_request.diff_head_sha)
+      end
+    end
+
     trait :with_test_reports do
       after(:build) do |merge_request|
         merge_request.head_pipeline = build(

spec/graphql/resolvers/concerns/looks_ahead_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe LooksAhead do
+  include GraphqlHelpers
+
+  let_it_be(:the_user) { create(:user) }
+  let_it_be(:label_a) { create(:label) }
+  let_it_be(:label_b) { create(:label) }
+  let_it_be(:issue_a) { create(:issue, author: the_user, labels: [label_a, label_b]) }
+  let_it_be(:issue_b) { create(:issue, author: the_user, labels: [label_a]) }
+  let_it_be(:issue_c) { create(:issue, author: the_user, labels: [label_b]) }
+
+  # Simplified schema to test lookahead
+  let_it_be(:schema) do
+    issues_resolver = Class.new(Resolvers::BaseResolver) do
+      include LooksAhead
+
+      def resolve_with_lookahead(**args)
+        apply_lookahead(object.issues)
+      end
+
+      def preloads
+        { labels: [:labels] }
+      end
+    end
+
+    label = Class.new(GraphQL::Schema::Object) do
+      graphql_name 'Label'
+      field :id, Integer, null: false
+    end
+    issue = Class.new(GraphQL::Schema::Object) do
+      graphql_name 'Issue'
+      field :title, String, null: true
+      field :labels, label.connection_type, null: true
+    end
+    user = Class.new(GraphQL::Schema::Object) do
+      graphql_name 'User'
+      field :name, String, null: true
+      field :issues, issue.connection_type,
+        null: true
+      field :issues_with_lookahead, issue.connection_type,
+        extras: [:lookahead],
+        resolver: issues_resolver,
+        null: true
+    end
+
+    Class.new(GraphQL::Schema) do
+      query(Class.new(GraphQL::Schema::Object) do
+        graphql_name 'Query'
+        field :find_user, user, null: true do
+          argument :username, String, required: true
+        end
+
+        def find_user(username:)
+          context[:user_db].find { |u| u.username == username }
+        end
+      end)
+    end
+  end
+
+  def query(doc = document)
+    GraphQL::Query.new(schema,
+                       document: doc,
+                       context: { user_db: [the_user] },
+                       variables: { username: the_user.username })
+  end
+
+  let(:document) do
+    GraphQL.parse <<-GRAPHQL
+    query($username: String!){
+      findUser(username: $username) {
+        name
+        issues {
+          nodes {
+            title
+            labels { nodes { id } }
+          }
+        }
+        issuesWithLookahead {
+          nodes {
+            title
+            labels { nodes { id } }
+          }
+        }
+      }
+    }
+    GRAPHQL
+  end
+
+  def run_query(gql_query)
+    query(GraphQL.parse(gql_query)).result
+  end
+
+  shared_examples 'a working query on the test schema' do
+    it 'has a good test setup', :aggregate_failures do
+      expected_label_ids = [label_a, label_b].cycle.take(4).map(&:id)
+      issue_titles = [issue_a, issue_b, issue_c].map(&:title)
+
+      res = query.result
+
+      expect(res['errors']).to be_blank
+      expect(res.dig('data', 'findUser', 'name')).to eq(the_user.name)
+      %w(issues issuesWithLookahead).each do |field|
+        expect(all_issue_titles(res, field)).to match_array(issue_titles)
+        expect(all_label_ids(res, field)).to match_array(expected_label_ids)
+      end
+    end
+  end
+
+  it_behaves_like 'a working query on the test schema'
+
+  it 'preloads labels on issues' do
+    expect(the_user.issues).to receive(:preload).with(:labels)
+
+    query.result
+  end
+
+  context 'the feature flag is off' do
+    before do
+      stub_feature_flags(described_class::FEATURE_FLAG => false)
+    end
+
+    it_behaves_like 'a working query on the test schema'
+
+    it 'does not preload labels on issues' do
+      expect(the_user.issues).not_to receive(:preload).with(:labels)
+
+      query.result
+    end
+  end
+
+  it 'issues fewer queries than the naive approach' do
+    the_user.reload # ensure no attributes are loaded before we begin
+    naive = <<-GQL
+    query($username: String!){
+      findUser(username: $username) {
+        name
+        issues {
+          nodes {
+            labels { nodes { id } }
+          }
+        }
+      }
+    }
+    GQL
+    with_lookahead = <<-GQL
+    query($username: String!){
+      findUser(username: $username) {
+        name
+        issuesWithLookahead {
+          nodes {
+            labels { nodes { id } }
+          }
+        }
+      }
+    }
+    GQL
+
+    expect { run_query(with_lookahead) }.to issue_fewer_queries_than { run_query(naive) }
+  end
+
+  private
+
+  def all_label_ids(result, field_name)
+    result.dig('data', 'findUser', field_name, 'nodes').flat_map do |node|
+      node.dig('labels', 'nodes').map { |n| n['id'] }
+    end
+  end
+
+  def all_issue_titles(result, field_name)
+    result.dig('data', 'findUser', field_name, 'nodes').map do |node|
+      node['title']
+    end
+  end
+end

spec/requests/api/graphql/project_query_spec.rb

@@ -62,6 +62,54 @@ describe 'getting project information' do
     end
   end
 
+  describe 'performance' do
+    before do
+      project.add_developer(current_user)
+      mrs = create_list(:merge_request, 10, :closed, :with_head_pipeline,
+                        source_project: project,
+                        author: current_user)
+      mrs.each do |mr|
+        mr.assignees << create(:user)
+        mr.assignees << current_user
+      end
+    end
+
+    def run_query(number)
+      q = <<~GQL
+      query {
+        project(fullPath: "#{project.full_path}") {
+          mergeRequests(first: #{number}) {
+            nodes {
+              assignees { nodes { username } }
+              headPipeline { status }
+            }
+          }
+        }
+      }
+      GQL
+
+      post_graphql(q, current_user: current_user)
+    end
+
+    it 'returns appropriate results' do
+      run_query(2)
+
+      mrs = graphql_data.dig('project', 'mergeRequests', 'nodes')
+
+      expect(mrs.size).to eq(2)
+      expect(mrs).to all(
+        match(
+          a_hash_including(
+            'assignees' => { 'nodes' => all(match(a_hash_including('username' => be_present))) },
+            'headPipeline' => { 'status' => be_present }
+          )))
+    end
+
+    it 'can lookahead to eliminate N+1 queries', :use_clean_rails_memory_store_caching, :request_store do
+      expect { run_query(10) }.to issue_same_number_of_queries_as { run_query(1) }.or_fewer.ignoring_cached_queries
+    end
+  end
+
   context 'when the user does not have access to the project' do
     it 'returns an empty field' do
       post_graphql(query, current_user: current_user)

spec/support/matchers/exceed_query_limit.rb

@@ -59,11 +59,15 @@ module ExceedQueryLimitHelpers
 
   def verify_count(&block)
     @subject_block = block
-    actual_count > expected_count + threshold
+    actual_count > maximum
+  end
+
+  def maximum
+    expected_count + threshold
   end
 
   def failure_message
-    threshold_message = threshold > 0 ? " (+#{@threshold})" : ''
+    threshold_message = threshold > 0 ? " (+#{threshold})" : ''
     counts = "#{expected_count}#{threshold_message}"
     "Expected a maximum of #{counts} queries, got #{actual_count}:\n\n#{log_message}"
   end
@@ -73,6 +77,55 @@ module ExceedQueryLimitHelpers
   end
 end
 
+RSpec::Matchers.define :issue_fewer_queries_than do
+  supports_block_expectations
+
+  include ExceedQueryLimitHelpers
+
+  def control
+    block_arg
+  end
+
+  def control_recorder
+    @control_recorder ||= ActiveRecord::QueryRecorder.new(&control)
+  end
+
+  def expected_count
+    control_recorder.count
+  end
+
+  def verify_count(&block)
+    @subject_block = block
+
+    # These blocks need to be evaluated in an expected order, in case
+    # the events in expected affect the counts in actual
+    expected_count
+    actual_count
+
+    actual_count < expected_count
+  end
+
+  match do |block|
+    verify_count(&block)
+  end
+
+  def failure_message
+    <<~MSG
+    Expected to issue fewer than #{expected_count} queries, but got #{actual_count}
+
+    #{log_message}
+    MSG
+  end
+
+  failure_message_when_negated do |actual|
+    <<~MSG
+    Expected query count of #{actual_count} to be less than #{expected_count}
+
+    #{log_message}
+    MSG
+  end
+end
+
 RSpec::Matchers.define :issue_same_number_of_queries_as do
   supports_block_expectations
 
@@ -82,30 +135,66 @@ RSpec::Matchers.define :issue_same_number_of_queries_as do
     block_arg
   end
 
+  chain :or_fewer do
+    @or_fewer = true
+  end
+
+  chain :ignoring_cached_queries do
+    @skip_cached = true
+  end
+
   def control_recorder
     @control_recorder ||= ActiveRecord::QueryRecorder.new(&control)
   end
 
   def expected_count
-    @expected_count ||= control_recorder.count
+    control_recorder.count
   end
 
   def verify_count(&block)
     @subject_block = block
 
-    (expected_count - actual_count).abs <= threshold
+    # These blocks need to be evaluated in an expected order, in case
+    # the events in expected affect the counts in actual
+    expected_count
+    actual_count
+
+    if @or_fewer
+      actual_count <= expected_count
+    else
+      (expected_count - actual_count).abs <= threshold
+    end
   end
 
   match do |block|
     verify_count(&block)
   end
 
+  def failure_message
+    <<~MSG
+    Expected #{expected_count_message} queries, but got #{actual_count}
+
+    #{log_message}
+    MSG
+  end
+
   failure_message_when_negated do |actual|
-    failure_message
+    <<~MSG
+    Expected #{actual_count} not to equal #{expected_count_message}
+
+    #{log_message}
+    MSG
+  end
+
+  def expected_count_message
+    or_fewer_msg = "or fewer" if @or_fewer
+    threshold_msg = "(+/- #{threshold})" unless threshold.zero?
+
+    ["#{expected_count}", or_fewer_msg, threshold_msg].compact.join(' ')
   end
 
   def skip_cached
-    false
+    @skip_cached || false
   end
 end