Add ability to store CI trace chunks on dedicated Redis

For now this capability is hidden behind an ops feature flag:
dedicated_redis_trace_chunks.

app/models/ci/build_trace_chunk.rb

@@ -14,7 +14,13 @@ module Ci
 
     belongs_to :build, class_name: "Ci::Build", foreign_key: :build_id
 
-    default_value_for :data_store, :redis
+    default_value_for :data_store do
+      if Feature.enabled?(:dedicated_redis_trace_chunks, type: :ops)
+        :redis_trace_chunks
+      else
+        :redis
+      end
+    end
 
     after_create { metrics.increment_trace_operation(operation: :chunked) }
 
@@ -28,17 +34,19 @@ module Ci
     DATA_STORES = {
       redis: 1,
       database: 2,
-      fog: 3
+      fog: 3,
+      redis_trace_chunks: 4
     }.freeze
 
     STORE_TYPES = DATA_STORES.keys.to_h do |store|
-      [store, "Ci::BuildTraceChunks::#{store.capitalize}".constantize]
+      [store, "Ci::BuildTraceChunks::#{store.to_s.camelize}".constantize]
     end.freeze
+    LIVE_STORES = %i[redis redis_trace_chunks].freeze
 
     enum data_store: DATA_STORES
 
-    scope :live, -> { redis }
-    scope :persisted, -> { not_redis.order(:chunk_index) }
+    scope :live, -> { where(data_store: LIVE_STORES) }
+    scope :persisted, -> { where.not(data_store: LIVE_STORES).order(:chunk_index) }
 
     class << self
       def all_stores
@@ -192,7 +200,7 @@ module Ci
     end
 
     def flushed?
-      !redis?
+      !live?
     end
 
     def migrated?
@@ -200,7 +208,7 @@ module Ci
     end
 
     def live?
-      redis?
+      LIVE_STORES.include?(data_store.to_sym)
     end
 
     def <=>(other)

app/models/ci/build_trace_chunks/redis.rb

@@ -2,88 +2,11 @@
 
 module Ci
   module BuildTraceChunks
-    class Redis
-      CHUNK_REDIS_TTL = 1.week
-      LUA_APPEND_CHUNK = <<~EOS
-        local key, new_data, offset = KEYS[1], ARGV[1], ARGV[2]
-        local length = new_data:len()
-        local expire = #{CHUNK_REDIS_TTL.seconds}
-        local current_size = redis.call("strlen", key)
-        offset = tonumber(offset)
-
-        if offset == 0 then
-          -- overwrite everything
-          redis.call("set", key, new_data, "ex", expire)
-          return redis.call("strlen", key)
-        elseif offset > current_size then
-          -- offset range violation
-          return -1
-        elseif offset + length >= current_size then
-          -- efficiently append or overwrite and append
-          redis.call("expire", key, expire)
-          return redis.call("setrange", key, offset, new_data)
-        else
-          -- append and truncate
-          local current_data = redis.call("get", key)
-          new_data = current_data:sub(1, offset) .. new_data
-          redis.call("set", key, new_data, "ex", expire)
-          return redis.call("strlen", key)
-        end
-      EOS
-
-      def data(model)
-        Gitlab::Redis::SharedState.with do |redis|
-          redis.get(key(model))
-        end
-      end
-
-      def set_data(model, new_data)
-        Gitlab::Redis::SharedState.with do |redis|
-          redis.set(key(model), new_data, ex: CHUNK_REDIS_TTL)
-        end
-      end
-
-      def append_data(model, new_data, offset)
-        Gitlab::Redis::SharedState.with do |redis|
-          redis.eval(LUA_APPEND_CHUNK, keys: [key(model)], argv: [new_data, offset])
-        end
-      end
-
-      def size(model)
-        Gitlab::Redis::SharedState.with do |redis|
-          redis.strlen(key(model))
-        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|
-          # https://gitlab.com/gitlab-org/gitlab/-/issues/224171
-          Gitlab::Instrumentation::RedisClusterValidator.allow_cross_slot_commands do
-            redis.del(keys)
-          end
-        end
-      end
-
+    class Redis < RedisBase
       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}"
+      def with_redis
+        Gitlab::Redis::SharedState.with { |redis| yield(redis) }
       end
     end
   end

app/models/ci/build_trace_chunks/redis_base.rb

