Merge branch 'diffcollection-no-restarts' into 'master'

Fix buffering in DiffCollection

See merge request !11659

2 files changed, 58 insertions, 21 deletions

diff --git a/lib/gitlab/git/diff_collection.rb b/lib/gitlab/git/diff_collection.rb
index bce3aef8bd3..898a5ae15f2 100644
--- a/lib/gitlab/git/diff_collection.rb
+++ b/lib/gitlab/git/diff_collection.rb
@@ -19,22 +19,19 @@ module Gitlab
         @line_count = 0
         @byte_count = 0
         @overflow = false
+        @empty = true
         @array = Array.new
       end
 
       def each(&block)
-        if @populated
-          # @iterator.each is slower than just iterating the array in place
-          @array.each(&block)
-        else
-          Gitlab::GitalyClient.migrate(:commit_raw_diffs) do
-            each_patch(&block)
-          end
+        Gitlab::GitalyClient.migrate(:commit_raw_diffs) do
+          each_patch(&block)
         end
       end
 
       def empty?
-        !@iterator.any?
+        any? # Make sure the iterator has been exercised
+        @empty
       end
 
       def overflow?
@@ -60,7 +57,6 @@ module Gitlab
         collection = each_with_index do |element, i|
           @array[i] = yield(element)
         end
-        @populated = true
         collection
       end
 
@@ -72,7 +68,6 @@ module Gitlab
         return if @populated
 
         each { nil } # force a loop through all diffs
-        @populated = true
         nil
       end
 
@@ -81,15 +76,17 @@ module Gitlab
       end
 
       def each_patch
-        @iterator.each_with_index do |raw, i|
-          # First yield cached Diff instances from @array
-          if @array[i]
-            yield @array[i]
-            next
-          end
+        i = 0
+        @array.each do |diff|
+          yield diff
+          i += 1
+        end
+
+        return if @overflow
+        return if @iterator.nil?
 
-          # We have exhausted @array, time to create new Diff instances or stop.
-          break if @overflow
+        @iterator.each do |raw|
+          @empty = false
 
           if !@all_diffs && i >= @max_files
             @overflow = true
@@ -115,7 +112,13 @@ module Gitlab
           end
 
           yield @array[i] = diff
+          i += 1
         end
+
+        @populated = true
+
+        # Allow iterator to be garbage-collected. It cannot be reused anyway.
+        @iterator = nil
       end
     end
   end
diff --git a/spec/lib/gitlab/git/diff_collection_spec.rb b/spec/lib/gitlab/git/diff_collection_spec.rb
index 122c93dcd69..ae617b313c5 100644
--- a/spec/lib/gitlab/git/diff_collection_spec.rb
+++ b/spec/lib/gitlab/git/diff_collection_spec.rb
@@ -10,7 +10,7 @@ describe Gitlab::Git::DiffCollection, seed_helper: true do
       no_collapse: no_collapse
     )
   end
-  let(:iterator) { Array.new(file_count, fake_diff(line_length, line_count)) }
+  let(:iterator) { MutatingConstantIterator.new(file_count, fake_diff(line_length, line_count)) }
   let(:file_count) { 0 }
   let(:line_length) { 1 }
   let(:line_count) { 1 }
@@ -64,7 +64,15 @@ describe Gitlab::Git::DiffCollection, seed_helper: true do
           subject { super().real_size }
           it { is_expected.to eq('3') }
         end
-        it { expect(subject.size).to eq(3) }
+
+        describe '#size' do
+          it { expect(subject.size).to eq(3) }
+
+          it 'does not change after peeking' do
+            subject.any?
+            expect(subject.size).to eq(3)
+          end
+        end
 
         context 'when limiting is disabled' do
           let(:all_diffs) { true }
@@ -83,7 +91,15 @@ describe Gitlab::Git::DiffCollection, seed_helper: true do
             subject { super().real_size }
             it { is_expected.to eq('3') }
           end
-          it { expect(subject.size).to eq(3) }
+
+          describe '#size' do
+            it { expect(subject.size).to eq(3) }
+  
+            it 'does not change after peeking' do
+              subject.any?
+              expect(subject.size).to eq(3)
+            end
+          end
         end
       end
 
@@ -457,4 +473,22 @@ describe Gitlab::Git::DiffCollection, seed_helper: true do
   def fake_diff(line_length, line_count)
     { 'diff' => "#{'a' * line_length}\n" * line_count }
   end
+
+  class MutatingConstantIterator
+    include Enumerable
+
+    def initialize(count, value)
+      @count = count
+      @value = value
+    end
+
+    def each
+      loop do
+        break if @count.zero?
+        # It is critical to decrement before yielding. We may never reach the lines after 'yield'.
+        @count -= 1
+        yield @value
+      end
+    end
+  end
 end