1 files changed, 52 insertions, 22 deletions

diff --git a/lib/gitlab/git/diff_collection.rb b/lib/gitlab/git/diff_collection.rb
index 0d8fe185ac5..87ed9c3ea26 100644
--- a/lib/gitlab/git/diff_collection.rb
+++ b/lib/gitlab/git/diff_collection.rb
@@ -7,16 +7,28 @@ module Gitlab
 
       DEFAULT_LIMITS = { max_files: 100, max_lines: 5000 }.freeze
 
+      attr_reader :limits
+
+      delegate :max_files, :max_lines, :max_bytes, :safe_max_files, :safe_max_lines, :safe_max_bytes, to: :limits
+
+      def self.collection_limits(options = {})
+        limits = {}
+        limits[:max_files] = options.fetch(:max_files, DEFAULT_LIMITS[:max_files])
+        limits[:max_lines] = options.fetch(:max_lines, DEFAULT_LIMITS[:max_lines])
+        limits[:max_bytes] = limits[:max_files] * 5.kilobytes # Average 5 KB per file
+        limits[:safe_max_files] = [limits[:max_files], DEFAULT_LIMITS[:max_files]].min
+        limits[:safe_max_lines] = [limits[:max_lines], DEFAULT_LIMITS[:max_lines]].min
+        limits[:safe_max_bytes] = limits[:safe_max_files] * 5.kilobytes # Average 5 KB per file
+
+        OpenStruct.new(limits)
+      end
+
       def initialize(iterator, options = {})
         @iterator = iterator
-        @max_files = options.fetch(:max_files, DEFAULT_LIMITS[:max_files])
-        @max_lines = options.fetch(:max_lines, DEFAULT_LIMITS[:max_lines])
-        @max_bytes = @max_files * 5.kilobytes # Average 5 KB per file
-        @safe_max_files = [@max_files, DEFAULT_LIMITS[:max_files]].min
-        @safe_max_lines = [@max_lines, DEFAULT_LIMITS[:max_lines]].min
-        @safe_max_bytes = @safe_max_files * 5.kilobytes # Average 5 KB per file
+        @limits = self.class.collection_limits(options)
         @enforce_limits = !!options.fetch(:limits, true)
         @expanded = !!options.fetch(:expanded, true)
+        @from_gitaly = options.fetch(:from_gitaly, false)
 
         @line_count = 0
         @byte_count = 0
@@ -26,9 +38,23 @@ module Gitlab
       end
 
       def each(&block)
-        Gitlab::GitalyClient.migrate(:commit_raw_diffs) do
-          each_patch(&block)
+        @array.each(&block)
+
+        return if @overflow
+        return if @iterator.nil?
+
+        Gitlab::GitalyClient.migrate(:commit_raw_diffs) do |is_enabled|
+          if is_enabled && @from_gitaly
+            each_gitaly_patch(&block)
+          else
+            each_rugged_patch(&block)
+          end
         end
+
+        @populated = true
+
+        # Allow iterator to be garbage-collected. It cannot be reused anyway.
+        @iterator = nil
       end
 
       def empty?
@@ -74,23 +100,32 @@ module Gitlab
       end
 
       def over_safe_limits?(files)
-        files >= @safe_max_files || @line_count > @safe_max_lines || @byte_count >= @safe_max_bytes
+        files >= safe_max_files || @line_count > safe_max_lines || @byte_count >= safe_max_bytes
       end
 
-      def each_patch
-        i = 0
-        @array.each do |diff|
-          yield diff
+      def each_gitaly_patch
+        i = @array.length
+
+        @iterator.each do |raw|
+          diff = Gitlab::Git::Diff.new(raw, expanded: !@enforce_limits || @expanded)
+
+          if raw.overflow_marker
+            @overflow = true
+            break
+          end
+
+          yield @array[i] = diff
           i += 1
         end
+      end
 
-        return if @overflow
-        return if @iterator.nil?
+      def each_rugged_patch
+        i = @array.length
 
         @iterator.each do |raw|
           @empty = false
 
-          if @enforce_limits && i >= @max_files
+          if @enforce_limits && i >= max_files
             @overflow = true
             break
           end
@@ -106,7 +141,7 @@ module Gitlab
           @line_count += diff.line_count
           @byte_count += diff.diff.bytesize
 
-          if @enforce_limits && (@line_count >= @max_lines || @byte_count >= @max_bytes)
+          if @enforce_limits && (@line_count >= max_lines || @byte_count >= max_bytes)
             # This last Diff instance pushes us over the lines limit. We stop and
             # discard it.
             @overflow = true
@@ -116,11 +151,6 @@ module Gitlab
           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