Lock writes to trace streamlock-trace-writes

8 files changed, 66 insertions, 30 deletions

diff --git a/app/models/ci/build_trace_chunk.rb b/app/models/ci/build_trace_chunk.rb
index 108874b75a6..7c84bd734bb 100644
--- a/app/models/ci/build_trace_chunk.rb
+++ b/app/models/ci/build_trace_chunk.rb
@@ -76,7 +76,7 @@ module Ci
       raise ArgumentError, 'Offset is out of range' if offset > size || offset < 0
       raise ArgumentError, 'Chunk size overflow' if CHUNK_SIZE < (offset + new_data.bytesize)
 
-      in_lock(*lock_params) do # Write opetation is atomic
+      in_lock(*lock_params) do # Write operation is atomic
         unsafe_set_data!(data.byteslice(0, offset) + new_data)
       end
 
@@ -100,7 +100,7 @@ module Ci
     end
 
     def persist_data!
-      in_lock(*lock_params) do # Write opetation is atomic
+      in_lock(*lock_params) do # Write operation is atomic
         unsafe_persist_to!(self.class.persistable_store)
       end
     end
diff --git a/changelogs/unreleased/lock-trace-writes.yml b/changelogs/unreleased/lock-trace-writes.yml
new file mode 100644
index 00000000000..9c5239081b9
--- /dev/null
+++ b/changelogs/unreleased/lock-trace-writes.yml
@@ -0,0 +1,5 @@
+---
+title: Lock writes to trace stream
+merge_request:
+author:
+type: fixed
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..121a73db836 100644
--- a/lib/gitlab/ci/trace.rb
+++ b/lib/gitlab/ci/trace.rb
@@ -83,7 +83,7 @@ module Gitlab
       end
 
       def write(mode)
-        stream = Gitlab::Ci::Trace::Stream.new do
+        stream = Gitlab::Ci::Trace::Stream.new(lock_key: "build:trace:stream:lock:#{job.id}") do
           if trace_artifact
             raise AlreadyArchivedError, 'Could not write to the archived trace'
           elsif current_path
diff --git a/lib/gitlab/ci/trace/stream.rb b/lib/gitlab/ci/trace/stream.rb
index bd40fdf59b1..1db919e2c1c 100644
--- a/lib/gitlab/ci/trace/stream.rb
+++ b/lib/gitlab/ci/trace/stream.rb
@@ -5,10 +5,16 @@ module Gitlab
     class Trace
       # This was inspired from: http://stackoverflow.com/a/10219411/1520132
       class Stream
+        include ::Gitlab::ExclusiveLeaseHelpers
+
         BUFFER_SIZE = 4096
         LIMIT_SIZE = 500.kilobytes
+        LOCK_TTL = 30.seconds
+        LOCK_RETRIES = 2
+        LOCK_SLEEP = 0.001.seconds
 
         attr_reader :stream
+        attr_reader :lock_key
 
         delegate :close, :tell, :seek, :size, :url, :truncate, to: :stream, allow_nil: true
 
@@ -16,9 +22,10 @@ module Gitlab
 
         alias_method :present?, :valid?
 
-        def initialize
+        def initialize(lock_key: nil)
           @stream = yield
           @stream&.binmode
+          @lock_key = lock_key
         end
 
         def valid?
@@ -41,21 +48,13 @@ module Gitlab
         end
 
         def append(data, offset)
-          data = data.force_encoding(Encoding::BINARY)
-
-          stream.truncate(offset)
-          stream.seek(0, IO::SEEK_END)
-          stream.write(data)
-          stream.flush()
+          in_write_lock do
+            unsafe_append(data, offset)
+          end
         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)
@@ -115,6 +114,19 @@ module Gitlab
 
         private
 
+        def in_write_lock(&blk)
+          in_lock(lock_key, ttl: LOCK_TTL, retries: LOCK_RETRIES, sleep_sec: LOCK_SLEEP, &blk)
+        end
+
+        def unsafe_append(data, offset)
+          data = data.force_encoding(Encoding::BINARY)
+
+          stream.seek(offset, IO::SEEK_SET)
+          stream.write(data)
+          stream.truncate(offset + data.bytesize)
+          stream.flush()
+        end
+
         def each_line_with_pos
           stream.seek(0, IO::SEEK_SET)
           stream.each_line do |line|
diff --git a/lib/gitlab/exclusive_lease_helpers.rb b/lib/gitlab/exclusive_lease_helpers.rb
index 4aaf2474763..3c35cc6896b 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 FailedToObtainLockError unless key
+
       lease = Gitlab::ExclusiveLease.new(key, timeout: ttl)
 
       until uuid = lease.try_obtain
