Optimise trace handling code to use streaming instead of full read

8 files changed, 308 insertions, 87 deletions

diff --git a/lib/api/jobs.rb b/lib/api/jobs.rb
index ffab0aafe59..288b03d940c 100644
--- a/lib/api/jobs.rb
+++ b/lib/api/jobs.rb
@@ -118,7 +118,7 @@ module API
         content_type 'text/plain'
         env['api.format'] = :binary
 
-        trace = build.trace
+        trace = build.trace.raw
         body trace
       end
 
diff --git a/lib/api/runner.rb b/lib/api/runner.rb
index d288369e362..6fbb02cb3aa 100644
--- a/lib/api/runner.rb
+++ b/lib/api/runner.rb
@@ -115,7 +115,7 @@ module API
       put '/:id' do
         job = authenticate_job!
 
-        job.update_attributes(trace: params[:trace]) if params[:trace]
+        job.trace.set(params[:trace]) if params[:trace]
 
         Gitlab::Metrics.add_event(:update_build,
                                   project: job.project.path_with_namespace)
@@ -145,16 +145,14 @@ module API
         content_range = request.headers['Content-Range']
         content_range = content_range.split('-')
 
-        current_length = job.trace_length
-        unless current_length == content_range[0].to_i
-          return error!('416 Range Not Satisfiable', 416, { 'Range' => "0-#{current_length}" })
+        stream_size = job.trace.append(request.body.read, content_range[0].to_i)
+        if stream_size < 0
+          return error!('416 Range Not Satisfiable', 416, { 'Range' => "0-#{-stream_size}" })
         end
 
-        job.append_trace(request.body.read, content_range[0].to_i)
-
         status 202
         header 'Job-Status', job.status
-        header 'Range', "0-#{job.trace_length}"
+        header 'Range', "0-#{stream_size}"
       end
 
       desc 'Authorize artifacts uploading for job' do
diff --git a/lib/api/v3/builds.rb b/lib/api/v3/builds.rb
index 6f97102c6ef..4dd03cdf24b 100644
--- a/lib/api/v3/builds.rb
+++ b/lib/api/v3/builds.rb
@@ -120,7 +120,7 @@ module API
           content_type 'text/plain'
           env['api.format'] = :binary
 
-          trace = build.trace
+          trace = build.trace.raw
           body trace
         end
 
diff --git a/lib/ci/ansi2html.rb b/lib/ci/ansi2html.rb
index b3ccad7b28d..1020452480a 100644
--- a/lib/ci/ansi2html.rb
+++ b/lib/ci/ansi2html.rb
@@ -132,34 +132,54 @@ module Ci
 
       STATE_PARAMS = [:offset, :n_open_tags, :fg_color, :bg_color, :style_mask].freeze
 
-      def convert(raw, new_state)
+      def convert(stream, new_state)
         reset_state
-        restore_state(raw, new_state) if new_state.present?
-
-        start = @offset
-        ansi = raw[@offset..-1]
+        restore_state(new_state, stream) if new_state.present?
+
+        append = false
+        truncated = false
+
+        cur_offset = stream.tell
+        if cur_offset > @offset
+          @offset = cur_offset
+          truncated = true
+        else
+          stream.seek(@offset)
+          append = @offset > 0
+        end
+        start_offset = @offset
 
         open_new_tag
 
-        s = StringScanner.new(ansi)
-        until s.eos?
-          if s.scan(/\e([@-_])(.*?)([@-~])/)
-            handle_sequence(s)
-          elsif s.scan(/\e(([@-_])(.*?)?)?$/)
-            break
-          elsif s.scan(/</)
-            @out << '&lt;'
-          elsif s.scan(/\r?\n/)
-            @out << '<br>'
-          else
-            @out << s.scan(/./m)
+        stream.each_line do |line|
+          s = StringScanner.new(line)
+          until s.eos?
+            if s.scan(/\e([@-_])(.*?)([@-~])/)
+              handle_sequence(s)
+            elsif s.scan(/\e(([@-_])(.*?)?)?$/)
+              break
+            elsif s.scan(/</)
+              @out << '&lt;'
+            elsif s.scan(/\r?\n/)
+              @out << '<br>'
+            else
+              @out << s.scan(/./m)
+            end
+            @offset += s.matched_size
           end
-          @offset += s.matched_size
         end
 
         close_open_tags()
 
-        { state: state, html: @out, text: ansi[0, @offset - start], append: start > 0 }
+        OpenStruct.new(
+          html: @out,
+          state: state,
+          append: append,
+          truncated: truncated,
+          offset: start_offset,
+          size: stream.tell - start_offset,
+          total: stream.size
+        )
       end
 
       def handle_sequence(s)
@@ -240,10 +260,10 @@ module Ci
         Base64.urlsafe_encode64(state.to_json)
       end
 
