Lock writes to trace stream

4 files changed, 42 insertions, 38 deletions

diff --git a/lib/api/api.rb b/lib/api/api.rb
index 8e259961828..449faf5f8da 100644
--- a/lib/api/api.rb
+++ b/lib/api/api.rb
@@ -50,6 +50,10 @@ module API
       rack_response({ 'message' => '404 Not found' }.to_json, 404)
     end
 
+    rescue_from ::Gitlab::ExclusiveLeaseHelpers::FailedToObtainLockError do
+      rack_response({ 'message' => '409 Conflict: Resource lock' }.to_json, 409)
+    end
+
     rescue_from UploadedFile::InvalidPathError do |e|
       rack_response({ 'message' => e.message }.to_json, 400)
     end
diff --git a/lib/gitlab/ci/trace.rb b/lib/gitlab/ci/trace.rb
index 8eccd262db9..bf5f2a31f0e 100644
--- a/lib/gitlab/ci/trace.rb
+++ b/lib/gitlab/ci/trace.rb
@@ -3,9 +3,11 @@
 module Gitlab
   module Ci
     class Trace
-      include ExclusiveLeaseGuard
+      include ::Gitlab::ExclusiveLeaseHelpers
 
-      LEASE_TIMEOUT = 1.hour
+      LOCK_TTL = 1.minute
+      LOCK_RETRIES = 2
+      LOCK_SLEEP = 0.001.seconds
 
       ArchiveError = Class.new(StandardError)
       AlreadyArchivedError = Class.new(StandardError)
@@ -82,24 +84,10 @@ module Gitlab
         stream&.close
       end
 
-      def write(mode)
-        stream = Gitlab::Ci::Trace::Stream.new do
-          if trace_artifact
-            raise AlreadyArchivedError, 'Could not write to the archived trace'
-          elsif current_path
-            File.open(current_path, mode)
-          elsif Feature.enabled?('ci_enable_live_trace')
-            Gitlab::Ci::Trace::ChunkedIO.new(job)
-          else
-            File.open(ensure_path, mode)
-          end
+      def write(mode, &blk)
+        in_write_lock do
+          unsafe_write!(mode, &blk)
         end
-
-        yield(stream).tap do
-          job.touch if job.needs_touch?
-        end
-      ensure
-        stream&.close
       end
 
       def erase!
@@ -117,13 +105,33 @@ module Gitlab
       end
 
       def archive!
-        try_obtain_lease do
+        in_write_lock do
           unsafe_archive!
         end
       end
 
       private
 
+      def unsafe_write!(mode, &blk)
+        stream = Gitlab::Ci::Trace::Stream.new do
+          if trace_artifact
+            raise AlreadyArchivedError, 'Could not write to the archived trace'
+          elsif current_path
+            File.open(current_path, mode)
+          elsif Feature.enabled?('ci_enable_live_trace')
+            Gitlab::Ci::Trace::ChunkedIO.new(job)
+          else
+            File.open(ensure_path, mode)
+          end
+        end
+
+        yield(stream).tap do
+          job.touch if job.needs_touch?
+        end
+      ensure
+        stream&.close
+      end
+
       def unsafe_archive!
         raise AlreadyArchivedError, 'Could not archive again' if trace_artifact
         raise ArchiveError, 'Job is not finished yet' unless job.complete?
@@ -146,6 +154,11 @@ module Gitlab
         end
       end
 
+      def in_write_lock(&blk)
+        lock_key = "trace:write:lock:#{job.id}"
+        in_lock(lock_key, ttl: LOCK_TTL, retries: LOCK_RETRIES, sleep_sec: LOCK_SLEEP, &blk)
+      end
+
       def archive_stream!(stream)
         clone_file!(stream, JobArtifactUploader.workhorse_upload_path) do |clone_path|
           create_build_trace!(job, clone_path)
@@ -226,16 +239,6 @@ module Gitlab
       def trace_artifact
         job.job_artifacts_trace
       end
-
-      # For ExclusiveLeaseGuard concern
-      def lease_key
-        @lease_key ||= "trace:archive:#{job.id}"
-      end
-
-      # For ExclusiveLeaseGuard concern
-      def lease_timeout
-        LEASE_TIMEOUT
-      end
     end
   end
 end
diff --git a/lib/gitlab/ci/trace/stream.rb b/lib/gitlab/ci/trace/stream.rb
index bd40fdf59b1..0f23b95ba15 100644
--- a/lib/gitlab/ci/trace/stream.rb
+++ b/lib/gitlab/ci/trace/stream.rb
@@ -43,19 +43,14 @@ module Gitlab
         def append(data, offset)
           data = data.force_encoding(Encoding::BINARY)
 
-          stream.truncate(offset)
-          stream.seek(0, IO::SEEK_END)
+          stream.seek(offset, IO::SEEK_SET)
           stream.write(data)
+          stream.truncate(offset + data.bytesize)
           stream.flush()
         end
 
         def set(data)
-          data = data.force_encoding(Encoding::BINARY)
-
-          stream.seek(0, IO::SEEK_SET)
-          stream.write(data)
-          stream.truncate(data.bytesize)
-          stream.flush()
+          append(data, 0)
         end
 
         def raw(last_lines: nil)
diff --git a/lib/gitlab/exclusive_lease_helpers.rb b/lib/gitlab/exclusive_lease_helpers.rb
index 4aaf2474763..7961d4bbd6e 100644
--- a/lib/gitlab/exclusive_lease_helpers.rb
+++ b/lib/gitlab/exclusive_lease_helpers.rb
@@ -12,6 +12,8 @@ module Gitlab
     #       because it holds the connection until all `retries` is consumed.
     #       This could potentially eat up all connection pools.
     def in_lock(key, ttl: 1.minute, retries: 10, sleep_sec: 0.01.seconds)
+      raise ArgumentError, 'Key needs to be specified' unless key
+
       lease = Gitlab::ExclusiveLease.new(key, timeout: ttl)
 
       until uuid = lease.try_obtain