+# frozen_string_literal: true
+
+module Ci
+  module BuildTraceChunks
+    class RedisBase
+      CHUNK_REDIS_TTL = 1.week
+      LUA_APPEND_CHUNK = <<~EOS
+        local key, new_data, offset = KEYS[1], ARGV[1], ARGV[2]
+        local length = new_data:len()
+        local expire = #{CHUNK_REDIS_TTL.seconds}
+        local current_size = redis.call("strlen", key)
+        offset = tonumber(offset)
+
+        if offset == 0 then
+          -- overwrite everything
+          redis.call("set", key, new_data, "ex", expire)
+          return redis.call("strlen", key)
+        elseif offset > current_size then
+          -- offset range violation
+          return -1
+        elseif offset + length >= current_size then
+          -- efficiently append or overwrite and append
+          redis.call("expire", key, expire)
+          return redis.call("setrange", key, offset, new_data)
+        else
+          -- append and truncate
+          local current_data = redis.call("get", key)
+          new_data = current_data:sub(1, offset) .. new_data
+          redis.call("set", key, new_data, "ex", expire)
+          return redis.call("strlen", key)
+        end
+      EOS
+
+      def data(model)
+        with_redis do |redis|
+          redis.get(key(model))
+        end
+      end
+
+      def set_data(model, new_data)
+        with_redis do |redis|
+          redis.set(key(model), new_data, ex: CHUNK_REDIS_TTL)
+        end
+      end
+
+      def append_data(model, new_data, offset)
+        with_redis do |redis|
+          redis.eval(LUA_APPEND_CHUNK, keys: [key(model)], argv: [new_data, offset])
+        end
+      end
+
+      def size(model)
+        with_redis do |redis|
+          redis.strlen(key(model))
+        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) }
+
+        with_redis do |redis|
+          # https://gitlab.com/gitlab-org/gitlab/-/issues/224171
+          Gitlab::Instrumentation::RedisClusterValidator.allow_cross_slot_commands do
+            redis.del(keys)
+          end
+        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

app/models/ci/build_trace_chunks/redis_trace_chunks.rb

+# frozen_string_literal: true
+
+module Ci
+  module BuildTraceChunks
+    class RedisTraceChunks < RedisBase
+      private
+
+      def with_redis
+        Gitlab::Redis::TraceChunks.with { |redis| yield(redis) }
+      end
+    end
+  end
+end

config/feature_flags/ops/dedicated_redis_trace_chunks.yml

+---
+name: dedicated_redis_trace_chunks
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/62938
+rollout_issue_url: https://gitlab.com/gitlab-com/gl-infra/scalability/-/issues/1096
+milestone: '14.0'
+type: ops
+group: team::Scalability
+default_enabled: false

spec/models/ci/build_trace_chunk_spec.rb

@@ -2,13 +2,13 @@
 
 require 'spec_helper'
 
-RSpec.describe Ci::BuildTraceChunk, :clean_gitlab_redis_shared_state do
+RSpec.describe Ci::BuildTraceChunk, :clean_gitlab_redis_shared_state, :clean_gitlab_redis_trace_chunks do
   include ExclusiveLeaseHelpers
 
   let_it_be(:build) { create(:ci_build, :running) }
 
   let(:chunk_index) { 0 }
-  let(:data_store) { :redis }
+  let(:data_store) { :redis_trace_chunks }
   let(:raw_data) { nil }
 
   let(:build_trace_chunk) do
@@ -22,6 +22,13 @@ RSpec.describe Ci::BuildTraceChunk, :clean_gitlab_redis_shared_state do
     stub_artifacts_object_storage
   end
 
+  def redis_instance
+    {
+      redis: Gitlab::Redis::SharedState,
+      redis_trace_chunks: Gitlab::Redis::TraceChunks
+    }[data_store]
+  end
+
   describe 'chunk creation' do
     let(:metrics) { spy('metrics') }
 
@@ -85,7 +92,7 @@ RSpec.describe Ci::BuildTraceChunk, :clean_gitlab_redis_shared_state do
     end
 
     def external_data_counter
-      Gitlab::Redis::SharedState.with do |redis|
+      redis_instance.with do |redis|
         redis.scan_each(match: "gitlab:ci:trace:*:chunks:*").to_a.size
       end
     end
@@ -101,24 +108,16 @@ RSpec.describe Ci::BuildTraceChunk, :clean_gitlab_redis_shared_state do
     subject { described_class.all_stores }
 
     it 'returns a correctly ordered array' do
