Validate build trace in an exclusive trace lock

This makes sure that traces are not being updated while we validate
them. In case of a deadlock being detected we will increase a relevant
metric and ask GitLab Runner to retry the request.

app/services/ci/update_build_state_service.rb

@@ -2,7 +2,8 @@
 
 module Ci
   class UpdateBuildStateService
-    include Gitlab::Utils::StrongMemoize
+    include ::Gitlab::Utils::StrongMemoize
+    include ::Gitlab::ExclusiveLeaseHelpers
 
     Result = Struct.new(:status, :backoff, keyword_init: true)
 
@@ -18,25 +19,45 @@ module Ci
 
     def execute
       overwrite_trace! if has_trace?
-      create_pending_state! if accept_available?
 
-      if accept_request?
-        accept_build_state!
-      else
-        validate_build_trace!
-        update_build_state!
+      unless accept_available?
+        return update_build_state!
+      end
+
+      ensure_pending_state!
+
+      in_build_trace_lock do
+        process_build_state!
       end
     end
 
     private
 
-    def accept_build_state!
-      if Time.current - pending_state.created_at > ACCEPT_TIMEOUT
-        metrics.increment_trace_operation(operation: :discarded)
+    def overwrite_trace!
+      metrics.increment_trace_operation(operation: :overwrite)
 
-        return update_build_state!
+      build.trace.set(params[:trace]) if Gitlab::Ci::Features.trace_overwrite?
+    end
+
+    def ensure_pending_state!
+      pending_state.created_at
     end
 
+    def process_build_state!
+      if live_chunks_pending?
+        if pending_state_outdated?
+          discard_build_trace!
+          update_build_state!
+        else
+          accept_build_state!
+        end
+      else
+        validate_build_trace!
+        update_build_state!
+      end
+    end
+
+    def accept_build_state!
       build.trace_chunks.live.find_each do |chunk|
         chunk.schedule_to_persist!
       end
@@ -48,26 +69,14 @@ module Ci
       end
     end
 
-    def overwrite_trace!
-      metrics.increment_trace_operation(operation: :overwrite)
-
-      build.trace.set(params[:trace]) if Gitlab::Ci::Features.trace_overwrite?
-    end
-
-    def create_pending_state!
-      pending_state.created_at
-    end
-
     def validate_build_trace!
-      return unless accept_available?
+      if chunks_persisted?
+        metrics.increment_trace_operation(operation: :finalized)
+      end
 
       unless ::Gitlab::Ci::Trace::Checksum.new(build).valid?
         metrics.increment_trace_operation(operation: :invalid)
       end
-
-      return unless chunks_persisted?
-
-      metrics.increment_trace_operation(operation: :finalized)
     end
 
     def update_build_state!
@@ -89,12 +98,24 @@ module Ci
       end
     end
 
+    def discard_build_trace!
+      metrics.increment_trace_operation(operation: :discarded)
+    end
+
     def accept_available?
       !build_running? && has_checksum? && chunks_migration_enabled?
     end
 
-    def accept_request?
-      accept_available? && live_chunks_pending?
+    def live_chunks_pending?
+      build.trace_chunks.live.any?
+    end
+
+    def chunks_persisted?
+      build.trace_chunks.any? && !live_chunks_pending?
+    end
+
+    def pending_state_outdated?
+      Time.current - pending_state.created_at > ACCEPT_TIMEOUT
     end
 
     def build_state
@@ -109,14 +130,6 @@ module Ci
       params.dig(:checksum).present?
     end
 
-    def live_chunks_pending?
-      build.trace_chunks.live.any?
-    end
-
-    def chunks_persisted?
-      build.trace_chunks.any? && !live_chunks_pending?
-    end
-
     def build_running?
       build_state == 'running'
     end
@@ -138,6 +151,14 @@ module Ci
       build.pending_state
     end
 
+    def in_build_trace_lock(&block)
+      build.trace.lock(&block) # rubocop:disable CodeReuse/ActiveRecord
+    rescue FailedToObtainLockError
+      metrics.increment_trace_operation(operation: :locked)
+
+      accept_build_state!
+    end
+
     def chunks_migration_enabled?
       ::Gitlab::Ci::Features.accept_trace?(build.project)
     end

lib/gitlab/ci/trace.rb

@@ -130,6 +130,10 @@ module Gitlab
         end
       end
 
+      def lock(&block)
+        in_write_lock(&block)
+      end
+
       private
 
       def read_stream

lib/gitlab/ci/trace/metrics.rb

@@ -7,7 +7,8 @@ module Gitlab
         extend Gitlab::Utils::StrongMemoize
 
         OPERATIONS = [:appended, :streamed, :chunked, :mutated, :overwrite,
-                      :accepted, :finalized, :discarded, :conflict, :invalid].freeze
+                      :accepted, :finalized, :discarded, :conflict, :locked,
+                      :invalid].freeze
 
         def increment_trace_operation(operation: :unknown)
           unless OPERATIONS.include?(operation)

spec/lib/gitlab/ci/trace_spec.rb

@@ -111,4 +111,13 @@ RSpec.describe Gitlab::Ci::Trace, :clean_gitlab_redis_shared_state do
       end
     end
   end
+
+  describe '#lock' do
+    it 'acquires an exclusive lease on the trace' do
+      trace.lock do
+        expect { trace.lock }
+          .to raise_error ::Gitlab::ExclusiveLeaseHelpers::FailedToObtainLockError
+      end
+    end
+  end
 end

spec/services/ci/update_build_state_service_spec.rb

@@ -135,6 +135,26 @@ RSpec.describe Ci::UpdateBuildStateService do
             .with(operation: :invalid)
         end
       end
+
+      context 'when failed to acquire a build trace lock' do
+        it 'accepts a state update request' do
+          build.trace.lock do
+            result = subject.execute
+
+            expect(result.status).to eq 202
+          end
+        end
+
+        it 'increment locked trace metric' do
+          build.trace.lock do
+            execute_with_stubbed_metrics!
+
+            expect(metrics)
+              .to have_received(:increment_trace_operation)
+              .with(operation: :locked)
+          end
+        end
+      end
     end
 
     context 'when build trace has not been migrated yet' do