Lock writes to trace stream

1 files changed, 33 insertions, 30 deletions

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