diff_collection.rb 3.87 KB
Newer Older
1 2
# Gitaly note: JV: no RPC's here.

Robert Speicher's avatar
Robert Speicher committed
3 4 5 6 7 8 9
module Gitlab
  module Git
    class DiffCollection
      include Enumerable

      DEFAULT_LIMITS = { max_files: 100, max_lines: 5000 }.freeze

10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25
      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

Robert Speicher's avatar
Robert Speicher committed
26 27
      def initialize(iterator, options = {})
        @iterator = iterator
28
        @limits = self.class.collection_limits(options)
Douwe Maan's avatar
Douwe Maan committed
29
        @enforce_limits = !!options.fetch(:limits, true)
30
        @expanded = !!options.fetch(:expanded, true)
31
        @from_gitaly = options.fetch(:from_gitaly, false)
Robert Speicher's avatar
Robert Speicher committed
32 33 34 35

        @line_count = 0
        @byte_count = 0
        @overflow = false
36
        @empty = true
Robert Speicher's avatar
Robert Speicher committed
37 38 39 40
        @array = Array.new
      end

      def each(&block)
41 42 43 44 45 46 47 48 49 50 51
        @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
Robert Speicher's avatar
Robert Speicher committed
52
        end
53 54 55 56 57

        @populated = true

        # Allow iterator to be garbage-collected. It cannot be reused anyway.
        @iterator = nil
Robert Speicher's avatar
Robert Speicher committed
58 59 60
      end

      def empty?
61 62
        any? # Make sure the iterator has been exercised
        @empty
Robert Speicher's avatar
Robert Speicher committed
63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90
      end

      def overflow?
        populate!
        !!@overflow
      end

      def size
        @size ||= count # forces a loop using each method
      end

      def real_size
        populate!

        if @overflow
          "#{size}+"
        else
          size.to_s
        end
      end

      def decorate!
        collection = each_with_index do |element, i|
          @array[i] = yield(element)
        end
        collection
      end

Douwe Maan's avatar
Douwe Maan committed
91 92
      alias_method :to_ary, :to_a

Robert Speicher's avatar
Robert Speicher committed
93 94 95 96 97 98 99 100 101 102
      private

      def populate!
        return if @populated

        each { nil } # force a loop through all diffs
        nil
      end

      def over_safe_limits?(files)
103
        files >= safe_max_files || @line_count > safe_max_lines || @byte_count >= safe_max_bytes
Robert Speicher's avatar
Robert Speicher committed
104 105
      end

106 107 108 109 110 111 112 113 114 115 116 117
      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
118 119
          i += 1
        end
120
      end
121

122 123
      def each_rugged_patch
        i = @array.length
Robert Speicher's avatar
Robert Speicher committed
124

125 126
        @iterator.each do |raw|
          @empty = false
Robert Speicher's avatar
Robert Speicher committed
127

128
          if @enforce_limits && i >= max_files
Robert Speicher's avatar
Robert Speicher committed
129 130 131 132
            @overflow = true
            break
          end

Douwe Maan's avatar
Douwe Maan committed
133
          expanded = !@enforce_limits || @expanded
Robert Speicher's avatar
Robert Speicher committed
134

135
          diff = Gitlab::Git::Diff.new(raw, expanded: expanded)
Robert Speicher's avatar
Robert Speicher committed
136

Douwe Maan's avatar
Douwe Maan committed
137
          if !expanded && over_safe_limits?(i) && diff.line_count > 0
138
            diff.collapse!
Robert Speicher's avatar
Robert Speicher committed
139 140 141 142 143
          end

          @line_count += diff.line_count
          @byte_count += diff.diff.bytesize

144
          if @enforce_limits && (@line_count >= max_lines || @byte_count >= max_bytes)
Robert Speicher's avatar
Robert Speicher committed
145 146 147 148 149 150 151
            # This last Diff instance pushes us over the lines limit. We stop and
            # discard it.
            @overflow = true
            break
          end

          yield @array[i] = diff
152
          i += 1
Robert Speicher's avatar
Robert Speicher committed
153 154 155 156 157
        end
      end
    end
  end
end