Improve search api preloading for merge_requests scope

Create MergeRequestPreloader for Search API
to be used with merge_requests scope.

The preloader lazily loads attributes using
BatchLoader and utilizes the lazy attributes in
the MergeRequestBasic Entity if lazy_load option
is set. This improves response times by avoiding
N+1 queries.

app/services/search_service.rb

@@ -5,7 +5,6 @@ class SearchService
 
   SEARCH_TERM_LIMIT = 64
   SEARCH_CHAR_LIMIT = 4096
-
   DEFAULT_PER_PAGE = Gitlab::SearchResults::DEFAULT_PER_PAGE
   MAX_PER_PAGE = 200
 
@@ -62,8 +61,8 @@ class SearchService
     @search_results ||= search_service.execute
   end
 
-  def search_objects
-    @search_objects ||= redact_unauthorized_results(search_results.objects(scope, page: params[:page], per_page: per_page))
+  def search_objects(preload_method = nil)
+    @search_objects ||= redact_unauthorized_results(search_results.objects(scope, page: params[:page], per_page: per_page, preload_method: preload_method))
   end
 
   private

ee/changelogs/unreleased/215709-improve-performance-of-search-api-advanced-merge_requests-scope.yml

+---
+title: Fix N+1 queries for Elastic Search merge_requests scope.
+merge_request: 30546
+author:
+type: performance

ee/lib/gitlab/elastic/search_results.rb

@@ -24,20 +24,20 @@ module Gitlab
         @public_and_internal_projects = public_and_internal_projects
       end
 
-      def objects(scope, page: 1, per_page: DEFAULT_PER_PAGE)
+      def objects(scope, page: 1, per_page: DEFAULT_PER_PAGE, preload_method: nil)
         page = (page || 1).to_i
 
         case scope
         when 'projects'
-          eager_load(projects, page, per_page, eager: [:route, :namespace, :compliance_framework_setting])
+          eager_load(projects, page, per_page, preload_method, [:route, :namespace, :compliance_framework_setting])
         when 'issues'
-          eager_load(issues, page, per_page, eager: { project: [:route, :namespace], labels: [], timelogs: [], assignees: [] })
+          eager_load(issues, page, per_page, preload_method, project: [:route, :namespace], labels: [], timelogs: [], assignees: [])
         when 'merge_requests'
-          eager_load(merge_requests, page, per_page, eager: { target_project: [:route, :namespace] })
+          eager_load(merge_requests, page, per_page, preload_method, target_project: [:route, :namespace])
         when 'milestones'
-          eager_load(milestones, page, per_page, eager: { project: [:route, :namespace] })
+          eager_load(milestones, page, per_page, preload_method, project: [:route, :namespace])
         when 'notes'
-          eager_load(notes, page, per_page, eager: { project: [:route, :namespace] })
+          eager_load(notes, page, per_page, preload_method, project: [:route, :namespace])
         when 'blobs'
           blobs(page: page, per_page: per_page)
         when 'wiki_blobs'
@@ -166,10 +166,12 @@ module Gitlab
       private
 
       # Apply some eager loading to the `records` of an ES result object without
-      # losing pagination information
-      def eager_load(es_result, page, per_page, eager:)
+      # losing pagination information. Also, take advantage of preload method if
+      # provided by the caller.
+      def eager_load(es_result, page, per_page, preload_method, eager)
         paginated_base = es_result.page(page).per(per_page)
         relation = paginated_base.records.includes(eager) # rubocop:disable CodeReuse/ActiveRecord
