Fix archive method. Fix trace spec

4 files changed, 97 insertions, 20 deletions

diff --git a/lib/gitlab/ci/trace.rb b/lib/gitlab/ci/trace.rb
index c4000cc787a..3dc4848c23d 100644
--- a/lib/gitlab/ci/trace.rb
+++ b/lib/gitlab/ci/trace.rb
@@ -54,7 +54,7 @@ module Gitlab
       end
 
       def exist?
-        trace_artifact&.exists? || current_path.present? || old_trace.present?
+        trace_artifact&.exists? || ChunkedFile::LiveTrace.exist?(job.id) || current_path.present? || old_trace.present?
       end
 
       def read
@@ -62,7 +62,7 @@ module Gitlab
           if trace_artifact
             trace_artifact.open
           elsif ChunkedFile::LiveTrace.exist?(job.id)
-            ChunkedFile::LiveTrace.new(job.id, "rb")
+            ChunkedFile::LiveTrace.new(job.id, nil, "rb")
           elsif current_path
             File.open(current_path, "rb")
           elsif old_trace
@@ -81,7 +81,7 @@ module Gitlab
             if current_path
               current_path
             else
-              ChunkedFile::LiveTrace.new(job.id, "a+b")
+              ChunkedFile::LiveTrace.new(job.id, nil, "a+b")
             end
           else
             File.open(ensure_path, "a+b")
@@ -110,9 +110,12 @@ module Gitlab
         raise ArchiveError, 'Job is not finished yet' unless job.complete?
 
         if ChunkedFile::LiveTrace.exist?(job.id)
-          ChunkedFile::LiveTrace.open(job.id, 'a+b') do |stream|
-            archive_stream!(stream)
-            stream.delete
+          ChunkedFile::LiveTrace.new(job.id, nil, 'a+b') do |live_trace_stream|
+            StringIO.new(live_trace_stream.read, 'rb').tap do |stream|
+              archive_stream!(stream)
+            end
+
+            live_trace_stream.delete
           end
         elsif current_path
           File.open(current_path) do |stream|
diff --git a/lib/gitlab/ci/trace/chunked_file/chunked_io.rb b/lib/gitlab/ci/trace/chunked_file/chunked_io.rb
index f3d3aae5a5b..f9adbffc25a 100644
--- a/lib/gitlab/ci/trace/chunked_file/chunked_io.rb
+++ b/lib/gitlab/ci/trace/chunked_file/chunked_io.rb
@@ -81,7 +81,7 @@ module Gitlab
             end
           end
 
-          def read(length = nil)
+          def read(length = nil, outbuf = nil)
             out = ""
 
             until eof? || (length && out.length >= length)
diff --git a/lib/gitlab/ci/trace/chunked_file/live_trace.rb b/lib/gitlab/ci/trace/chunked_file/live_trace.rb
index 1e1700b43c9..bf918fd4ace 100644
--- a/lib/gitlab/ci/trace/chunked_file/live_trace.rb
+++ b/lib/gitlab/ci/trace/chunked_file/live_trace.rb
@@ -5,7 +5,7 @@ module Gitlab
         class LiveTrace < ChunkedIO
           class << self
             def exist?(job_id)
-              ChunkStore::Redis.chunks_count(job_id) > 0 || ChunkStore::Database.chunks_count(job_id) > 0
+              ChunkedFile::ChunkStore::Redis.chunks_count(job_id) > 0 || ChunkedFile::ChunkStore::Database.chunks_count(job_id) > 0
             end
           end
 
@@ -14,7 +14,7 @@ module Gitlab
           def stash_to_database(store)
             # Once data is filled into redis, move the data to database
             if store.filled?
-              ChunkStore::Database.open(job_id, chunk_index, params_for_store) do |to_store|
+              ChunkedFile::ChunkStore::Database.open(job_id, chunk_index, params_for_store) do |to_store|
                 to_store.write!(store.get)
                 store.delete!
               end
@@ -33,22 +33,22 @@ module Gitlab
           end
 
           def delete
-            ChunkStore::Redis.delete_all(job_id)
-            ChunkStore::Database.delete_all(job_id)
+            ChunkedFile::ChunkStore::Redis.delete_all(job_id)
+            ChunkedFile::ChunkStore::Database.delete_all(job_id)
           end
 
           private
 
           def calculate_size(job_id)
-            ChunkStore::Redis.chunks_size(job_id) +
-              ChunkStore::Database.chunks_size(job_id)
+            ChunkedFile::ChunkStore::Redis.chunks_size(job_id) +
+              ChunkedFile::ChunkStore::Database.chunks_size(job_id)
           end
 
           def chunk_store
             if last_range.include?(tell)
-              ChunkStore::Redis
+              ChunkedFile::ChunkStore::Redis
             else
-              ChunkStore::Database
+              ChunkedFile::ChunkStore::Database
             end
           end
 
diff --git a/spec/lib/gitlab/ci/trace_spec.rb b/spec/lib/gitlab/ci/trace_spec.rb
index 3a9371ed2e8..c246ce9cdf3 100644
--- a/spec/lib/gitlab/ci/trace_spec.rb
+++ b/spec/lib/gitlab/ci/trace_spec.rb
@@ -1,9 +1,13 @@
 require 'spec_helper'
 
