feat: prevent a cross slot error by using pipelining when calling a command with multiple keys (#356)

Author: supercaracal

lib/redis_client/cluster/command.rb

@@ -97,6 +97,12 @@ def exists?(name)
         @commands.key?(::RedisClient::Cluster::NormalizedCmdName.instance.get_by_name(name))
       end
 
+      def determine_key_step(command)
+        name = ::RedisClient::Cluster::NormalizedCmdName.instance.get_by_command(command)
+        # Some commands like EVALSHA have zero as the step in COMMANDS somehow.
+        @commands[name].key_step == 0 ? 1 : @commands[name].key_step
+      end
+
       private
 
       def determine_first_key_position(command) # rubocop:disable Metrics/CyclomaticComplexity
@@ -143,12 +149,6 @@ def determine_last_key_position(command, keys_start) # rubocop:disable Metrics/A
         end
       end
 
-      def determine_key_step(command)
-        name = ::RedisClient::Cluster::NormalizedCmdName.instance.get_by_command(command)
-        # Some commands like EVALSHA have zero as the step in COMMANDS somehow.
-        @commands[name].key_step == 0 ? 1 : @commands[name].key_step
-      end
-
       def determine_optional_key_position(command, option_name) # rubocop:disable Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity
         idx = command&.flatten&.map(&:to_s)&.map(&:downcase)&.index(option_name&.downcase)
         idx.nil? ? 0 : idx + 1

lib/redis_client/cluster/router.rb

@@ -10,13 +10,19 @@
 require 'redis_client/cluster/normalized_cmd_name'
 require 'redis_client/cluster/transaction'
 require 'redis_client/cluster/optimistic_locking'
+require 'redis_client/cluster/pipeline'
 require 'redis_client/cluster/error_identification'
 
 class RedisClient
   class Cluster
     class Router
       ZERO_CURSOR_FOR_SCAN = '0'
       TSF = ->(f, x) { f.nil? ? x : f.call(x) }.curry
+      MULTIPLE_KEYS_COMMAND_TO_PIPELINE = {
+        'mset' => 'set',
+        'mget' => 'get',
+        'del' => 'del'
+      }.freeze
 
       def initialize(config, concurrent_worker, pool: nil, **kwargs)
         @config = config.dup
@@ -48,6 +54,8 @@ def send_command(method, command, *args, &block) # rubocop:disable Metrics/AbcSi
         when 'script'   then send_script_command(method, command, args, &block)
         when 'pubsub'   then send_pubsub_command(method, command, args, &block)
         when 'watch'    then send_watch_command(command, &block)
+        when 'mset', 'mget', 'del'
+          send_multiple_keys_command(cmd, method, command, args, &block)
         when 'acl', 'auth', 'bgrewriteaof', 'bgsave', 'quit', 'save'
           @node.call_all(method, command, args).first.then(&TSF.call(block))
         when 'flushall', 'flushdb'
@@ -326,6 +334,25 @@ def send_watch_command(command)
         end
       end
 
+      def send_multiple_keys_command(cmd, method, command, args, &block) # rubocop:disable Metrics/AbcSize, Metrics/CyclomaticComplexity
+        if command.size < 3 || !::RedisClient::Cluster::KeySlotConverter.extract_hash_tag(command[1]).empty? # rubocop:disable Style/IfUnlessModifier
+          return try_send(assign_node(command), method, command, args, &block)
+        end
+
+        single_key_cmd = MULTIPLE_KEYS_COMMAND_TO_PIPELINE[cmd]
+        key_step = @command.determine_key_step(cmd)
+        seed = @config.use_replica? && @config.replica_affinity == :random ? nil : Random.new_seed
+        pipeline = ::RedisClient::Cluster::Pipeline.new(self, @command_builder, @concurrent_worker, exception: true, seed: seed)
+        command[1..].each_slice(key_step) { |*v| pipeline.call(single_key_cmd, *v) }
+        replies = pipeline.execute
+        result = case cmd
+                 when 'mset' then replies.first
+                 when 'del' then replies.sum
+                 else replies
+                 end
+        block_given? ? yield(result) : result
+      end
+
       def update_cluster_info!
         @node.reload!
       end

test/redis_client/cluster/test_command.rb

@@ -147,6 +147,20 @@ def test_exists?
         end
       end
 
+      def test_determine_key_step
+        cmd = ::RedisClient::Cluster::Command.load(@raw_clients)
+        [
+          { name: 'MSET', want: 2 },
+          { name: 'MGET', want: 1 },
+          { name: 'DEL', want: 1 },
+          { name: 'EVALSHA', want: 1 }
+        ].each_with_index do |c, idx|
+          msg = "Case: #{idx}"
+          got = cmd.determine_key_step(c[:name])
+          assert_equal(c[:want], got, msg)
+        end
+      end
+
       def test_determine_first_key_position
         cmd = ::RedisClient::Cluster::Command.load(@raw_clients)
         [

test/redis_client/test_cluster.rb

@@ -686,6 +686,26 @@ def test_with_method
         assert_raises(NotImplementedError) { @client.with }
       end
 
+      def test_dedicated_multiple_keys_command
+        [
+          { command: %w[MSET key1 val1], want: 'OK' },
+          { command: %w[MGET key1], want: %w[val1] },
+          { command: %w[DEL key1], want: 1 },
+          { command: %w[MSET {key}1 val1 {key}2 val2], want: 'OK' },
+          { command: %w[MGET {key}1 {key}2], want: %w[val1 val2] },
+          { command: %w[DEL {key}1 {key}2], want: 2 },
+          { command: %w[MSET key1 val1 key2 val2], want: 'OK' },
+          { command: %w[MGET key1 key2], want: %w[val1 val2] },
+          { command: %w[DEL key1 key2], want: 2 },
+          { command: %w[MSET key1 val1 key2 val2], block: ->(r) { "#{r}!" }, want: 'OK!' },
+          { command: %w[MGET key1 key2], block: ->(r) { r.map { |e| "#{e}!" } }, want: %w[val1! val2!] },
+          { command: %w[DEL key1 key2], block: ->(r) { r == 2 }, want: true }
+        ].each_with_index do |c, i|
+          block = c.key?(:block) ? c.fetch(:block) : nil
+          assert_equal(c.fetch(:want), @client.call_v(c.fetch(:command), &block), i + 1)
+        end
+      end
+
       def test_dedicated_commands # rubocop:disable Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity
         10.times { |i| @client.call('SET', "key#{i}", i) }
         wait_for_replication