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,18 @@ 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
   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 +799,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