Allow to store BuildTraceChunks on Object Storage

6 files changed, 189 insertions, 61 deletions

diff --git a/app/models/ci/build_trace_chunk.rb b/app/models/ci/build_trace_chunk.rb
index 4856f10846c..6b89efe8eb0 100644
--- a/app/models/ci/build_trace_chunk.rb
+++ b/app/models/ci/build_trace_chunk.rb
@@ -10,45 +10,48 @@ module Ci
     WriteError = Class.new(StandardError)
 
     CHUNK_SIZE = 128.kilobytes
-    CHUNK_REDIS_TTL = 1.week
     WRITE_LOCK_RETRY = 10
     WRITE_LOCK_SLEEP = 0.01.seconds
     WRITE_LOCK_TTL = 1.minute
 
     enum data_store: {
       redis: 1,
-      db: 2
+      database: 2,
+      fog: 3
     }
 
     class << self
-      def redis_data_key(build_id, chunk_index)
-        "gitlab:ci:trace:#{build_id}:chunks:#{chunk_index}"
+      def all_stores
+        @all_stores ||= self.data_stores.keys
       end
 
-      def redis_data_keys
-        redis.pluck(:build_id, :chunk_index).map do |data|
-          redis_data_key(data.first, data.second)
-        end
+      def persist_store
+        # get first available store from the back of the list
+        all_stores.reverse.find { |store| get_store_class(store).available? }
       end
 
-      def redis_delete_data(keys)
-        return if keys.empty?
-
-        Gitlab::Redis::SharedState.with do |redis|
-          redis.del(keys)
-        end
+      def get_store_class(store)
+        @stores ||= {}
+        @stores[store] ||= "Ci::BuildTraceChunks::#{store.capitalize}".constantize.new
       end
 
       ##
       # FastDestroyAll concerns
       def begin_fast_destroy
-        redis_data_keys
+        all_stores.each_with_object({}) do |result, store|
+          relation = public_send(store)
+          keys = get_store_class(store).keys(relation)
+
+          result[store] = keys if keys.present?
+        end
       end
 
       ##
       # FastDestroyAll concerns
       def finalize_fast_destroy(keys)
-        redis_delete_data(keys)
+        keys.each do |store, value|
+          get_store_class(store).delete_keys(value)
+        end
       end
     end
 
@@ -69,7 +72,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)
 
-      set_data(data.byteslice(0, offset) + new_data)
+      set_data!(data.byteslice(0, offset) + new_data)
     end
 
     def size
@@ -87,51 +90,53 @@ module Ci
     def range
       (start_offset...end_offset)
     end
+    
+    def persisted?
+      !redis?
+    end
 
-    def use_database!
+    def persist!
       in_lock do
-        break if db?
-        break unless size > 0
-
-        self.update!(raw_data: data, data_store: :db)
-        self.class.redis_delete_data([redis_data_key])
+        unsafe_move_to!(self.class.persist_store)
       end
     end
 
     private
 
+    def unsafe_move_to!(new_store)
+      return if data_store == new_store.to_s
+      return unless size > 0
+
+      old_store_class = self.class.get_store_class(data_store)
+
+      self.get_data.tap do |the_data|
+        self.raw_data = nil
+        self.data_store = new_store
+        self.set_data!(the_data)
+      end
+
+      old_store_class.delete_data(self)
+    end
+
     def get_data
-      if redis?
-        redis_data
-      elsif db?
-        raw_data
-      else
-        raise 'Unsupported data store'
-      end&.force_encoding(Encoding::BINARY) # Redis/Database return UTF-8 string as default
+      self.class.get_store_class(data_store).data(self)&.force_encoding(Encoding::BINARY) # Redis/Database return UTF-8 string as default
     end
 
-    def set_data(value)
+    def set_data!(value)
       raise ArgumentError, 'too much data' if value.bytesize > CHUNK_SIZE
 
       in_lock do
-        if redis?
-          redis_set_data(value)
-        elsif db?
-          self.raw_data = value
-        else
-          raise 'Unsupported data store'
-        end
-
+        self.class.get_store_class(data_store).set_data(self, value)
         @data = value
 
         save! if changed?
       end
 
-      schedule_to_db if full?
+      schedule_to_persist if full?
     end
 
-    def schedule_to_db
-      return if db?
+    def schedule_to_persist
+      return if persisted?
 
       Ci::BuildTraceChunkFlushWorker.perform_async(id)
     end
@@ -140,22 +145,6 @@ module Ci
       size == CHUNK_SIZE
     end
 
-    def redis_data
-      Gitlab::Redis::SharedState.with do |redis|
-        redis.get(redis_data_key)
-      end
-    end
-
-    def redis_set_data(data)
-      Gitlab::Redis::SharedState.with do |redis|
-        redis.set(redis_data_key, data, ex: CHUNK_REDIS_TTL)
-      end
-    end
-
-    def redis_data_key
-      self.class.redis_data_key(build_id, chunk_index)
-    end
-
     def in_lock
       write_lock_key = "trace_write:#{build_id}:chunks:#{chunk_index}"
 
