1 files changed, 29 insertions, 10 deletions

diff --git a/app/models/ci/build_trace_chunk.rb b/app/models/ci/build_trace_chunk.rb
index 0a7a0e0772b..407802baf09 100644
--- a/app/models/ci/build_trace_chunk.rb
+++ b/app/models/ci/build_trace_chunk.rb
@@ -75,18 +75,16 @@ module Ci
 
     def append(new_data, offset)
       raise ArgumentError, 'New data is missing' unless new_data
-      raise ArgumentError, 'Offset is out of range' if offset > size || offset < 0
+      raise ArgumentError, 'Offset is out of range' if offset < 0 || offset > size
       raise ArgumentError, 'Chunk size overflow' if CHUNK_SIZE < (offset + new_data.bytesize)
 
-      in_lock(*lock_params) do # Write operation is atomic
-        unsafe_set_data!(data.byteslice(0, offset) + new_data)
-      end
+      in_lock(*lock_params) { unsafe_append_data!(new_data, offset) }
 
       schedule_to_persist if full?
     end
 
     def size
-      data&.bytesize.to_i
+      @size ||= current_store.size(self) || data&.bytesize
     end
 
     def start_offset
@@ -118,7 +116,7 @@ module Ci
         raise FailedToPersistDataError, 'Data is not fulfilled in a bucket'
       end
 
-      old_store_class = self.class.get_store_class(data_store)
+      old_store_class = current_store
 
       self.raw_data = nil
       self.data_store = new_store
@@ -128,16 +126,33 @@ module Ci
     end
 
     def get_data
-      self.class.get_store_class(data_store).data(self)&.force_encoding(Encoding::BINARY) # Redis/Database return UTF-8 string as default
-    rescue Excon::Error::NotFound
-      # If the data store is :fog and the file does not exist in the object storage, this method returns nil.
+      current_store.data(self)&.force_encoding(Encoding::BINARY) # Redis/Database return UTF-8 string as default
     end
 
     def unsafe_set_data!(value)
       raise ArgumentError, 'New data size exceeds chunk size' if value.bytesize > CHUNK_SIZE
 
-      self.class.get_store_class(data_store).set_data(self, value)
+      current_store.set_data(self, value)
+
       @data = value
+      @size = value.bytesize
+
+      save! if changed?
+    end
+
+    def unsafe_append_data!(value, offset)
+      new_size = value.bytesize + offset
+
+      if new_size > CHUNK_SIZE
+        raise ArgumentError, 'New data size exceeds chunk size'
+      end
+
+      current_store.append_data(self, value, offset).then do |stored|
+        raise ArgumentError, 'Trace appended incorrectly' if stored != new_size
+      end
+
+      @data = nil
+      @size = new_size
 
       save! if changed?
     end
@@ -156,6 +171,10 @@ module Ci
       size == CHUNK_SIZE
     end
 
+    def current_store
+      self.class.get_store_class(data_store)
+    end
+
     def lock_params
       ["trace_write:#{build_id}:chunks:#{chunk_index}",
        { ttl: WRITE_LOCK_TTL,