Optimize limited_notes_countjprovazn-scoped-limit

For getting notes, complex query with multiple UNIONs (for each of
noteable type) is constructed, then using LIMIT on such query is still
slow because DB has to compute all UNIONs first.

For getting only limited count, we can break the UNION query into
multiple smaller queries on which we can use LIMIT, also we don't have
to search all noteable types if we reach the limit earlier.

4 files changed, 69 insertions, 2 deletions

diff --git a/app/finders/notes_finder.rb b/app/finders/notes_finder.rb
index 33ee1e975b9..35f4ff2f62f 100644
--- a/app/finders/notes_finder.rb
+++ b/app/finders/notes_finder.rb
@@ -48,11 +48,23 @@ class NotesFinder
   def init_collection
     if target
       notes_on_target
+    elsif target_type
+      notes_of_target_type
     else
       notes_of_any_type
     end
   end
 
+  def notes_of_target_type
+    notes = notes_for_type(target_type)
+
+    search(notes)
+  end
+
+  def target_type
+    @params[:target_type]
+  end
+
   def notes_of_any_type
     types = %w(commit issue merge_request snippet)
     note_relations = types.map { |t| notes_for_type(t) }
diff --git a/lib/gitlab/project_search_results.rb b/lib/gitlab/project_search_results.rb
index 23801afb71c..29277ec6481 100644
--- a/lib/gitlab/project_search_results.rb
+++ b/lib/gitlab/project_search_results.rb
@@ -30,7 +30,17 @@ module Gitlab
     end
 
     def limited_notes_count
-      @limited_notes_count ||= notes.limit(count_limit).count
+      return @limited_notes_count if defined?(@limited_notes_count)
+
+      types = %w(issue merge_request commit snippet)
+      @limited_notes_count = 0
+
+      types.each do |type|
+        @limited_notes_count += notes_finder(type).limit(count_limit).count
+        break if @limited_notes_count >= count_limit
+      end
+
+      @limited_notes_count
     end
 
     def wiki_blobs_count
@@ -107,7 +117,11 @@ module Gitlab
     end
 
     def notes
-      @notes ||= NotesFinder.new(project, @current_user, search: query).execute.user.order('updated_at DESC')
+      @notes ||= notes_finder(nil)
+    end
+
+    def notes_finder(type)
+      NotesFinder.new(project, @current_user, search: query, target_type: type).execute.user.order('updated_at DESC')
     end
 
     def commits
diff --git a/spec/finders/notes_finder_spec.rb b/spec/finders/notes_finder_spec.rb
index 7b43494eea2..f1ae2c7ab65 100644
--- a/spec/finders/notes_finder_spec.rb
+++ b/spec/finders/notes_finder_spec.rb
@@ -75,6 +75,18 @@ describe NotesFinder do
       end
     end
 
+    context 'for target type' do
+      let(:project) { create(:project, :repository) }
+      let!(:note1) { create :note_on_issue, project: project }
+      let!(:note2) { create :note_on_commit, project: project }
+
+      it 'finds only notes for the selected type' do
+        notes = described_class.new(project, user, target_type: 'issue').execute
+
+        expect(notes).to eq([note1])
+      end
+    end
+
     context 'for target' do
       let(:project) { create(:project, :repository) }
       let(:note1) { create :note_on_commit, project: project }
diff --git a/spec/lib/gitlab/project_search_results_spec.rb b/spec/lib/gitlab/project_search_results_spec.rb
index 28e7ce78ffd..c46bb8edebf 100644
--- a/spec/lib/gitlab/project_search_results_spec.rb
+++ b/spec/lib/gitlab/project_search_results_spec.rb
@@ -304,6 +304,35 @@ describe Gitlab::ProjectSearchResults do
     end
   end
 
+  describe '#limited_notes_count' do
+    let(:project) { create(:project, :public) }
+    let(:note) { create(:note_on_issue, project: project) }
+    let(:results) { described_class.new(user, project, note.note) }
+
+    context 'when count_limit is lower than total amount' do
+      before do
+        allow(results).to receive(:count_limit).and_return(1)
+      end
+
+      it 'calls note finder once to get the limited amount of notes' do
+        expect(results).to receive(:notes_finder).once.and_call_original
+        expect(results.limited_notes_count).to eq(1)
+      end
+    end
+
+    context 'when count_limit is higher than total amount' do
+      it 'calls note finder multiple times to get the limited amount of notes' do
+        project = create(:project, :public)
+        note = create(:note_on_issue, project: project)
+
+        results = described_class.new(user, project, note.note)
+
+        expect(results).to receive(:notes_finder).exactly(4).times.and_call_original
+        expect(results.limited_notes_count).to eq(1)
+      end
+    end
+  end
+
   # Examples for commit access level test
   #
   # params: