fix: handle redirection errors in pipelining (#131)

Author: supercaracal

.rubocop.yml

@@ -26,7 +26,7 @@ Metrics/MethodLength:
     - 'test/**/*'
 
 Metrics/BlockLength:
-  Max: 25
+  Max: 40
   Exclude:
     - 'test/**/*'
 

lib/redis_client/cluster/node/latency_replica.rb

@@ -34,7 +34,7 @@ def find_node_key_of_replica(primary_node_key, seed: nil) # rubocop:disable Lint
 
         def any_replica_node_key(seed: nil)
           random = seed.nil? ? Random : Random.new(seed)
-          @existed_replicas.sample(random: random).first
+          @existed_replicas.sample(random: random)&.first
         end
 
         private

lib/redis_client/cluster/pipeline.rb

@@ -2,70 +2,141 @@
 
 require 'redis_client'
 require 'redis_client/cluster/errors'
+require 'redis_client/connection_mixin'
 require 'redis_client/middlewares'
 require 'redis_client/pooled'
 
 class RedisClient
   class Cluster
     class Pipeline
+      class Extended < ::RedisClient::Pipeline
+        attr_reader :outer_indices
+
+        def initialize(command_builder)
+          super
+          @outer_indices = nil
+        end
+
+        def add_outer_index(index)
+          @outer_indices ||= []
+          @outer_indices << index
+        end
+
+        def get_inner_index(outer_index)
+          @outer_indices&.find_index(outer_index)
+        end
+
+        def get_callee_method(inner_index)
+          if @timeouts.is_a?(Array) && !@timeouts[inner_index].nil?
+            :blocking_call_v
+          elsif _retryable?
+            :call_once_v
+          else
+            :call_v
+          end
+        end
+
+        def get_command(inner_index)
+          @commands.is_a?(Array) ? @commands[inner_index] : nil
+        end
+
+        def get_timeout(inner_index)
+          @timeouts.is_a?(Array) ? @timeouts[inner_index] : nil
+        end
+
+        def get_block(inner_index)
+          @blocks.is_a?(Array) ? @blocks[inner_index] : nil
+        end
+      end
+
+      ::RedisClient::ConnectionMixin.module_eval do
+        def call_pipelined_aware_of_redirection(commands, timeouts) # rubocop:disable Metrics/AbcSize
+          size = commands.size
+          results = Array.new(commands.size)
+          @pending_reads += size
+          write_multi(commands)
+
+          redirection_indices = nil
+          size.times do |index|
+            timeout = timeouts && timeouts[index]
+            result = read(timeout)
+            @pending_reads -= 1
+            if result.is_a?(CommandError)
+              result._set_command(commands[index])
+              if result.message.start_with?('MOVED', 'ASK')
+                redirection_indices ||= []
+                redirection_indices << index
+              end
+            end
+
+            results[index] = result
+          end
+
+          return results if redirection_indices.nil?
+
+          err = ::RedisClient::Cluster::Pipeline::RedirectionNeeded.new
+          err.replies = results
+          err.indices = redirection_indices
+          raise err
+        end
+      end
+
       ReplySizeError = Class.new(::RedisClient::Error)
+
+      class RedirectionNeeded < ::RedisClient::Error
+        attr_accessor :replies, :indices
+      end
+
       MAX_THREADS = Integer(ENV.fetch('REDIS_CLIENT_MAX_THREADS', 5))
 
       def initialize(router, command_builder, seed: Random.new_seed)
         @router = router
         @command_builder = command_builder
         @seed = seed
-        @pipelines = @indices = nil
+        @pipelines = nil
         @size = 0
       end
 
       def call(*args, **kwargs, &block)
         command = @command_builder.generate(args, kwargs)
         node_key = @router.find_node_key(command, seed: @seed)
-        get_pipeline(node_key).call_v(command, &block)
-        index_pipeline(node_key)
+        append_pipeline(node_key).call_v(command, &block)
       end
 
       def call_v(args, &block)
         command = @command_builder.generate(args)
         node_key = @router.find_node_key(command, seed: @seed)
-        get_pipeline(node_key).call_v(command, &block)
-        index_pipeline(node_key)
+        append_pipeline(node_key).call_v(command, &block)
       end
 
       def call_once(*args, **kwargs, &block)
         command = @command_builder.generate(args, kwargs)
         node_key = @router.find_node_key(command, seed: @seed)
-        get_pipeline(node_key).call_once_v(command, &block)
-        index_pipeline(node_key)
+        append_pipeline(node_key).call_once_v(command, &block)
       end
 
       def call_once_v(args, &block)
         command = @command_builder.generate(args)
         node_key = @router.find_node_key(command, seed: @seed)
-        get_pipeline(node_key).call_once_v(command, &block)
-        index_pipeline(node_key)
+        append_pipeline(node_key).call_once_v(command, &block)
       end
 
       def blocking_call(timeout, *args, **kwargs, &block)
         command = @command_builder.generate(args, kwargs)
         node_key = @router.find_node_key(command, seed: @seed)
-        get_pipeline(node_key).blocking_call_v(timeout, command, &block)
-        index_pipeline(node_key)
+        append_pipeline(node_key).blocking_call_v(timeout, command, &block)
       end
 
       def blocking_call_v(timeout, args, &block)
         command = @command_builder.generate(args)
         node_key = @router.find_node_key(command, seed: @seed)
-        get_pipeline(node_key).blocking_call_v(timeout, command, &block)
-        index_pipeline(node_key)
+        append_pipeline(node_key).blocking_call_v(timeout, command, &block)
       end
 
       def empty?
         @size.zero?
       end
 
-      # TODO: https://github.com/redis-rb/redis-cluster-client/issues/37 handle redirections
       def execute # rubocop:disable Metrics/AbcSize, Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity
         all_replies = errors = nil
         @pipelines&.each_slice(MAX_THREADS) do |chuncked_pipelines|
@@ -77,40 +148,47 @@ def execute # rubocop:disable Metrics/AbcSize, Metrics/CyclomaticComplexity, Met
               raise ReplySizeError, "commands: #{pl._size}, replies: #{replies.size}" if pl._size != replies.size
 
               Thread.current.thread_variable_set(:replies, replies)
+            rescue ::RedisClient::Cluster::Pipeline::RedirectionNeeded => e
+              Thread.current.thread_variable_set(:redirection_needed, e)
             rescue StandardError => e
               Thread.current.thread_variable_set(:error, e)
             end
           end
 
           threads.each do |t|
             t.join
+
             if t.thread_variable?(:replies)
               all_replies ||= Array.new(@size)
-              @indices[t.thread_variable_get(:node_key)].each_with_index { |gi, i| all_replies[gi] = t.thread_variable_get(:replies)[i] }
+              @pipelines[t.thread_variable_get(:node_key)]
+                .outer_indices
+                .each_with_index { |outer, inner| all_replies[outer] = t.thread_variable_get(:replies)[inner] }
+            elsif t.thread_variable?(:redirection_needed)
+              all_replies ||= Array.new(@size)
+              pipeline = @pipelines[t.thread_variable_get(:node_key)]
+              err = t.thread_variable_get(:redirection_needed)
+              err.indices.each { |i| err.replies[i] = handle_redirection(err.replies[i], pipeline, i) }
+              pipeline.outer_indices.each_with_index { |outer, inner| all_replies[outer] = err.replies[inner] }
             elsif t.thread_variable?(:error)
               errors ||= {}
               errors[t.thread_variable_get(:node_key)] = t.thread_variable_get(:error)
             end
           end
         end
 
-        return all_replies if errors.nil?
+        raise ::RedisClient::Cluster::ErrorCollection, errors unless errors.nil?
 
-        raise ::RedisClient::Cluster::ErrorCollection, errors
+        all_replies
       end
 
       private
 
-      def get_pipeline(node_key)
+      def append_pipeline(node_key)
         @pipelines ||= {}
-        @pipelines[node_key] ||= ::RedisClient::Pipeline.new(@command_builder)
-      end
-
-      def index_pipeline(node_key)
-        @indices ||= {}
-        @indices[node_key] ||= []
-        @indices[node_key] << @size
+        @pipelines[node_key] ||= ::RedisClient::Cluster::Pipeline::Extended.new(@command_builder)
+        @pipelines[node_key].add_outer_index(@size)
         @size += 1
+        @pipelines[node_key]
       end
 
       def do_pipelining(client, pipeline)
@@ -125,12 +203,52 @@ def send_pipeline(client, pipeline)
         results = client.send(:ensure_connected, retryable: pipeline._retryable?) do |connection|
           commands = pipeline._commands
           ::RedisClient::Middlewares.call_pipelined(commands, client.config) do
-            connection.call_pipelined(commands, pipeline._timeouts)
+            connection.call_pipelined_aware_of_redirection(commands, pipeline._timeouts)
           end
         end
 
         pipeline._coerce!(results)
       end
+
+      def handle_redirection(err, pipeline, inner_index)
+        return err unless err.is_a?(::RedisClient::CommandError)
+
+        if err.message.start_with?('MOVED')
+          node = @router.assign_redirection_node(err.message)
+          try_redirection(node, pipeline, inner_index)
+        elsif err.message.start_with?('ASK')
+          node = @router.assign_asking_node(err.message)
+          try_asking(node) ? try_redirection(node, pipeline, inner_index) : err
+        else
+          err
+        end
+      end
+
+      def try_redirection(node, pipeline, inner_index)
+        redirect_command(node, pipeline, inner_index)
+      rescue StandardError => e
+        e
+      end
+
+      def redirect_command(node, pipeline, inner_index)
+        method = pipeline.get_callee_method(inner_index)
+        command = pipeline.get_command(inner_index)
+        timeout = pipeline.get_timeout(inner_index)
+        block = pipeline.get_block(inner_index)
+        args = timeout.nil? ? [] : [timeout]
+
+        if block.nil?
+          @router.try_send(node, method, command, args)
+        else
+          @router.try_send(node, method, command, args, &block)
+        end
+      end
+
+      def try_asking(node)
+        node.call('ASKING') == 'OK'
+      rescue StandardError
+        false
+      end
     end
   end
 end

lib/redis_client/cluster/router.rb

@@ -172,7 +172,7 @@ def find_node_key(command, seed: nil)
       def find_node(node_key, retry_count: 3)
         @node.find_by(node_key)
       rescue ::RedisClient::Cluster::Node::ReloadNeeded
-        raise ::RedieClient::Cluster::NodeMightBeDown if retry_count <= 0
+        raise ::RedisClient::Cluster::NodeMightBeDown if retry_count <= 0
 
         update_cluster_info!
         retry_count -= 1
@@ -183,6 +183,18 @@ def command_exists?(name)
         @command.exists?(name)
       end
 
+      def assign_redirection_node(err_msg)
+        _, slot, node_key = err_msg.split
+        slot = slot.to_i
+        @node.update_slot(slot, node_key)
+        find_node(node_key)
+      end
+
+      def assign_asking_node(err_msg)
+        _, _, node_key = err_msg.split
+        find_node(node_key)
+      end
+
       private
 
       def send_wait_command(method, command, args, retry_count: 3, &block)
@@ -258,18 +270,6 @@ def send_pubsub_command(method, command, args, &block) # rubocop:disable Metrics
         end
       end
 
-      def assign_redirection_node(err_msg)
-        _, slot, node_key = err_msg.split
-        slot = slot.to_i
-        @node.update_slot(slot, node_key)
-        find_node(node_key)
-      end
-
-      def assign_asking_node(err_msg)
-        _, _, node_key = err_msg.split
-        find_node(node_key)
-      end
-
       def fetch_cluster_info(config, pool: nil, **kwargs)
         node_info = ::RedisClient::Cluster::Node.load_info(config.per_node_key, **kwargs)
         node_addrs = node_info.map { |info| ::RedisClient::Cluster::NodeKey.hashify(info[:node_key]) }

test/redis_client/test_cluster.rb

@@ -140,16 +140,22 @@ def test_pipelined
       end
 
       def test_pipelined_with_errors
-        assert_raises(::RedisClient::Cluster::ErrorCollection) do
-          @client.pipelined do |pipeline|
-            10.times do |i|
-              pipeline.call('SET', "string#{i}", i)
-              pipeline.call('SET', "string#{i}", i, 'too many args')
-              pipeline.call('SET', "string#{i}", i + 10)
-            end
+        got = @client.pipelined do |pipeline|
+          10.times do |i|
+            pipeline.call('SET', "string#{i}", i)
+            pipeline.call('SET', "string#{i}", i, 'too many args')
+            pipeline.call('SET', "string#{i}", i + 10)
           end
         end
 
+        assert_equal(30, got.size)
+
+        10.times do |i|
+          assert_equal('OK', got[(3 * i) + 0])
+          assert_instance_of(::RedisClient::CommandError, got[(3 * i) + 1])
+          assert_equal('OK', got[(3 * i) + 2])
+        end
+
         wait_for_replication
 
         10.times { |i| assert_equal((i + 10).to_s, @client.call('GET', "string#{i}")) }

test/test_against_cluster_state.rb

@@ -47,8 +47,6 @@ def test_the_state_of_cluster_resharding
     end
 
     def test_the_state_of_cluster_resharding_with_pipelining
-      skip('TODO: https://github.com/redis-rb/redis-cluster-client/issues/37')
-
       do_resharding_test do |keys|
         values = @client.pipelined do |pipeline|
           keys.each { |key| pipeline.call('GET', key) }