Add feature flag for using new Fog implementation

The CI trace Fog implementation was rewritten to be vendor agnostic.
Introduce a feature flag under `ci_trace_new_fog_store` to ensure a
smooth rollout.

app/models/ci/build_trace_chunk.rb

@@ -45,6 +45,10 @@ module Ci
 
       def get_store_class(store)
         @stores ||= {}
+
+        # Can't memoize this because the feature flag may alter this
+        return fog_store_class.new if store == :fog
+
         @stores[store] ||= "Ci::BuildTraceChunks::#{store.capitalize}".constantize.new
       end
 
@@ -74,6 +78,14 @@ module Ci
       def metadata_attributes
         attribute_names - %w[raw_data]
       end
+
+      def fog_store_class
+        if Feature.enabled?(:ci_trace_new_fog_store)
+          Ci::BuildTraceChunks::Fog
+        else
+          Ci::BuildTraceChunks::LegacyFog
+        end
+      end
     end
 
     def data

app/models/ci/build_trace_chunks/legacy_fog.rb

+# frozen_string_literal: true
+
+module Ci
+  module BuildTraceChunks
+    class LegacyFog
+      def available?
+        object_store.enabled
+      end
+
+      def data(model)
+        connection.get_object(bucket_name, key(model))[:body]
+      rescue Excon::Error::NotFound
+        # If the object does not exist in the object storage, this method returns nil.
+      end
+
+      def set_data(model, new_data)
+        connection.put_object(bucket_name, key(model), new_data)
+      end
+
+      def append_data(model, new_data, offset)
+        if offset > 0
+          truncated_data = data(model).to_s.byteslice(0, offset)
+          new_data = truncated_data + new_data
+        end
+
+        set_data(model, new_data)
+        new_data.bytesize
+      end
+
+      def size(model)
+        data(model).to_s.bytesize
+      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/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

config/feature_flags/development/ci_trace_new_fog_store.yml

+---
+name: ci_trace_new_fog_store
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/46209
+rollout_issue_url: 
+type: development
+group: group::testing
+default_enabled: false

spec/models/ci/build_trace_chunk_spec.rb

@@ -135,11 +135,31 @@ RSpec.describe Ci::BuildTraceChunk, :clean_gitlab_redis_shared_state do
     context 'when data_store is fog' do
       let(:data_store) { :fog }
 
-      before do
-        build_trace_chunk.send(:unsafe_set_data!, +'Sample data in fog')
+      context 'when legacy Fog is enabled' do
+        before do
+          stub_feature_flags(ci_trace_new_fog_store: false)
+          build_trace_chunk.send(:unsafe_set_data!, +'Sample data in fog')
+        end
+
+        it { is_expected.to eq('Sample data in fog') }
+
+        it 'returns a LegacyFog store' do
+          expect(described_class.get_store_class(data_store)).to be_a(Ci::BuildTraceChunks::LegacyFog)
+        end
       end
 
-      it { is_expected.to eq('Sample data in fog') }
+      context 'when new Fog is enabled' do
+        before do
+          stub_feature_flags(ci_trace_new_fog_store: true)
+          build_trace_chunk.send(:unsafe_set_data!, +'Sample data in fog')
+        end
+
+        it { is_expected.to eq('Sample data in fog') }
+
+        it 'returns a new Fog store' do
+          expect(described_class.get_store_class(data_store)).to be_a(Ci::BuildTraceChunks::Fog)
+        end
+      end
     end
   end
 

