Use Gitlab::Redis::ClusterUtil in redis cache store patch

Changelog: other

lib/gitlab/patch/redis_cache_store.rb

@@ -17,15 +17,9 @@ def read_multi_mget(*names) # rubocop:disable Style/ArgumentsForwarding
       def delete_multi_entries(entries, **options)
         return super unless enable_rails_cache_pipeline_patch?
 
-        delete_count = 0
         redis.with do |conn|
-          entries.each_slice(pipeline_batch_size) do |subset|
-            delete_count += conn.pipelined do |pipeline|
-              subset.each { |entry| pipeline.del(entry) }
-            end.sum
-          end
+          ::Gitlab::Redis::ClusterUtil.batch_del(entries, conn)
         end
-        delete_count
       end
 
       # Copied from https://github.com/rails/rails/blob/v6.1.6.1/activesupport/lib/active_support/cache/redis_cache_store.rb
@@ -42,7 +36,7 @@ def patched_read_multi_mget(*names)
 
         values = failsafe(:patched_read_multi_mget, returning: {}) do
           redis.with do |c|
-            pipeline_mget(c, keys)
+            ::Gitlab::Redis::ClusterUtil.batch_get(keys, c)
           end
         end
 
@@ -56,20 +50,8 @@ def patched_read_multi_mget(*names)
         end
       end
 
-      def pipeline_mget(conn, keys)
-        keys.each_slice(pipeline_batch_size).flat_map do |subset|
-          conn.pipelined do |p|
-            subset.each { |key| p.get(key) }
-          end
-        end
-      end
-
       private
 
-      def pipeline_batch_size
-        @pipeline_batch_size ||= [ENV['GITLAB_REDIS_CLUSTER_PIPELINE_BATCH_LIMIT'].to_i, 1000].max
-      end
-
       def enable_rails_cache_pipeline_patch?
         redis.with { |c| ::Gitlab::Redis::ClusterUtil.cluster?(c) }
       end

lib/gitlab/redis/cluster_util.rb

@@ -19,7 +19,7 @@ def cluster?(obj)
 
         def batch_unlink(keys, redis)
           expired_count = 0
-          keys.each_slice(1000) do |subset|
+          keys.each_slice(pipeline_batch_size) do |subset|
             expired_count += redis.pipelined do |pipeline|
               subset.each { |key| pipeline.unlink(key) }
             end.sum
@@ -27,14 +27,30 @@ def batch_unlink(keys, redis)
           expired_count
         end
 
+        def batch_del(keys, redis)
+          expired_count = 0
+          keys.each_slice(pipeline_batch_size) do |subset|
+            expired_count += redis.pipelined do |pipeline|
+              subset.each { |key| pipeline.del(key) }
+            end.sum
+          end
+          expired_count
+        end
+
         # Redis cluster alternative to mget
         def batch_get(keys, redis)
-          keys.each_slice(1000).flat_map do |subset|
+          keys.each_slice(pipeline_batch_size).flat_map do |subset|
             redis.pipelined do |pipeline|
               subset.map { |key| pipeline.get(key) }
             end
           end
         end
+
+        private
+
+        def pipeline_batch_size
+          @pipeline_batch_size ||= [ENV['GITLAB_REDIS_CLUSTER_PIPELINE_BATCH_LIMIT'].to_i, 1000].max
+        end
       end
     end
   end

spec/lib/gitlab/patch/redis_cache_store_spec.rb

@@ -64,25 +64,6 @@
           end
         end
 
-        context 'when GITLAB_REDIS_CLUSTER_PIPELINE_BATCH_LIMIT is smaller than the default' do
-          before do
-            stub_env('GITLAB_REDIS_CLUSTER_PIPELINE_BATCH_LIMIT', 10)
-          end
-
-          it_behaves_like 'read large amount of keys'
-        end
-
-        context 'when GITLAB_REDIS_CLUSTER_PIPELINE_BATCH_LIMIT is larger than the default' do
-          let(:input_size) { 4000 }
-          let(:chunk_size) { 2000 }
-
-          before do
-            stub_env('GITLAB_REDIS_CLUSTER_PIPELINE_BATCH_LIMIT', chunk_size)
-          end
-
-          it_behaves_like 'read large amount of keys'
-        end
-
         it_behaves_like 'read large amount of keys'
       end
     end

