Migrate DiffCollection limiting logic to Gitaly

GITALY_SERVER_VERSION

-0.21.2
+0.22.0

Gemfile

@@ -383,7 +383,7 @@ gem 'vmstat', '~> 2.3.0'
 gem 'sys-filesystem', '~> 1.1.6'
 
 # Gitaly GRPC client
-gem 'gitaly', '~> 0.14.0'
+gem 'gitaly', '~> 0.17.0'
 
 gem 'toml-rb', '~> 0.3.15', require: false
 

Gemfile.lock

@@ -269,7 +269,7 @@ GEM
       po_to_json (>= 1.0.0)
       rails (>= 3.2.0)
     gherkin-ruby (0.3.2)
-    gitaly (0.14.0)
+    gitaly (0.17.0)
       google-protobuf (~> 3.1)
       grpc (~> 1.0)
     github-linguist (4.7.6)
@@ -971,7 +971,7 @@ DEPENDENCIES
   gettext (~> 3.2.2)
   gettext_i18n_rails (~> 1.8.0)
   gettext_i18n_rails_js (~> 1.2.0)
-  gitaly (~> 0.14.0)
+  gitaly (~> 0.17.0)
   github-linguist (~> 4.7.0)
   gitlab-flowdock-git-hook (~> 1.0.1)
   gitlab-markup (~> 1.5.1)

lib/gitlab/git/diff.rb

@@ -234,6 +234,8 @@ module Gitlab
         @new_file = diff.from_id == BLANK_SHA
         @renamed_file = diff.from_path != diff.to_path
         @deleted_file = diff.to_id == BLANK_SHA
+
+        collapse! if diff.respond_to?(:collapsed) && diff.collapsed
       end
 
       def prune_diff_if_eligible

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

lib/gitlab/gitaly_client/commit_service.rb

@@ -23,9 +23,13 @@ module Gitlab
       def diff_from_parent(commit, options = {})
         request_params = commit_diff_request_params(commit, options)
         request_params[:ignore_whitespace_change] = options.fetch(:ignore_whitespace_change, false)
+        request_params[:enforce_limits] = options.fetch(:limits, true)
+        request_params[:collapse_diffs] = request_params[:enforce_limits] || !options.fetch(:expanded, true)
+        request_params.merge!(Gitlab::Git::DiffCollection.collection_limits(options).to_h)
+
         request = Gitaly::CommitDiffRequest.new(request_params)
         response = GitalyClient.call(@repository.storage, :diff_service, :commit_diff, request)
-        Gitlab::Git::DiffCollection.new(GitalyClient::DiffStitcher.new(response), options)
+        Gitlab::Git::DiffCollection.new(GitalyClient::DiffStitcher.new(response), options.merge(from_gitaly: true))
       end
 
       def commit_deltas(commit)

lib/gitlab/gitaly_client/diff.rb

 module Gitlab
   module GitalyClient
     class Diff
-      FIELDS = %i(from_path to_path old_mode new_mode from_id to_id patch).freeze
+      FIELDS = %i(from_path to_path old_mode new_mode from_id to_id patch overflow_marker collapsed).freeze
 
       attr_accessor(*FIELDS)
 

spec/lib/gitlab/git/diff_collection_spec.rb

@@ -484,6 +484,8 @@ describe Gitlab::Git::DiffCollection, seed_helper: true do
     end
 
     def each
+      return enum_for(:each) unless block_given?
+
       loop do
         break if @count.zero?
         # It is critical to decrement before yielding. We may never reach the lines after 'yield'.

spec/lib/gitlab/gitaly_client/commit_service_spec.rb

@@ -12,7 +12,10 @@ describe Gitlab::GitalyClient::CommitService do
         request = Gitaly::CommitDiffRequest.new(
           repository: repository_message,
           left_commit_id: 'cfe32cf61b73a0d5e9f13e774abde7ff789b1660',
-          right_commit_id: commit.id
+          right_commit_id: commit.id,
+          collapse_diffs: true,
+          enforce_limits: true,
+          **Gitlab::Git::DiffCollection.collection_limits.to_h
         )
 
         expect_any_instance_of(Gitaly::DiffService::Stub).to receive(:commit_diff).with(request, kind_of(Hash))
@@ -27,7 +30,10 @@ describe Gitlab::GitalyClient::CommitService do
         request        = Gitaly::CommitDiffRequest.new(
           repository: repository_message,
           left_commit_id: '4b825dc642cb6eb9a060e54bf8d69288fbee4904',
-          right_commit_id: initial_commit.id
+          right_commit_id: initial_commit.id,
+          collapse_diffs: true,
+          enforce_limits: true,
+          **Gitlab::Git::DiffCollection.collection_limits.to_h
         )
 
         expect_any_instance_of(Gitaly::DiffService::Stub).to receive(:commit_diff).with(request, kind_of(Hash))
@@ -43,7 +49,7 @@ describe Gitlab::GitalyClient::CommitService do
     end
 
     it 'passes options to Gitlab::Git::DiffCollection' do
-      options = { max_files: 31, max_lines: 13 }
+      options = { max_files: 31, max_lines: 13, from_gitaly: true }
 
       expect(Gitlab::Git::DiffCollection).to receive(:new).with(kind_of(Enumerable), options)