Allow ReferenceExtractor to efficiently load references from multiple texts at once

4 files changed, 36 insertions, 19 deletions

diff --git a/lib/gitlab/markdown/reference_filter.rb b/lib/gitlab/markdown/reference_filter.rb
index 0ae0d93f70d..ede9f8865ff 100644
--- a/lib/gitlab/markdown/reference_filter.rb
+++ b/lib/gitlab/markdown/reference_filter.rb
@@ -12,7 +12,14 @@ module Gitlab
     #   :project (required) - Current project, ignored if reference is cross-project.
     #   :only_path          - Generate path-only links.
     class ReferenceFilter < HTML::Pipeline::Filter
-      LazyReference = Struct.new(:klass, :ids)
+      LazyReference = Struct.new(:klass, :ids) do
+        def self.load(refs)
+          refs.group_by(&:klass).flat_map do |klass, refs|
+            ids = refs.flat_map(&:ids)
+            klass.where(id: ids)
+          end
+        end
+      end
 
       def self.user_can_reference?(user, node, context)
         if node.has_attribute?('data-project')
diff --git a/lib/gitlab/markdown/reference_gatherer_filter.rb b/lib/gitlab/markdown/reference_gatherer_filter.rb
index cf9a2303db8..18df5db94d6 100644
--- a/lib/gitlab/markdown/reference_gatherer_filter.rb
+++ b/lib/gitlab/markdown/reference_gatherer_filter.rb
@@ -21,7 +21,7 @@ module Gitlab
           gather_references(node)
         end
 
-        load_lazy_references
+        load_lazy_references unless context[:load_lazy_references] == false
 
         doc
       end
@@ -56,11 +56,8 @@ module Gitlab
       # Will load all references of one type using one query.
       def load_lazy_references
         result[:lazy_references].each do |type, refs|
-          refs.group_by(&:klass).each do |klass, refs|
-            ids = refs.map(&:ids).flatten
-            values = klass.find(ids)
-            result[:references][type].push(*values)
-          end
+          values = ReferenceFilter::LazyReference.load(refs)
+          result[:references][type].concat(values)
         end
       end
 
diff --git a/lib/gitlab/reference_extractor.rb b/lib/gitlab/reference_extractor.rb
index 2e546ef0d54..8100f2675a7 100644
--- a/lib/gitlab/reference_extractor.rb
+++ b/lib/gitlab/reference_extractor.rb
@@ -10,9 +10,10 @@ module Gitlab
       @current_user = current_user
     end
 
-    def analyze(text)
+    def analyze(texts)
       references.clear
-      @text = Gitlab::Markdown.render_without_gfm(text)
+      texts = Array(texts)
+      @texts = texts.map { |text| Gitlab::Markdown.render_without_gfm(text) }
     end
 
     %i(user label issue merge_request snippet commit commit_range).each do |type|
@@ -47,13 +48,25 @@ module Gitlab
         current_user: current_user,
         # We don't actually care about the links generated
         only_path: true,
-        ignore_blockquotes: true
+        ignore_blockquotes: true,
+        load_lazy_references: false
       }
 
       pipeline = HTML::Pipeline.new([filter, Gitlab::Markdown::ReferenceGathererFilter], context)
-      result = pipeline.call(@text)
 
-      result[:references][filter_type]
+      values = []
+      lazy_references = []
+
+      @texts.each do |text|
+        result = pipeline.call(text)
+
+        values.concat(result[:references][filter_type])
+        lazy_references.concat(result[:lazy_references][filter_type])
+      end
+
+      lazy_values = Gitlab::Markdown::ReferenceFilter::LazyReference.load(lazy_references)
+      values.concat(lazy_values)
+      values
     end
   end
 end
diff --git a/spec/lib/gitlab/reference_extractor_spec.rb b/spec/lib/gitlab/reference_extractor_spec.rb
index 088e34f050c..ad84d2274e8 100644
--- a/spec/lib/gitlab/reference_extractor_spec.rb
+++ b/spec/lib/gitlab/reference_extractor_spec.rb
@@ -13,7 +13,7 @@ describe Gitlab::ReferenceExtractor do
     project.team << [@u_bar, :guest]
 
     subject.analyze('@foo, @baduser, @bar, and @offteam')
-    expect(subject.users).to eq([@u_foo, @u_bar, @u_offteam])
+    expect(subject.users).to match_array([@u_foo, @u_bar, @u_offteam])
   end
 
   it 'ignores user mentions inside specific elements' do
@@ -37,7 +37,7 @@ describe Gitlab::ReferenceExtractor do
 
       > @offteam
     })
-    expect(subject.users).to eq([])
+    expect(subject.users).to match_array([])
   end
 
   it 'accesses valid issue objects' do
@@ -45,7 +45,7 @@ describe Gitlab::ReferenceExtractor do
     @i1 = create(:issue, project: project)
 
     subject.analyze("#{@i0.to_reference}, #{@i1.to_reference}, and #{Issue.reference_prefix}999.")
-    expect(subject.issues).to eq([@i0, @i1])
+    expect(subject.issues).to match_array([@i0, @i1])
   end
 
   it 'accesses valid merge requests' do
@@ -53,7 +53,7 @@ describe Gitlab::ReferenceExtractor do
     @m1 = create(:merge_request, source_project: project, target_project: project, source_branch: 'feature_conflict')
 
     subject.analyze("!999, !#{@m1.iid}, and !#{@m0.iid}.")
-    expect(subject.merge_requests).to eq([@m1, @m0])
+    expect(subject.merge_requests).to match_array([@m1, @m0])
   end
 
   it 'accesses valid labels' do
@@ -62,7 +62,7 @@ describe Gitlab::ReferenceExtractor do
     @l2 = create(:label)
 
     subject.analyze("~#{@l0.id}, ~999, ~#{@l2.id}, ~#{@l1.id}")
-    expect(subject.labels).to eq([@l0, @l1])
+    expect(subject.labels).to match_array([@l0, @l1])
   end
 
   it 'accesses valid snippets' do
@@ -71,7 +71,7 @@ describe Gitlab::ReferenceExtractor do
     @s2 = create(:project_snippet)
 
     subject.analyze("$#{@s0.id}, $999, $#{@s2.id}, $#{@s1.id}")
-    expect(subject.snippets).to eq([@s0, @s1])
+    expect(subject.snippets).to match_array([@s0, @s1])
   end
 
   it 'accesses valid commits' do
@@ -109,7 +109,7 @@ describe Gitlab::ReferenceExtractor do
       subject.analyze("this refers issue #{issue.to_reference(project)}")
       extracted = subject.issues
       expect(extracted.size).to eq(1)
-      expect(extracted).to eq([issue])
+      expect(extracted).to match_array([issue])
     end
   end
 end