spec/lib/gitlab/redis/cluster_util_spec.rb

@@ -6,6 +6,7 @@
   using RSpec::Parameterized::TableSyntax
 
   let(:router_stub) { instance_double(::RedisClient::Cluster::Router) }
+  let(:array) { Array.new(10000, &:to_s) }
 
   before do
     allow(::RedisClient::Cluster::Router).to receive(:new).and_return(router_stub)
@@ -54,4 +55,103 @@
       end
     end
   end
+
+  shared_examples 'batches commands' do
+    it 'calls pipelined multiple times' do
+      Gitlab::Instrumentation::RedisClusterValidator.allow_cross_slot_commands do
+        Gitlab::Redis::Cache.with do |c|
+          expect(c).to receive(:pipelined).exactly(times).times.and_call_original
+
+          described_class.send(cmd, Array.new(size) { |i| i }, c)
+        end
+      end
+    end
+  end
+
+  shared_examples 'batches pipelined commands' do
+    let(:times) { 1 }
+    let(:size) { 1000 }
+
+    it_behaves_like 'batches commands'
+
+    context 'when larger than batch limit' do
+      let(:times) { 2 }
+      let(:size) { 1001 }
+
+      it_behaves_like 'batches commands'
+    end
+
+    context 'when smaller than batch limit' do
+      let(:times) { 1 }
+      let(:size) { 999 }
+
+      it_behaves_like 'batches commands'
+    end
+  end
+
+  describe '.batch_get' do
+    let(:cmd) { :batch_get }
+
+    before do
+      Gitlab::Instrumentation::RedisClusterValidator.allow_cross_slot_commands do
+        Gitlab::Redis::Cache.with do |c|
+          c.pipelined { |p| array.each { |i| p.set(i, i) } }
+        end
+      end
+    end
+
+    it_behaves_like 'batches pipelined commands'
+
+    it 'gets multiple keys' do
+      Gitlab::Instrumentation::RedisClusterValidator.allow_cross_slot_commands do
+        Gitlab::Redis::Cache.with do |c|
+          expect(described_class.batch_get(array, c)).to eq(array)
+        end
+      end
+    end
+  end
+
+  describe '.batch_del' do
+    let(:cmd) { :batch_del }
+
+    before do
+      Gitlab::Instrumentation::RedisClusterValidator.allow_cross_slot_commands do
+        Gitlab::Redis::Cache.with do |c|
+          c.pipelined { |p| array.each { |i| p.set(i, i) } }
+        end
+      end
+    end
+
+    it 'deletes multiple keys' do
+      Gitlab::Instrumentation::RedisClusterValidator.allow_cross_slot_commands do
+        Gitlab::Redis::Cache.with do |c|
+          expect(described_class.batch_del(array, c)).to eq(array.size)
+        end
+      end
+    end
+
+    it_behaves_like 'batches pipelined commands'
+  end
+
+  describe '.batch_unlink' do
+    let(:cmd) { :batch_del }
+
+    before do
+      Gitlab::Instrumentation::RedisClusterValidator.allow_cross_slot_commands do
+        Gitlab::Redis::Cache.with do |c|
+          c.pipelined { |p| array.each { |i| p.set(i, i) } }
+        end
+      end
+    end
+
+    it 'unlinks multiple keys' do
+      Gitlab::Instrumentation::RedisClusterValidator.allow_cross_slot_commands do
+        Gitlab::Redis::Cache.with do |c|
+          expect(described_class.batch_unlink(array, c)).to eq(array.size)
+        end
+      end
+    end
+
+    it_behaves_like 'batches pipelined commands'
+  end
 end