diff --git a/app/models/ci/build_trace_chunks/database.rb b/app/models/ci/build_trace_chunks/database.rb
new file mode 100644
index 00000000000..3666d77c790
--- /dev/null
+++ b/app/models/ci/build_trace_chunks/database.rb
@@ -0,0 +1,29 @@
+module Ci
+  module BuildTraceChunks
+    class Database
+      def available?
+        true
+      end
+
+      def keys(relation)
+        []
+      end
+
+      def delete_keys(keys)
+        # no-op
+      end
+
+      def data(model)
+        model.raw_data
+      end
+
+      def set_data(model, data)
+        model.raw_data = data
+      end
+
+      def delete_data(model)
+        model.update_columns(raw_data: nil) unless model.raw_data.nil?
+      end
+    end
+  end
+end
diff --git a/app/models/ci/build_trace_chunks/fog.rb b/app/models/ci/build_trace_chunks/fog.rb
new file mode 100644
index 00000000000..18b09347381
--- /dev/null
+++ b/app/models/ci/build_trace_chunks/fog.rb
@@ -0,0 +1,59 @@
+module Ci
+  module BuildTraceChunks
+    class Fog
+      def available?
+        object_store.enabled
+      end
+
+      def data(model)
+        connection.get_object(bucket_name, key(model)).body
+      end
+
+      def set_data(model, data)
+        connection.put_object(bucket_name, key(model), data)
+      end
+
+      def delete_data(model)
+        delete_keys([[model.build_id, model.chunk_index]])
+      end
+
+      def keys(relation)
+        return [] unless available?
+
+        relation.pluck(:build_id, :chunk_index)
+      end
+
+      def delete_keys(keys)
+        keys.each do |key|
+          connection.delete_object(bucket_name, key_raw(*key))
+        end
+      end
+
+      private
+
+      def key(model)
+        key_raw(model.build_id, model.chunk_index)
+      end
+
+      def key_raw(build_id, chunk_index)
+        "tmp/chunks/builds/#{build_id.to_i}/chunks/#{chunk_index.to_i}.log"
+      end
+
+      def bucket_name
+        return unless available?
+
+        object_store.remote_directory
+      end
+
+      def connection
+        return unless available?
+
+        @connection ||= ::Fog::Storage.new(object_store.connection.to_hash.deep_symbolize_keys)
+      end
+
+      def object_store
+        Gitlab.config.artifacts.object_store
+      end
+    end
+  end
+end
diff --git a/app/models/ci/build_trace_chunks/redis.rb b/app/models/ci/build_trace_chunks/redis.rb
new file mode 100644
index 00000000000..fdb6065e2a0
--- /dev/null
+++ b/app/models/ci/build_trace_chunks/redis.rb
@@ -0,0 +1,51 @@
+module Ci
+  module BuildTraceChunks
+    class Redis
+      CHUNK_REDIS_TTL = 1.week
+
+      def available?
+        true
+      end
+
+      def data(model)
+        Gitlab::Redis::SharedState.with do |redis|
+          redis.get(key(model))
+        end
+      end
+
+      def set_data(model, data)
+        Gitlab::Redis::SharedState.with do |redis|
+          redis.set(key(model), data, ex: CHUNK_REDIS_TTL)
+        end
+      end
+
+      def delete_data(model)
+        delete_keys([[model.build_id, model.chunk_index]])
+      end
+
+      def keys(relation)
+        relation.pluck(:build_id, :chunk_index)
+      end
+
+      def delete_keys(keys)
+        return if keys.empty?
+
+        keys = keys.map { |key| key_raw(*key) }
+
+        Gitlab::Redis::SharedState.with do |redis|
+          redis.del(keys)
+        end
+      end
+
+      private
+
+      def key(model)
+        key_raw(model.build_id, model.chunk_index)
+      end
+
+      def key_raw(build_id, chunk_index)
+        "gitlab:ci:trace:#{build_id.to_i}:chunks:#{chunk_index.to_i}"
+      end
+    end
+  end
+end
diff --git a/app/workers/ci/build_trace_chunk_flush_worker.rb b/app/workers/ci/build_trace_chunk_flush_worker.rb
index 218d6688bd9..8e08ccbc414 100644
--- a/app/workers/ci/build_trace_chunk_flush_worker.rb
+++ b/app/workers/ci/build_trace_chunk_flush_worker.rb
@@ -5,7 +5,7 @@ module Ci
 
     def perform(build_trace_chunk_id)
       ::Ci::BuildTraceChunk.find_by(id: build_trace_chunk_id).try do |build_trace_chunk|
-        build_trace_chunk.use_database!
+        build_trace_chunk.persist!
       end
     end
   end
diff --git a/spec/models/ci/build_trace_chunk_spec.rb b/spec/models/ci/build_trace_chunk_spec.rb
index cbcf1e55979..3f89898ed06 100644
--- a/spec/models/ci/build_trace_chunk_spec.rb
+++ b/spec/models/ci/build_trace_chunk_spec.rb
@@ -294,8 +294,8 @@ describe Ci::BuildTraceChunk, :clean_gitlab_redis_shared_state do
     end
   end
 
-  describe '#use_database!' do
-    subject { build_trace_chunk.use_database! }
+  describe '#persist!' do
+    subject { build_trace_chunk.persist! }
 
     context 'when data_store is redis' do
       let(:data_store) { :redis }