-      is_expected.to eq(%i[redis database fog])
-    end
-
-    it 'returns redis store as the lowest precedence' do
-      expect(subject.first).to eq(:redis)
-    end
-
-    it 'returns fog store as the highest precedence' do
-      expect(subject.last).to eq(:fog)
+      is_expected.to eq(%i[redis database fog redis_trace_chunks])
     end
   end
 
   describe '#data' do
     subject { build_trace_chunk.data }
 
-    context 'when data_store is redis' do
-      let(:data_store) { :redis }
+    where(:data_store) { %i[redis redis_trace_chunks] }
 
+    with_them do
       before do
         build_trace_chunk.send(:unsafe_set_data!, +'Sample data in redis')
       end
@@ -148,6 +147,22 @@ RSpec.describe Ci::BuildTraceChunk, :clean_gitlab_redis_shared_state do
     end
   end
 
+  describe '#data_store' do
+    subject { described_class.new.data_store }
+
+    context 'default value' do
+      it { expect(subject).to eq('redis_trace_chunks') }
+
+      context 'when dedicated_redis_trace_chunks is disabled' do
+        before do
+          stub_feature_flags(dedicated_redis_trace_chunks: false)
+        end
+
+        it { expect(subject).to eq('redis') }
+      end
+    end
+  end
+
   describe '#get_store_class' do
     using RSpec::Parameterized::TableSyntax
 
@@ -155,6 +170,7 @@ RSpec.describe Ci::BuildTraceChunk, :clean_gitlab_redis_shared_state do
       :redis | Ci::BuildTraceChunks::Redis
       :database | Ci::BuildTraceChunks::Database
       :fog | Ci::BuildTraceChunks::Fog
+      :redis_trace_chunks | Ci::BuildTraceChunks::RedisTraceChunks
     end
 
     with_them do
@@ -302,9 +318,9 @@ RSpec.describe Ci::BuildTraceChunk, :clean_gitlab_redis_shared_state do
       end
     end
 
-    context 'when data_store is redis' do
-      let(:data_store) { :redis }
+    where(:data_store) { %i[redis redis_trace_chunks] }
 
+    with_them do
       context 'when there are no data' do
         let(:data) { +'' }
 
@@ -441,8 +457,9 @@ RSpec.describe Ci::BuildTraceChunk, :clean_gitlab_redis_shared_state do
       end
     end
 
-    context 'when data_store is redis' do
-      let(:data_store) { :redis }
+    where(:data_store) { %i[redis redis_trace_chunks] }
+
+    with_them do
       let(:data) { +'Sample data in redis' }
 
       before do
@@ -475,9 +492,9 @@ RSpec.describe Ci::BuildTraceChunk, :clean_gitlab_redis_shared_state do
   describe '#size' do
     subject { build_trace_chunk.size }
 
-    context 'when data_store is redis' do
-      let(:data_store) { :redis }
+    where(:data_store) { %i[redis redis_trace_chunks] }
 
+    with_them do
       context 'when data exists' do
         let(:data) { +'Sample data in redis' }
 
@@ -537,9 +554,14 @@ RSpec.describe Ci::BuildTraceChunk, :clean_gitlab_redis_shared_state do
 
     subject { build_trace_chunk.persist_data! }
 
-    context 'when data_store is redis' do
-      let(:data_store) { :redis }
+    where(:data_store, :redis_class) do
+      [
+        [:redis, Ci::BuildTraceChunks::Redis],
+        [:redis_trace_chunks, Ci::BuildTraceChunks::RedisTraceChunks]
+      ]
+    end
 
+    with_them do
       context 'when data exists' do
         before do
           build_trace_chunk.send(:unsafe_set_data!, data)
@@ -549,15 +571,15 @@ RSpec.describe Ci::BuildTraceChunk, :clean_gitlab_redis_shared_state do
           let(:data) { +'a' * described_class::CHUNK_SIZE }
 
           it 'persists the data' do
-            expect(build_trace_chunk.redis?).to be_truthy
-            expect(Ci::BuildTraceChunks::Redis.new.data(build_trace_chunk)).to eq(data)
+            expect(build_trace_chunk.data_store).to eq(data_store.to_s)
+            expect(redis_class.new.data(build_trace_chunk)).to eq(data)
             expect(Ci::BuildTraceChunks::Database.new.data(build_trace_chunk)).to be_nil
             expect(Ci::BuildTraceChunks::Fog.new.data(build_trace_chunk)).to be_nil
 
             subject
 
             expect(build_trace_chunk.fog?).to be_truthy