-      def restore_state(raw, new_state)
+      def restore_state(new_state, stream)
         state = Base64.urlsafe_decode64(new_state)
         state = JSON.parse(state, symbolize_names: true)
-        return if state[:offset].to_i > raw.length
+        return if state[:offset].to_i > stream.size
 
         STATE_PARAMS.each do |param|
           send("#{param}=".to_sym, state[param])
diff --git a/lib/ci/api/builds.rb b/lib/ci/api/builds.rb
index 95cc6308c3b..67b269b330c 100644
--- a/lib/ci/api/builds.rb
+++ b/lib/ci/api/builds.rb
@@ -61,7 +61,7 @@ module Ci
 
           update_runner_info
 
-          build.update_attributes(trace: params[:trace]) if params[:trace]
+          build.trace.set(params[:trace]) if params[:trace]
 
           Gitlab::Metrics.add_event(:update_build,
                                     project: build.project.path_with_namespace)
@@ -92,16 +92,14 @@ module Ci
           content_range = request.headers['Content-Range']
           content_range = content_range.split('-')
 
-          current_length = build.trace_length
-          unless current_length == content_range[0].to_i
-            return error!('416 Range Not Satisfiable', 416, { 'Range' => "0-#{current_length}" })
+          stream_size = build.trace.append(request.body.read, content_range[0].to_i)
+          if stream_size < 0
+            return error!('416 Range Not Satisfiable', 416, { 'Range' => "0-#{-stream_size}" })
           end
 
-          build.append_trace(request.body.read, content_range[0].to_i)
-
           status 202
           header 'Build-Status', build.status
-          header 'Range', "0-#{build.trace_length}"
+          header 'Range', "0-#{stream_size}"
         end
 
         # Authorize artifacts uploading for build - Runners only