diff --git a/spec/lib/gitlab/ci/trace/stream_spec.rb b/spec/lib/gitlab/ci/trace/stream_spec.rb
index 4f49958dd33..c9f2fce8929 100644
--- a/spec/lib/gitlab/ci/trace/stream_spec.rb
+++ b/spec/lib/gitlab/ci/trace/stream_spec.rb
@@ -120,7 +120,7 @@ describe Gitlab::Ci::Trace::Stream, :clean_gitlab_redis_cache do
       let(:tempfile) { Tempfile.new }
 
       let(:stream) do
-        described_class.new do
+        described_class.new(lock_key: 'key') do
           tempfile.write("12345678")
           tempfile.rewind
           tempfile
@@ -136,7 +136,7 @@ describe Gitlab::Ci::Trace::Stream, :clean_gitlab_redis_cache do
 
     context 'when stream is ChunkedIO' do
       let(:stream) do
-        described_class.new do
+        described_class.new(lock_key: 'key') do
           Gitlab::Ci::Trace::ChunkedIO.new(build).tap do |chunked_io|
             chunked_io.write('12345678')
             chunked_io.seek(0, IO::SEEK_SET)
@@ -164,7 +164,7 @@ describe Gitlab::Ci::Trace::Stream, :clean_gitlab_redis_cache do
 
     context 'when stream is StringIO' do
       let(:stream) do
-        described_class.new do
+        described_class.new(lock_key: 'key') do
           StringIO.new("12345678")
         end
       end
@@ -174,7 +174,7 @@ describe Gitlab::Ci::Trace::Stream, :clean_gitlab_redis_cache do
 
     context 'when stream is ChunkedIO' do
       let(:stream) do
-        described_class.new do
+        described_class.new(lock_key: 'key') do
           Gitlab::Ci::Trace::ChunkedIO.new(build).tap do |chunked_io|
             chunked_io.write('12345678')
             chunked_io.seek(0, IO::SEEK_SET)
@@ -257,7 +257,8 @@ describe Gitlab::Ci::Trace::Stream, :clean_gitlab_redis_cache do
         let!(:last_result) { stream.html_with_state }
 
         before do
-          stream.append("5678", 4)
+          data_stream.seek(4, IO::SEEK_SET)
+          data_stream.write("5678")
           stream.seek(0)
         end
 
@@ -271,25 +272,29 @@ describe Gitlab::Ci::Trace::Stream, :clean_gitlab_redis_cache do
     end
 
     context 'when stream is StringIO' do
+      let(:data_stream) do
+        StringIO.new("1234")
+      end
+
       let(:stream) do
-        described_class.new do
-          StringIO.new("1234")
-        end
+        described_class.new { data_stream }
       end
 
       it_behaves_like 'html_with_states'
     end
 
     context 'when stream is ChunkedIO' do
-      let(:stream) do
-        described_class.new do
-          Gitlab::Ci::Trace::ChunkedIO.new(build).tap do |chunked_io|
-            chunked_io.write("1234")
-            chunked_io.seek(0, IO::SEEK_SET)
-          end
+      let(:data_stream) do
+        Gitlab::Ci::Trace::ChunkedIO.new(build).tap do |chunked_io|
+          chunked_io.write("1234")
+          chunked_io.seek(0, IO::SEEK_SET)
         end
       end
 
+      let(:stream) do
+        described_class.new { data_stream }
+      end
+
       it_behaves_like 'html_with_states'
     end
   end
diff --git a/spec/lib/gitlab/exclusive_lease_helpers_spec.rb b/spec/lib/gitlab/exclusive_lease_helpers_spec.rb
index 2e3656b52fb..8f463775a5b 100644
--- a/spec/lib/gitlab/exclusive_lease_helpers_spec.rb
+++ b/spec/lib/gitlab/exclusive_lease_helpers_spec.rb
@@ -11,6 +11,14 @@ describe Gitlab::ExclusiveLeaseHelpers, :clean_gitlab_redis_shared_state do
 
     let(:options) { {} }
 
+    context 'when unique key is not set' do
+      let(:unique_key) { }
+
+      it 'raises an error' do
+        expect { subject }.to raise_error described_class::FailedToObtainLockError
+      end
+    end
+
     context 'when the lease is not obtained yet' do
       before do
         stub_exclusive_lease(unique_key, 'uuid')