Merge branch 'ajk-gql-lookahead' into 'master'

[GraphQL] Lookahead optimizations for merge requests

See merge request gitlab-org/gitlab!32373

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
+                @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
 
+    # 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