Use primary database when modifying trace metadata

This commits adds a new class that disables database load balancing in
EE when reading build trace metadata (chunks, checksums) before a
database write happens. This makes build logs more resilient to a
replication lag problem.

app/models/ci/build_trace_chunk.rb

@@ -154,8 +154,15 @@ module Ci
       in_lock(lock_key, **lock_params) do # exclusive Redis lock is acquired first
         raise FailedToPersistDataError, 'Modifed build trace chunk detected' if has_changes_to_save?
 
-        self.reset.then do |chunk|     # we ensure having latest lock_version
-          chunk.unsafe_persist_data!   # we migrate the data and update data store
+        ##
+        # We use primary data source (in EE it is a primary database) and
+        # reload `lock_version` before we migrate the data to persistent
+        # storage and update data store location of the chunk.
+        #
+        # TODO we also need to stick the chunk to primary
+        #
+        ::Gitlab::Ci::Trace::Metadata.using_primary_source do
+          self.reset.then { |chunk| chunk.unsafe_persist_data! }
         end
       end
     rescue FailedToObtainLockError

ee/lib/ee/gitlab/ci/trace/metadata.rb

+# frozen_string_literal: true
+
+module EE
+  module Gitlab
+    module Ci
+      module Trace
+        module Metadata
+          extend ::Gitlab::Utils::Override
+
+          ##
+          # In EE we are disabling the database load balancing for requests
+          # that attempt to read trace metadata before we actually perform a
+          # write.
+          #
+          # This ensures that we are reading the latest build trace chunks,
+          # pending states and checksums.
+          #
+          override :use_primary
+          def use_primary(&block)
+            ::Gitlab::Database::LoadBalancing::Session
+              .current.use_primary(&block)
+          end
+        end
+      end
+    end
+  end
+end

lib/gitlab/ci/trace/checksum.rb

@@ -30,7 +30,11 @@ module Gitlab
         end
 
         def state_crc32
-          strong_memoize(:state_crc32) { build.pending_state&.crc32 }
+          strong_memoize(:state_crc32) do
+            Trace::Metadata.using_primary_source do
+              build.pending_state&.crc32
+            end
+          end
         end
 
         def chunks_crc32
@@ -59,8 +63,10 @@ module Gitlab
         #
         def trace_chunks
           strong_memoize(:trace_chunks) do
-            build.trace_chunks.persisted
-              .select(::Ci::BuildTraceChunk.metadata_attributes)
+            Trace::Metadata.using_primary_source do
+              build.trace_chunks.persisted
+                .select(::Ci::BuildTraceChunk.metadata_attributes)
+            end
           end
         end
 

lib/gitlab/ci/trace/metadata.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Ci
+    class Trace
+      ##
+      # Class that describes CI/CD build logs metadata (chunks, pending
+      # states, checksums) and their datastore.
+      #
+      class Metadata
+        def self.using_primary_source(&block)
+          self.new.use_primary(&block)
+        end
+
+        ##
+        # This method is overriden in EE to make sure we can disable database
+        # load balancing when reading trace metadata. Here, it simply yields.
+        #
+        def use_primary(&block)
+          yield
+        end
+      end
+    end
+  end
+end
+
+::Gitlab::Ci::Trace::Metadata.prepend_if_ee('EE::Gitlab::Ci::Trace::Metadata')