Retry a build trace chunk migration in case of an exception

This commit adds a retry logic to build chunks migration. In case of not
being able to migrate a chunk, we want to retry this operation only
once.

app/models/ci/build_trace_chunk.rb

@@ -126,12 +126,18 @@ module Ci
       Ci::BuildTraceChunkFlushWorker.perform_async(id)
     end
 
-    def persisted?
-      !redis?
-    end
-
-    def live?
-      redis?
+    ##
+    # It is possible that we run into two concurrent migrations. It might
+    # happen that a chunk gets migrated after being loaded by another worker
+    # but before the worker acquires a lock to perform the migration.
+    #
+    # We want to reset a chunk in that case and retry migration. If it fails
+    # again, we want to re-raise the exception.
+    #
+    def flush!
+      persist_data!
+    rescue FailedToPersistDataError
+      self.reset.persist_data!
     end
 
     ##
@@ -144,6 +150,14 @@ module Ci
         build.trace_chunks.maximum(:chunk_index).to_i == chunk_index
     end
 
+    def persisted?
+      !redis?
+    end
+
+    def live?
+      redis?
+    end
+
     def <=>(other)
       return unless self.build_id == other.build_id
 

app/workers/ci/build_trace_chunk_flush_worker.rb

@@ -9,9 +9,7 @@ module Ci
 
     # rubocop: disable CodeReuse/ActiveRecord
     def perform(chunk_id)
-      ::Ci::BuildTraceChunk.find_by(id: chunk_id).try do |chunk|
-        chunk.persist_data!
-      end
+      ::Ci::BuildTraceChunk.find_by(id: chunk_id).try(&:flush!)
     end
     # rubocop: enable CodeReuse/ActiveRecord
   end

spec/models/ci/build_trace_chunk_spec.rb

@@ -780,6 +780,51 @@ RSpec.describe Ci::BuildTraceChunk, :clean_gitlab_redis_shared_state do
     end
   end
 
+  describe '#flush!' do
+    context 'when chunk can be flushed without problems' do
+      before do
+        allow(build_trace_chunk).to receive(:persist_data!)
+      end
+
+      it 'completes migration successfully' do
+        expect { build_trace_chunk.flush! }.not_to raise_error
+      end
+    end
+
+    context 'when the flush operation fails at first' do
+      it 'retries reloads the chunk' do
+        expect(build_trace_chunk)
+          .to receive(:persist_data!)
+          .and_raise(described_class::FailedToPersistDataError)
+          .ordered
+        expect(build_trace_chunk).to receive(:reset)
+          .and_return(build_trace_chunk)
+          .ordered
+        expect(build_trace_chunk)
+          .to receive(:persist_data!)
+          .ordered
+
+        build_trace_chunk.flush!
+      end
+    end
+
+    context 'when the flush constatly fails' do
+      before do
+        allow(build_trace_chunk)
+          .to receive(:persist_data!)
+          .and_raise(described_class::FailedToPersistDataError)
+      end
+
+      it 'attems to reset the chunk but eventually fails too' do
+        expect(build_trace_chunk).to receive(:reset)
+          .and_return(build_trace_chunk)
+
+        expect { build_trace_chunk.flush! }
+          .to raise_error(described_class::FailedToPersistDataError)
+      end
+    end
+  end
+
   describe 'comparable build trace chunks' do
     describe '#<=>' do
       context 'when chunks are associated with different builds' do