spec/models/ci/build_trace_chunks/legacy_fog_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Ci::BuildTraceChunks::LegacyFog do
+  let(:data_store) { described_class.new }
+
+  before do
+    stub_artifacts_object_storage
+  end
+
+  describe '#available?' do
+    subject { data_store.available? }
+
+    context 'when object storage is enabled' do
+      it { is_expected.to be_truthy }
+    end
+
+    context 'when object storage is disabled' do
+      before do
+        stub_artifacts_object_storage(enabled: false)
+      end
+
+      it { is_expected.to be_falsy }
+    end
+  end
+
+  describe '#data' do
+    subject { data_store.data(model) }
+
+    context 'when data exists' do
+      let(:model) { create(:ci_build_trace_chunk, :fog_with_data, initial_data: 'sample data in fog') }
+
+      it 'returns the data' do
+        is_expected.to eq('sample data in fog')
+      end
+    end
+
+    context 'when data does not exist' do
+      let(:model) { create(:ci_build_trace_chunk, :fog_without_data) }
+
+      it 'returns nil' do
+        expect(data_store.data(model)).to be_nil
+      end
+    end
+  end
+
+  describe '#set_data' do
+    let(:new_data) { 'abc123' }
+
+    context 'when data exists' do
+      let(:model) { create(:ci_build_trace_chunk, :fog_with_data, initial_data: 'sample data in fog') }
+
+      it 'overwrites data' do
+        expect(data_store.data(model)).to eq('sample data in fog')
+
+        data_store.set_data(model, new_data)
+
+        expect(data_store.data(model)).to eq new_data
+      end
+    end
+
+    context 'when data does not exist' do
+      let(:model) { create(:ci_build_trace_chunk, :fog_without_data) }
+
+      it 'sets new data' do
+        expect(data_store.data(model)).to be_nil
+
+        data_store.set_data(model, new_data)
+
+        expect(data_store.data(model)).to eq new_data
+      end
+    end
+  end
+
+  describe '#delete_data' do
+    subject { data_store.delete_data(model) }
+
+    context 'when data exists' do
+      let(:model) { create(:ci_build_trace_chunk, :fog_with_data, initial_data: 'sample data in fog') }
+
+      it 'deletes data' do
+        expect(data_store.data(model)).to eq('sample data in fog')
+
+        subject
+
+        expect(data_store.data(model)).to be_nil
+      end
+    end
+
+    context 'when data does not exist' do
+      let(:model) { create(:ci_build_trace_chunk, :fog_without_data) }
+
+      it 'does nothing' do
+        expect(data_store.data(model)).to be_nil
+
+        subject
+
+        expect(data_store.data(model)).to be_nil
+      end
+    end
+  end
+
+  describe '#size' do
+    context 'when data exists' do
+      let(:model) { create(:ci_build_trace_chunk, :fog_with_data, initial_data: 'üabcd') }
+
+      it 'returns data bytesize correctly' do
+        expect(data_store.size(model)).to eq 6
+      end
+    end
+
+    context 'when data does not exist' do
+      let(:model) { create(:ci_build_trace_chunk, :fog_without_data) }
+
+      it 'returns zero' do
+        expect(data_store.size(model)).to be_zero
+      end
+    end
+  end
+
+  describe '#keys' do
+    subject { data_store.keys(relation) }
+
+    let(:build) { create(:ci_build) }
+    let(:relation) { build.trace_chunks }
+
+    before do
+      create(:ci_build_trace_chunk, :fog_with_data, chunk_index: 0, build: build)
+      create(:ci_build_trace_chunk, :fog_with_data, chunk_index: 1, build: build)
+    end
+
+    it 'returns keys' do
+      is_expected.to eq([[build.id, 0], [build.id, 1]])
+    end
+  end
+
+  describe '#delete_keys' do
+    subject { data_store.delete_keys(keys) }
+
+    let(:build) { create(:ci_build) }
+    let(:relation) { build.trace_chunks }
+    let(:keys) { data_store.keys(relation) }
+
+    before do
+      create(:ci_build_trace_chunk, :fog_with_data, chunk_index: 0, build: build)
+      create(:ci_build_trace_chunk, :fog_with_data, chunk_index: 1, build: build)
+    end
+
+    it 'deletes multiple data' do
+      ::Fog::Storage.new(JobArtifactUploader.object_store_credentials).tap do |connection|
+        expect(connection.get_object('artifacts', "tmp/builds/#{build.id}/chunks/0.log")[:body]).to be_present
+        expect(connection.get_object('artifacts', "tmp/builds/#{build.id}/chunks/1.log")[:body]).to be_present
+      end
+
+      subject
+
+      ::Fog::Storage.new(JobArtifactUploader.object_store_credentials).tap do |connection|
+        expect { connection.get_object('artifacts', "tmp/builds/#{build.id}/chunks/0.log")[:body] }.to raise_error(Excon::Error::NotFound)
+        expect { connection.get_object('artifacts', "tmp/builds/#{build.id}/chunks/1.log")[:body] }.to raise_error(Excon::Error::NotFound)
+      end
+    end
+  end
+end