-            expect(Ci::BuildTraceChunks::Redis.new.data(build_trace_chunk)).to be_nil
+            expect(redis_class.new.data(build_trace_chunk)).to be_nil
             expect(Ci::BuildTraceChunks::Database.new.data(build_trace_chunk)).to be_nil
             expect(Ci::BuildTraceChunks::Fog.new.data(build_trace_chunk)).to eq(data)
           end
@@ -575,8 +597,8 @@ RSpec.describe Ci::BuildTraceChunk, :clean_gitlab_redis_shared_state do
           it 'does not persist the data and the orignal data is intact' do
             expect { subject }.to raise_error(described_class::FailedToPersistDataError)
 
-            expect(build_trace_chunk.redis?).to be_truthy
-            expect(Ci::BuildTraceChunks::Redis.new.data(build_trace_chunk)).to eq(data)
+            expect(build_trace_chunk.data_store).to eq(data_store.to_s)
+            expect(redis_class.new.data(build_trace_chunk)).to eq(data)
             expect(Ci::BuildTraceChunks::Database.new.data(build_trace_chunk)).to be_nil
             expect(Ci::BuildTraceChunks::Fog.new.data(build_trace_chunk)).to be_nil
           end
@@ -810,7 +832,7 @@ RSpec.describe Ci::BuildTraceChunk, :clean_gitlab_redis_shared_state do
 
     shared_examples_for 'deletes all build_trace_chunk and data in redis' do
       it 'deletes all build_trace_chunk and data in redis', :sidekiq_might_not_need_inline do
-        Gitlab::Redis::SharedState.with do |redis|
+        redis_instance.with do |redis|
           expect(redis.scan_each(match: "gitlab:ci:trace:*:chunks:*").to_a.size).to eq(3)
         end
 
@@ -820,7 +842,7 @@ RSpec.describe Ci::BuildTraceChunk, :clean_gitlab_redis_shared_state do
 
         expect(described_class.count).to eq(0)
 
-        Gitlab::Redis::SharedState.with do |redis|
+        redis_instance.with do |redis|
           expect(redis.scan_each(match: "gitlab:ci:trace:*:chunks:*").to_a.size).to eq(0)
         end
       end
@@ -902,4 +924,38 @@ RSpec.describe Ci::BuildTraceChunk, :clean_gitlab_redis_shared_state do
       end
     end
   end
+
+  describe '#live?' do
+    subject { build_trace_chunk.live? }
+
+    where(:data_store, :value) do
+      [
+        [:redis, true],
+        [:redis_trace_chunks, true],
+        [:database, false],
+        [:fog, false]
+      ]
+    end
+
+    with_them do
+      it { is_expected.to eq(value) }
+    end
+  end
+
+  describe '#flushed?' do
+    subject { build_trace_chunk.flushed? }
+
+    where(:data_store, :value) do
+      [
+        [:redis, false],
+        [:redis_trace_chunks, false],
+        [:database, true],
+        [:fog, true]
+      ]
+    end
+
+    with_them do
+      it { is_expected.to eq(value) }
+    end
+  end
 end

spec/services/ci/append_build_trace_service_spec.rb

@@ -44,7 +44,7 @@ RSpec.describe Ci::AppendBuildTraceService do
 
       expect(::Gitlab::ErrorTracking)
         .to receive(:log_exception)
-        .with(anything, hash_including(chunk_index: 0, chunk_store: 'redis'))
+        .with(anything, hash_including(chunk_index: 0, chunk_store: 'redis_trace_chunks'))
 
       result = described_class
         .new(build, content_range: '0-128')

spec/support/redis.rb

@@ -30,4 +30,12 @@ RSpec.configure do |config|
 
     redis_queues_cleanup!
   end
+
+  config.around(:each, :clean_gitlab_redis_trace_chunks) do |example|
+    redis_trace_chunks_cleanup!
+
+    example.run
+
+    redis_trace_chunks_cleanup!
+  end
 end

spec/support/redis/redis_helpers.rb

@@ -17,4 +17,9 @@ module RedisHelpers
   def redis_shared_state_cleanup!
     Gitlab::Redis::SharedState.with(&:flushall)
   end
+
+  # Usage: CI trace chunks
+  def redis_trace_chunks_cleanup!
+    Gitlab::Redis::TraceChunks.with(&:flushall)
+  end
 end