wire up retry attempt earlier to handle errors during connection prelude

Author: underwoo16

lib/redis_client.rb

@@ -736,8 +736,7 @@ def ensure_connected(retryable: true)
       connection = nil
       preferred_error = nil
       begin
-        connection = raw_connection
-        connection.retry_attempt = config.retriable?(tries) ? tries : nil
+        connection = raw_connection(config.retriable?(tries) ? tries : nil)
         if block_given?
           yield connection
         else
@@ -776,17 +775,19 @@ def ensure_connected(retryable: true)
     end
   end
 
-  def raw_connection
+  def raw_connection(retry_attempt = nil)
     if @raw_connection.nil? || !@raw_connection.revalidate
-      connect
+      connect(retry_attempt)
     end
+    @raw_connection.retry_attempt = retry_attempt
     @raw_connection
   end
 
-  def connect
+  def connect(retry_attempt = nil)
     @pid = PIDCache.pid
 
     if @raw_connection
+      @raw_connection.retry_attempt = retry_attempt
       @middlewares.connect(config) do
         @raw_connection.reconnect
       end
@@ -799,6 +800,7 @@ def connect
           write_timeout: write_timeout,
         )
       end
+      @raw_connection.retry_attempt = retry_attempt
     end
 
     prelude = config.connection_prelude.dup

lib/redis_client/ruby_connection.rb

@@ -76,6 +76,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
 
@@ -88,6 +92,10 @@ def write_multi(commands)
         @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
 
@@ -101,6 +109,10 @@ def read(timeout = nil)
       raise RedisClient::ProtocolError.with_config(error.message, config)
     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
 
     def measure_round_trip_delay

test/redis_client/middlewares_test.rb

@@ -129,6 +129,31 @@ 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
+
     module DummyMiddleware
       def call(command, _config, &_)
         command