diff --git a/lib/gitlab/ci/trace.rb b/lib/gitlab/ci/trace.rb
new file mode 100644
index 00000000000..5b835bb669a
--- /dev/null
+++ b/lib/gitlab/ci/trace.rb
@@ -0,0 +1,136 @@
+module Gitlab
+  module Ci
+    class Trace
+      attr_reader :job
+
+      delegate :old_trace, to: :job
+
+      def initialize(job)
+        @job = job
+      end
+
+      def html(last_lines: nil)
+        read do |stream|
+          stream.html(last_lines: last_lines)
+        end
+      end
+
+      def raw(last_lines: nil)
+        read do |stream|
+          stream.raw(last_lines: last_lines)
+        end
+      end
+
+      def extract_coverage(regex)
+        read do |stream|
+          stream.extract_coverage(regex)
+        end
+      end
+
+      def set(data)
+        write do |stream|
+          data = job.hide_secrets(data)
+          stream.set(data)
+        end
+      end
+
+      def append(data, offset)
+        write do |stream|
+          current_length = stream.size
+          return -current_length unless current_length == offset
+
+          data = job.hide_secrets(data)
+          stream.append(data, offset)
+          stream.size
+        end
+      end
+
+      def exist?
+        current_path.present? || old_trace.present?
+      end
+
+      def read
+        stream = Gitlab::Ci::Trace::Stream.new do
+          if current_path
+            File.open(current_path, "rb")
+          elsif old_trace
+            StringIO.new(old_trace)
+          end
+        end
+
+        yield stream
+      ensure
+        stream&.close
+      end
+
+      def write
+        stream = Gitlab::Ci::Trace::Stream.new do
+          File.open(ensure_path, "a+b")
+        end
+
+        yield(stream).tap do
+          job.touch if job.needs_touch?
+        end
+      ensure
+        stream&.close
+      end
+
+      def erase!
+        paths.each do |trace_path|
+          FileUtils.rm(trace_path, force: true)
+        end
+
+        job.erase_old_trace!
+      end
+
+      private
+
+      def ensure_path
+        return current_path if current_path
+
+        ensure_directory
+        default_path
+      end
+
+      def ensure_directory
+        unless Dir.exist?(default_directory)
+          FileUtils.mkdir_p(default_directory)
+        end
+      end
+
+      def current_path
+        @current_path ||= paths.find do |trace_path|
+          File.exist?(trace_path)
+        end
+      end
+
+      def paths
+        [
+          default_path,
+          deprecated_path
+        ].compact
+      end
+
+      def default_directory
+        File.join(
+          Settings.gitlab_ci.builds_path,
+          job.created_at.utc.strftime("%Y_%m"),
+          job.project_id.to_s
+        )
+      end
+
+      def default_path
+        File.join(default_directory, "#{job.id}.log")
+      end
+
+      def deprecated_path
+        File.join(
+          Settings.gitlab_ci.builds_path,
+          job.created_at.utc.strftime("%Y_%m"),
+          job.project.ci_id.to_s,
+          "#{job.id}.log"
+        ) if job.project&.ci_id
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/ci/trace/stream.rb b/lib/gitlab/ci/trace/stream.rb
new file mode 100644
index 00000000000..2af94e2c60e
--- /dev/null
+++ b/lib/gitlab/ci/trace/stream.rb
@@ -0,0 +1,119 @@
+module Gitlab
+  module Ci
+    class Trace
+      # This was inspired from: http://stackoverflow.com/a/10219411/1520132
+      class Stream
+        BUFFER_SIZE = 4096
+        LIMIT_SIZE = 50.kilobytes
+
+        attr_reader :stream
+
+        delegate :close, :tell, :seek, :size, :path, :truncate, to: :stream, allow_nil: true
+
+        delegate :valid?, to: :stream, as: :present?, allow_nil: true
+
+        def initialize
+          @stream = yield
+        end
+
+        def valid?
+          self.stream.present?
+        end
+
+        def file?
+          self.path.present?
+        end
+
+        def limit(last_bytes = LIMIT_SIZE)
+          stream_size = size
+          if stream_size < last_bytes
+            last_bytes = stream_size
+          end
+          stream.seek(-last_bytes, IO::SEEK_END)
+        end
+
+        def append(data, offset)
+          stream.truncate(offset)
+          stream.seek(0, IO::SEEK_END)
+          stream.write(data)
+          stream.flush()
+        end
+
+        def set(data)
+          truncate(0)
+          stream.write(data)
+          stream.flush()
+        end
+
+        def raw(last_lines: nil)
+          return unless valid?
+
+          if last_lines.to_i > 0
+            read_last_lines(last_lines)
+          else
+            stream.read
+          end
+        end
+
+        def html_with_state(state = nil)
+          ::Ci::Ansi2html.convert(stream, state)
+        end
+
+        def html(last_lines: nil)
+          text = raw(last_lines: last_lines)
+          stream = StringIO.new(text)
+          ::Ci::Ansi2html.convert(stream).html
+        end
+
+        def extract_coverage(regex)
+          return unless valid?
+          return unless regex
+
+          regex = Regexp.new(regex)
+
+          match = ""
+
+          stream.each_line do |line|
+            matches = line.scan(regex)
+            next unless matches.is_a?(Array)
+
+            match = matches.flatten.last
+            coverage = match.gsub(/\d+(\.\d+)?/).first
+            return coverage.to_f if coverage.present?
+          end
+        rescue
+          # if bad regex or something goes wrong we dont want to interrupt transition
+          # so we just silentrly ignore error for now
+        end
+
+        private
+
+        def read_last_lines(last_lines)
+          chunks = []
+          pos = lines = 0
+          max = stream.size
+
+          # We want an extra line to make sure fist line has full contents
+          while lines <= last_lines && pos < max
+            pos += BUFFER_SIZE
+
+            buf =
+              if pos <= max
+                stream.seek(-pos, IO::SEEK_END)
+                stream.read(BUFFER_SIZE)
+              else # Reached the head, read only left
+                stream.seek(0)
+                stream.read(BUFFER_SIZE - (pos - max))
+              end
+
+            lines += buf.count("\n")
+            chunks.unshift(buf)
+          end
+
+          chunks.join.lines.last(last_lines).join
+            .force_encoding(Encoding.default_external)
+        end
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/ci/trace_reader.rb b/lib/gitlab/ci/trace_reader.rb
deleted file mode 100644
index 1d7ddeb3e0f..00000000000
--- a/lib/gitlab/ci/trace_reader.rb
+++ /dev/null
@@ -1,50 +0,0 @@
-module Gitlab
-  module Ci
-    # This was inspired from: http://stackoverflow.com/a/10219411/1520132
-    class TraceReader
-      BUFFER_SIZE = 4096
-
-      attr_accessor :path, :buffer_size
-
-      def initialize(new_path, buffer_size: BUFFER_SIZE)
-        self.path = new_path
-        self.buffer_size = Integer(buffer_size)
-      end
-
-      def read(last_lines: nil)
-        if last_lines
-          read_last_lines(last_lines)
-        else
-          File.read(path)
-        end
-      end
-
-      def read_last_lines(max_lines)
-        File.open(path) do |file|
-          chunks = []
-          pos = lines = 0
-          max = file.size
-
-          # We want an extra line to make sure fist line has full contents
-          while lines <= max_lines && pos < max
-            pos += buffer_size
-
-            buf = if pos <= max
-                    file.seek(-pos, IO::SEEK_END)
-                    file.read(buffer_size)
-                  else # Reached the head, read only left
-                    file.seek(0)
-                    file.read(buffer_size - (pos - max))
-                  end
-
-            lines += buf.count("\n")
-            chunks.unshift(buf)
-          end
-
-          chunks.join.lines.last(max_lines).join
-            .force_encoding(Encoding.default_external)
-        end
-      end
-    end
-  end
-end