Set retry_attempt more consistently for Error#final?

CHANGELOG.md

@@ -1,5 +1,6 @@
 # Unreleased
 
+- Fix a few corner cases where `RedisClient::Error#final?` was innacurate.
 - hiredis-client: Properly reconnect to the new leader after a sentinel failover.
 
 # 0.26.0

lib/redis_client.rb

@@ -117,6 +117,23 @@ def final?
     end
   end
 
+  module Final
+    def _set_retry_attempt(_retry_attempt)
+    end
+
+    def retry_attempt
+      0
+    end
+
+    def retriable?
+      false
+    end
+
+    def final?
+      true
+    end
+  end
+
   class Error < StandardError
     include HasConfig
     include Retriable
@@ -161,6 +178,7 @@ def initialize(message = nil, code = nil)
   class CommandError < Error
     include HasCommand
     include HasCode
+    include Final
 
     class << self
       def parse(error_message)
@@ -231,6 +249,7 @@ def initialize(config, **)
     @middlewares = config.middlewares_stack.new(self)
     @raw_connection = nil
     @disable_reconnection = false
+    @retry_attempt = nil
   end
 
   def inspect
@@ -736,8 +755,8 @@ def ensure_connected(retryable: true)
       connection = nil
       preferred_error = nil
       begin
+        @retry_attempt = config.retriable?(tries) ? tries : nil
         connection = raw_connection
-        connection.retry_attempt = config.retriable?(tries) ? tries : nil
         if block_given?
           yield connection
         else
@@ -780,6 +799,7 @@ def raw_connection
     if @raw_connection.nil? || !@raw_connection.revalidate
       connect
     end
+    @raw_connection.retry_attempt = @retry_attempt
     @raw_connection
   end
 
@@ -800,6 +820,7 @@ def connect
         )
       end
     end
+    @raw_connection.retry_attempt = @retry_attempt
 
     prelude = config.connection_prelude.dup
 

lib/redis_client/connection_mixin.rb

@@ -92,7 +92,9 @@ def connection_timeout(timeout)
     end
 
     def protocol_error(message)
-      ProtocolError.with_config(message, config)
+      error = ProtocolError.with_config(message, config)
+      error._set_retry_attempt(@retry_attempt)
+      error
     end
 
     def connection_error(message)

lib/redis_client/ruby_connection.rb

@@ -73,6 +73,10 @@ def write(command)
         @io.write(buffer)
       rescue SystemCallError, IOError, OpenSSL::SSL::SSLError => error
         raise connection_error(error.message)
+      rescue Error => error
+        error._set_config(config)
+        error._set_retry_attempt(@retry_attempt)
+        raise error
       end
     end
 
@@ -94,8 +98,8 @@ def read(timeout = nil)
       else
         @io.with_timeout(timeout) { RESP3.load(@io) }
       end
-    rescue RedisClient::RESP3::UnknownType => error
-      raise RedisClient::ProtocolError.with_config(error.message, config)
+    rescue RedisClient::RESP3::Error => error
+      raise protocol_error(error.message)
     rescue SystemCallError, IOError, OpenSSL::SSL::SSLError => error
       raise connection_error(error.message)
     end

test/redis_client/middlewares_test.rb

@@ -129,6 +129,90 @@ def test_final_errors
       TestMiddleware.calls.clear
     end
 
+    def test_final_errors_during_reconnect
+      client = new_client(reconnect_attempts: 1)
+      simulate_network_errors(client, ["PING", "HELLO"]) do
+        assert_raises ConnectionError do
+          client.call("PING")
+        end
+      end
+
+      calls = TestMiddleware.calls.select { |type, _| type == :call }
+      assert_equal 1, calls.size
+
+      call = calls[0]
+      assert_equal :error, call[1]
+      assert_equal ["PING"], call[2]
+      refute_predicate call[3], :final?
+
+      pipeline_calls = TestMiddleware.calls.select { |type, _| type == :pipeline }
+      assert_equal 2, pipeline_calls.size
+
+      failing_pipeline = pipeline_calls[1]
+      assert_equal :error, failing_pipeline[1]
+      assert_equal [["HELLO", "3"]], failing_pipeline[2]
+      assert_predicate failing_pipeline[3], :final?
+    end
+
+    def test_command_error_final
+      tcp_server = TCPServer.new("127.0.0.1", 0)
+      tcp_server.setsockopt(Socket::SOL_SOCKET, Socket::SO_REUSEADDR, true)
+      port = tcp_server.addr[1]
+
+      server_thread = Thread.new do
+        session = tcp_server.accept
+        session.write("-Whoops\r\n")
+        session.close
+      end
+
+      assert_raises CommandError do
+        new_client(host: "127.0.0.1", port: port, reconnect_attempts: 1, protocol: 2).call("PING")
+      end
+
+      calls = TestMiddleware.calls.select { |type, _| type == :call }
+      assert_equal 1, calls.size
+      call = calls[0]
+      assert_equal :error, call[1]
+      assert_equal ["PING"], call[2]
+      assert_predicate call[3], :final?
+    ensure
+      server_thread&.kill
+    end
+
+    def test_protocol_error
+      tcp_server = TCPServer.new("127.0.0.1", 0)
+      tcp_server.setsockopt(Socket::SOL_SOCKET, Socket::SO_REUSEADDR, true)
+      port = tcp_server.addr[1]
+
+      server_thread = Thread.new do
+        2.times do
+          session = tcp_server.accept
+          session.write("garbage\r\n")
+          session.flush
+          session.close
+        end
+      end
+
+      assert_raises ProtocolError do
+        new_client(host: "127.0.0.1", port: port, reconnect_attempts: 1, protocol: 2).call("PING")
+      end
+
+      calls = TestMiddleware.calls.select { |type, _| type == :call }
+      assert_equal 2, calls.size
+
+      call = calls[0]
+      assert_equal :error, call[1]
+      assert_equal ["PING"], call[2]
+      refute_predicate call[3], :final?
+
+      call = calls[1]
+      assert_equal :error, call[1]
+      assert_equal ["PING"], call[2]
+      assert_predicate call[3], :final?
+    ensure
+      server_thread&.kill
+    end
+
     module DummyMiddleware
       def call(command, _config, &_)
         command