-describe Gitlab::Ci::Trace do
+describe Gitlab::Ci::Trace, :clean_gitlab_redis_cache do
   let(:build) { create(:ci_build) }
   let(:trace) { described_class.new(build) }
 
+  before do
+    stub_feature_flags(ci_enable_live_trace: true)
+  end
+
   describe "associations" do
     it { expect(trace).to respond_to(:job) }
     it { expect(trace).to delegate_method(:old_trace).to(:job) }
@@ -403,6 +407,10 @@ describe Gitlab::Ci::Trace do
   describe '#archive!' do
     subject { trace.archive! }
 
+    before do
+      stub_feature_flags(ci_enable_live_trace: false)
+    end
+
     shared_examples 'archive trace file' do
       it do
         expect { subject }.to change { Ci::JobArtifact.count }.by(1)
@@ -455,11 +463,44 @@ describe Gitlab::Ci::Trace do
       end
     end
 
+    shared_examples 'archive trace file in ChunkedIO' do
+      it do
+        expect { subject }.to change { Ci::JobArtifact.count }.by(1)
+
+        build.reload
+        expect(build.trace.exist?).to be_truthy
+        expect(build.job_artifacts_trace.file.exists?).to be_truthy
+        expect(build.job_artifacts_trace.file.filename).to eq('job.log')
+        expect(Gitlab::Ci::Trace::ChunkedFile::LiveTrace.exist?(build.id)).to be_falsy
+        expect(src_checksum)
+          .to eq(Digest::SHA256.file(build.job_artifacts_trace.file.path).hexdigest)
+        expect(build.job_artifacts_trace.file_sha256).to eq(src_checksum)
+      end
+    end
+
+    shared_examples 'source trace in ChunkedIO stays intact' do |error:|
+      it do
+        expect { subject }.to raise_error(error)
+
+        build.reload
+        expect(build.trace.exist?).to be_truthy
+        expect(build.job_artifacts_trace).to be_nil
+        Gitlab::Ci::Trace::ChunkedFile::LiveTrace.new(build.id, nil, 'rb') do |stream|
+          expect(stream.read).to eq(trace_raw)
+        end
+      end
+    end
+
     context 'when job does not have trace artifact' do
       context 'when trace file stored in default path' do
-        let!(:build) { create(:ci_build, :success, :trace_live) }
-        let!(:src_path) { trace.read { |s| return s.path } }
-        let!(:src_checksum) { Digest::SHA256.file(src_path).hexdigest }
+        let(:build) { create(:ci_build, :success, :trace_live) }
+        let(:src_path) { trace.read { |s| return s.path } }
+        let(:src_checksum) { Digest::SHA256.file(src_path).hexdigest }
+
+        before do
+          stub_feature_flags(ci_enable_live_trace: false)
+          build; src_path; src_checksum; # Initialize after set feature flag
+        end
 
         it_behaves_like 'archive trace file'
 
@@ -485,9 +526,11 @@ describe Gitlab::Ci::Trace do
       context 'when trace is stored in database' do
         let(:build) { create(:ci_build, :success) }
         let(:trace_content) { 'Sample trace' }
-        let!(:src_checksum) { Digest::SHA256.hexdigest(trace_content) }
+        let(:src_checksum) { Digest::SHA256.hexdigest(trace_content) }
 
         before do
+          stub_feature_flags(ci_enable_live_trace: false)
+          build; trace_content; src_checksum; # Initialize after set feature flag
           build.update_column(:trace, trace_content)
         end
 
@@ -533,6 +576,37 @@ describe Gitlab::Ci::Trace do
           it_behaves_like 'archive trace in database'
         end
       end
+
+      context 'when trace is stored in ChunkedIO' do
+        let(:build) { create(:ci_build, :success, :trace_live) }
+        let(:trace_raw) { build.trace.raw }
+        let(:src_checksum) { Digest::SHA256.hexdigest(trace_raw) }
+
+        before do
+          stub_feature_flags(ci_enable_live_trace: true)
+          build; trace_raw; src_checksum; # Initialize after set feature flag
+        end
+
+        it_behaves_like 'archive trace file in ChunkedIO'
+
+        context 'when failed to create clone file' do
+          before do
+            allow(IO).to receive(:copy_stream).and_return(0)
+          end
+
+          it_behaves_like 'source trace in ChunkedIO stays intact', error: Gitlab::Ci::Trace::ArchiveError
+        end
+
+        context 'when failed to create job artifact record' do
+          before do
+            allow_any_instance_of(Ci::JobArtifact).to receive(:save).and_return(false)
+            allow_any_instance_of(Ci::JobArtifact).to receive_message_chain(:errors, :full_messages)
+              .and_return(%w[Error Error])
+          end
+
+          it_behaves_like 'source trace in ChunkedIO stays intact', error: ActiveRecord::RecordInvalid
+        end
+      end
     end
 
     context 'when job has trace artifact' do