BigW Consortium Gitlab

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

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 committed
26 27
      def initialize(iterator, options = {})
        @iterator = iterator
28
        @limits = self.class.collection_limits(options)
29
        @enforce_limits = !!options.fetch(:limits, true)
30
        @expanded = !!options.fetch(:expanded, true)
Robert Speicher committed
31 32 33 34

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

      def each(&block)
40 41 42 43 44 45
        @array.each(&block)

        return if @overflow
        return if @iterator.nil?

        Gitlab::GitalyClient.migrate(:commit_raw_diffs) do |is_enabled|
46
          if is_enabled && @iterator.is_a?(Gitlab::GitalyClient::DiffStitcher)
47 48 49 50
            each_gitaly_patch(&block)
          else
            each_rugged_patch(&block)
          end
Robert Speicher committed
51
        end
52 53 54 55 56

        @populated = true

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

      def empty?
60 61
        any? # Make sure the iterator has been exercised
        @empty
Robert Speicher committed
62 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
      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 committed
90 91
      alias_method :to_ary, :to_a

Robert Speicher committed
92 93 94 95 96 97 98 99 100 101
      private

      def populate!
        return if @populated

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

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

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

121 122
      def each_rugged_patch
        i = @array.length
Robert Speicher committed
123

124 125
        @iterator.each do |raw|
          @empty = false
Robert Speicher committed
126

127
          if @enforce_limits && i >= max_files
Robert Speicher committed
128 129 130 131
            @overflow = true
            break
          end

132
          expanded = !@enforce_limits || @expanded
Robert Speicher committed
133

134
          diff = Gitlab::Git::Diff.new(raw, expanded: expanded)
Robert Speicher committed
135

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

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

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

          yield @array[i] = diff
151
          i += 1
Robert Speicher committed
152 153 154 155 156
        end
      end
    end
  end
end