+        relation = relation.public_send(preload_method) if preload_method # rubocop:disable GitlabSecurity/PublicSend
 
         Kaminari.paginate_array(
           relation,

ee/lib/gitlab/elastic/snippet_search_results.rb

@@ -3,10 +3,9 @@
 module Gitlab
   module Elastic
     class SnippetSearchResults < Gitlab::Elastic::SearchResults
-      def objects(scope, page: 1, per_page: DEFAULT_PER_PAGE)
+      def objects(scope, page: 1, per_page: DEFAULT_PER_PAGE, preload_method: nil)
         page = (page || 1).to_i
-
-        eager_load(snippet_titles, page, per_page, eager: { project: [:route, :namespace] })
+        eager_load(snippet_titles, page, per_page, preload_method, project: [:route, :namespace])
       end
 
       def formatted_count(scope)
@@ -26,10 +25,6 @@ module Gitlab
       def limited_snippet_titles_count
         @limited_snippet_titles_count ||= snippet_titles.total_count
       end
-
-      def paginated_objects(relation, page)
-        super.records
-      end
     end
   end
 end

ee/spec/requests/api/search_spec.rb

@@ -59,6 +59,30 @@ describe API::Search do
   end
 
   shared_examples 'elasticsearch enabled' do |level:|
+    context 'for merge_requests scope', :sidekiq_inline do
+      before do
+        create(:labeled_merge_request, target_branch: 'feature_1', source_project: project, labels: [create(:label), create(:label)])
+        create(:merge_request, target_branch: 'feature_2', source_project: project, author: create(:user))
+        create(:merge_request, target_branch: 'feature_3', source_project: project, milestone: create(:milestone, project: project))
+        create(:merge_request, target_branch: 'feature_4', source_project: project)
+        ensure_elasticsearch_index!
+      end
+
+      it 'avoids N+1 queries' do
+        control = ActiveRecord::QueryRecorder.new { get api(endpoint, user), params: { scope: 'merge_requests', search: '*' } }
+
+        create(:labeled_merge_request, target_branch: 'feature_5', source_project: project, labels: [create(:label), create(:label)])
+        create(:merge_request, target_branch: 'feature_6', source_project: project, author: create(:user))
+        create(:merge_request, target_branch: 'feature_7', source_project: project, milestone: create(:milestone, project: project))
+        create(:merge_request, target_branch: 'feature_8', source_project: project)
+
+        ensure_elasticsearch_index!
+
+        # Some N+1 queries still exist
+        expect { get api(endpoint, user), params: { scope: 'merge_requests', search: '*' } }.not_to exceed_query_limit(control.count + 16)
+      end
+    end
+
     context 'for wiki_blobs scope', :sidekiq_might_not_need_inline do
       before do
         wiki = create(:project_wiki, project: project)

ee/spec/requests/search_controller_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe SearchController, type: :request do
+  let_it_be(:user) { create(:user) }
+  let_it_be(:group) { create(:group) }
+  let(:project) { create(:project, :public, :repository, :wiki_repo, name: 'awesome project', group: group) }
+
+  before_all do
+    login_as(user)
+  end
+
+  def send_search_request(params)
+    get search_path, params: params
+  end
+
+  describe 'GET /search' do
+    context 'when elasticsearch is enabled', :elastic, :sidekiq_inline do
+      before do
+        stub_ee_application_setting(elasticsearch_search: true, elasticsearch_indexing: true)
+      end
+      context 'for merge_request scope' do
+        before do
+          create(:merge_request, target_branch: 'feature_1', source_project: project)
+          create(:merge_request, target_branch: 'feature_2', source_project: project)
+          create(:merge_request, target_branch: 'feature_3', source_project: project)
+          create(:merge_request, target_branch: 'feature_4', source_project: project)
+          ensure_elasticsearch_index!
+        end
+        it 'avoids N+1 queries' do
+          control = ActiveRecord::QueryRecorder.new(skip_cached: false) { send_search_request(scope: 'merge_requests', search: '*') }
+
+          create(:merge_request, target_branch: 'feature_5', source_project: project)
+          create(:merge_request, target_branch: 'feature_6', source_project: project)
+          create(:merge_request, target_branch: 'feature_7', source_project: project)
+          create(:merge_request, target_branch: 'feature_8', source_project: project)
+
+          ensure_elasticsearch_index!
+
+          # some N+1 queries still exist
+          expect { send_search_request(scope: 'merge_requests', search: '*') }
+              .not_to exceed_all_query_limit(control.count + 2)
+        end
+      end
+    end
+  end
+end

lib/api/entities/merge_request_basic.rb

@@ -6,19 +6,15 @@ module API
       expose :merged_by, using: Entities::UserBasic do |merge_request, _options|
         merge_request.metrics&.merged_by
       end
-
       expose :merged_at do |merge_request, _options|
         merge_request.metrics&.merged_at
       end
-
       expose :closed_by, using: Entities::UserBasic do |merge_request, _options|
         merge_request.metrics&.latest_closed_by
       end
-
       expose :closed_at do |merge_request, _options|
         merge_request.metrics&.latest_closed_at
       end
-
       expose :title_html, if: -> (_, options) { options[:render_html] } do |entity|
         MarkupHelper.markdown_field(entity, :title)
       end
@@ -33,7 +29,6 @@ module API
         merge_request.assignee
       end
       expose :author, :assignees, using: Entities::UserBasic
-
       expose :source_project_id, :target_project_id
       expose :labels do |merge_request, options|
         if options[:with_labels_details]
@@ -85,11 +80,8 @@ module API
       end
 
       expose :squash
-
       expose :task_completion_status
-
       expose :cannot_be_merged?, as: :has_conflicts
-
       expose :mergeable_discussions_state?, as: :blocking_discussions_resolved
     end
   end

lib/api/search.rb

@@ -20,6 +20,10 @@ module API
         users: Entities::UserBasic
       }.freeze
 
+      SCOPE_PRELOAD_METHOD = {
+          merge_requests: :with_api_entity_associations
+      }.freeze
+
       def search(additional_params = {})
         search_params = {
           scope: params[:scope],
@@ -29,7 +33,7 @@ module API
           per_page: params[:per_page]
         }.merge(additional_params)
 
-        results = SearchService.new(current_user, search_params).search_objects
+        results = SearchService.new(current_user, search_params).search_objects(preload_method)
 
         paginate(results)
       end
@@ -42,6 +46,10 @@ module API
         SCOPE_ENTITY[params[:scope].to_sym]
       end
 
+      def preload_method
+        SCOPE_PRELOAD_METHOD[params[:scope].to_sym]
+      end
+
       def verify_search_scope!(resource:)
         # In EE we have additional validation requirements for searches.
         # Defining this method here as a noop allows us to easily extend it in

lib/gitlab/project_search_results.rb

@@ -11,7 +11,7 @@ module Gitlab
       @query = query
     end
 
-    def objects(scope, page: nil, per_page: DEFAULT_PER_PAGE)
+    def objects(scope, page: nil, per_page: DEFAULT_PER_PAGE, preload_method: nil)
       case scope
       when 'notes'
         notes.page(page).per(per_page)

lib/gitlab/search_results.rb

@@ -26,7 +26,7 @@ module Gitlab
       @default_project_filter = default_project_filter
     end
 
-    def objects(scope, page: nil, per_page: DEFAULT_PER_PAGE, without_count: true)
+    def objects(scope, page: nil, per_page: DEFAULT_PER_PAGE, without_count: true, preload_method: nil)
       collection = case scope
                    when 'projects'
                      projects

lib/gitlab/snippet_search_results.rb

@@ -11,7 +11,7 @@ module Gitlab
       @query = query
     end
 
-    def objects(scope, page: nil, per_page: DEFAULT_PER_PAGE)
+    def objects(scope, page: nil, per_page: DEFAULT_PER_PAGE, preload_method: nil)
       paginated_objects(snippet_titles, page, per_page)
     end