Use the iterator inside DiffCollection only once

1 files changed, 21 insertions, 18 deletions

diff --git a/lib/gitlab/git/diff_collection.rb b/lib/gitlab/git/diff_collection.rb
index bcbad8ec829..398d6013d35 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
 
@@ -70,7 +66,6 @@ module Gitlab
         return if @populated
 
         each { nil } # force a loop through all diffs
-        @populated = true
         nil
       end
 
@@ -79,15 +74,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
@@ -113,7 +110,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