pass tries through middleware

Author: underwoo16

README.md

@@ -415,15 +415,15 @@ either globally or on a given configuration instance.
 
 ```ruby
 module MyGlobalRedisInstrumentation
-  def connect(redis_config)
+  def connect(redis_config, retry_attempts)
     MyMonitoringService.instrument("redis.connect") { super }
   end
 
-  def call(command, redis_config)
+  def call(command, redis_config, retry_attempts)
     MyMonitoringService.instrument("redis.query") { super }
   end
 
-  def call_pipelined(commands, redis_config)
+  def call_pipelined(commands, redis_config, retry_attempts)
     MyMonitoringService.instrument("redis.pipeline") { super }
   end
 end
@@ -443,15 +443,15 @@ If middlewares need a client-specific configuration, `Config#custom` can be used
 
 ```ruby
 module MyGlobalRedisInstrumentation
-  def connect(redis_config)
+  def connect(redis_config, retry_attempts)
     MyMonitoringService.instrument("redis.connect", tags: redis_config.custom[:tags]) { super }
   end
 
-  def call(command, redis_config)
+  def call(command, redis_config, retry_attempts)
     MyMonitoringService.instrument("redis.query", tags: redis_config.custom[:tags]) { super }
   end
 
-  def call_pipelined(commands, redis_config)
+  def call_pipelined(commands, redis_config, retry_attempts)
     MyMonitoringService.instrument("redis.pipeline", tags: redis_config.custom[:tags]) { super }
   end
 end

lib/redis_client.rb

@@ -287,17 +287,17 @@ def pubsub
   end
 
   def measure_round_trip_delay
-    ensure_connected do |connection|
-      @middlewares.call(["PING"], config) do
+    ensure_connected do |connection, retry_attempts|
+      @middlewares.call(["PING"], config, retry_attempts) do
         connection.measure_round_trip_delay
       end
     end
   end
 
   def call(*command, **kwargs)
     command = @command_builder.generate(command, kwargs)
-    result = ensure_connected do |connection|
-      @middlewares.call(command, config) do
+    result = ensure_connected do |connection, retry_attempts|
+      @middlewares.call(command, config, retry_attempts) do
         connection.call(command, nil)
       end
     end
@@ -311,8 +311,8 @@ def call(*command, **kwargs)
 
   def call_v(command)
     command = @command_builder.generate(command)
-    result = ensure_connected do |connection|
-      @middlewares.call(command, config) do
+    result = ensure_connected do |connection, retry_attempts|
+      @middlewares.call(command, config, retry_attempts) do
         connection.call(command, nil)
       end
     end
@@ -357,8 +357,8 @@ def call_once_v(command)
   def blocking_call(timeout, *command, **kwargs)
     command = @command_builder.generate(command, kwargs)
     error = nil
-    result = ensure_connected do |connection|
-      @middlewares.call(command, config) do
+    result = ensure_connected do |connection, retry_attempts|
+      @middlewares.call(command, config, retry_attempts) do
         connection.call(command, timeout)
       end
     rescue ReadTimeoutError => error
@@ -377,8 +377,8 @@ def blocking_call(timeout, *command, **kwargs)
   def blocking_call_v(timeout, command)
     command = @command_builder.generate(command)
     error = nil
-    result = ensure_connected do |connection|
-      @middlewares.call(command, config) do
+    result = ensure_connected do |connection, retry_attempts|
+      @middlewares.call(command, config, retry_attempts) do
         connection.call(command, timeout)
       end
     rescue ReadTimeoutError => error
@@ -450,9 +450,9 @@ def pipelined(exception: true)
     if pipeline._size == 0
       []
     else
-      results = ensure_connected(retryable: pipeline._retryable?) do |connection|
+      results = ensure_connected(retryable: pipeline._retryable?) do |connection, retry_attempts|
         commands = pipeline._commands
-        @middlewares.call_pipelined(commands, config) do
+        @middlewares.call_pipelined(commands, config, retry_attempts) do
           connection.call_pipelined(commands, pipeline._timeouts, exception: exception)
         end
       end
@@ -489,9 +489,9 @@ def multi(watch: nil, &block)
       if transaction._empty?
         []
       else
-        ensure_connected(retryable: transaction._retryable?) do |connection|
+        ensure_connected(retryable: transaction._retryable?) do |connection, retry_attempts|
           commands = transaction._commands
-          @middlewares.call_pipelined(commands, config) do
+          @middlewares.call_pipelined(commands, config, retry_attempts) do
             connection.call_pipelined(commands, nil)
           end.last
         end
@@ -707,7 +707,7 @@ def ensure_connected(retryable: true)
 
     if @disable_reconnection
       if block_given?
-        yield @raw_connection
+        yield @raw_connection, 0
       else
         @raw_connection
       end
@@ -716,9 +716,9 @@ def ensure_connected(retryable: true)
       connection = nil
       preferred_error = nil
       begin
-        connection = raw_connection
+        connection = raw_connection(tries)
         if block_given?
-          yield connection
+          yield connection, tries
         else
           connection
         end
@@ -744,7 +744,7 @@ def ensure_connected(retryable: true)
       connection = ensure_connected
       begin
         @disable_reconnection = true
-        yield connection
+        yield connection, 0
       rescue ConnectionError, ProtocolError
         close
         raise
@@ -754,22 +754,22 @@ def ensure_connected(retryable: true)
     end
   end
 
-  def raw_connection
+  def raw_connection(retry_attempts = 0)
     if @raw_connection.nil? || !@raw_connection.revalidate
-      connect
+      connect(retry_attempts)
     end
     @raw_connection
   end
 
-  def connect
+  def connect(retry_attempts = 0)
     @pid = PIDCache.pid
 
     if @raw_connection
-      @middlewares.connect(config) do
+      @middlewares.connect(config, retry_attempts) do
         @raw_connection.reconnect
       end
     else
-      @raw_connection = @middlewares.connect(config) do
+      @raw_connection = @middlewares.connect(config, retry_attempts) do
         config.driver.new(
           config,
           connect_timeout: connect_timeout,
@@ -788,13 +788,13 @@ def connect
     # The connection prelude is deliberately not sent to Middlewares
     if config.sentinel?
       prelude << ["ROLE"]
-      role, = @middlewares.call_pipelined(prelude, config) do
+      role, = @middlewares.call_pipelined(prelude, config, retry_attempts) do
         @raw_connection.call_pipelined(prelude, nil).last
       end
       config.check_role!(role)
     else
       unless prelude.empty?
-        @middlewares.call_pipelined(prelude, config) do
+        @middlewares.call_pipelined(prelude, config, retry_attempts) do
           @raw_connection.call_pipelined(prelude, nil)
         end
       end

lib/redis_client/circuit_breaker.rb

@@ -3,15 +3,15 @@
 class RedisClient
   class CircuitBreaker
     module Middleware
-      def connect(config)
+      def connect(config, _retry_attempts = 0)
         config.circuit_breaker.protect { super }
       end
 
-      def call(_command, config)
+      def call(_command, config, _retry_attempts = 0)
         config.circuit_breaker.protect { super }
       end
 
-      def call_pipelined(_commands, config)
+      def call_pipelined(_commands, config, _retry_attempts = 0)
         config.circuit_breaker.protect { super }
       end
     end

lib/redis_client/middlewares.rb

@@ -8,11 +8,11 @@ def initialize(client)
       @client = client
     end
 
-    def connect(_config)
+    def connect(_config, _retry_attempts = 0)
       yield
     end
 
-    def call(command, _config)
+    def call(command, _config, _retry_attempts = 0)
       yield command
     end
     alias_method :call_pipelined, :call

test/redis_client/middlewares_test.rb

@@ -14,6 +14,7 @@ def setup
       RedisClient.register(TestMiddleware)
       super
       TestMiddleware.calls.clear
+      InstrumentRetryAttemptsMiddleware.calls.clear
     end
 
     def teardown
@@ -22,6 +23,7 @@ def teardown
         RedisClient.const_set(:Middlewares, @original_module)
       end
       TestMiddleware.calls.clear
+      InstrumentRetryAttemptsMiddleware.calls.clear
       super
     end
 
@@ -70,11 +72,11 @@ def test_multi_instrumentation
     end
 
     module DummyMiddleware
-      def call(command, _config, &_)
+      def call(command, _config, _retry_attempts = 0, &_)
         command
       end
 
-      def call_pipelined(commands, _config, &_)
+      def call_pipelined(commands, _config, _retry_attempts = 0, &_)
         commands
       end
     end
@@ -85,6 +87,25 @@ def test_instance_middleware
       assert_equal([["GET", "2"]], second_client.pipelined { |p| p.call("GET", 2) })
     end
 
+    def test_retry_instruments_attempts
+      client = new_client(reconnect_attempts: 1, middlewares: [InstrumentRetryAttemptsMiddleware])
+
+      simulate_network_errors(client, ["PING"]) do
+        client.call("PING")
+      end
+
+      assert_includes InstrumentRetryAttemptsMiddleware.calls, [:call, :error, ["PING"], 0]
+      assert_includes InstrumentRetryAttemptsMiddleware.calls, [:call, :success, ["PING"], 1]
+    end
+
+    def test_connect_instruments_attempts
+      client = new_client(middlewares: [InstrumentRetryAttemptsMiddleware])
+
+      client.call("PING")
+
+      assert_includes InstrumentRetryAttemptsMiddleware.calls, [:connect, :success, 0]
+    end
+
     private
 
     def assert_call(call)
@@ -102,7 +123,7 @@ class << self
       end
       @calls = []
 
-      def connect(config)
+      def connect(config, _retry_attempts = 0)
         result = super
         TestMiddleware.calls << [:connect, :success, result, config]
         result
@@ -111,7 +132,7 @@ def connect(config)
         raise
       end
 
-      def call(command, config)
+      def call(command, config, _retry_attempts = 0)
         result = super
         TestMiddleware.calls << [:call, :success, command, result, config]
         result
@@ -120,7 +141,7 @@ def call(command, config)
         raise
       end
 
-      def call_pipelined(commands, config)
+      def call_pipelined(commands, config, _retry_attempts = 0)
         result = super
         TestMiddleware.calls << [:pipeline, :success, commands, result, config]
         result
@@ -129,5 +150,39 @@ def call_pipelined(commands, config)
         raise
       end
     end
+
+    module InstrumentRetryAttemptsMiddleware
+      class << self
+        attr_accessor :calls
+      end
+      @calls = []
+
+      def connect(config, retry_attempts = 0)
+        result = super
+        InstrumentRetryAttemptsMiddleware.calls << [:connect, :success, retry_attempts]
+        result
+      rescue
+        InstrumentRetryAttemptsMiddleware.calls << [:connect, :error, retry_attempts]
+        raise
+      end
+
+      def call(command, config, retry_attempts = 0)
+        result = super
+        InstrumentRetryAttemptsMiddleware.calls << [:call, :success, command, retry_attempts]
+        result
+      rescue
+        InstrumentRetryAttemptsMiddleware.calls << [:call, :error, command, retry_attempts]
+        raise
+      end
+
+      def call_pipelined(commands, config, retry_attempts = 0)
+        result = super
+        InstrumentRetryAttemptsMiddleware.calls << [:pipeline, :success, commands, retry_attempts]
+        result
+      rescue
+        InstrumentRetryAttemptsMiddleware.calls << [:pipeline, :error, commands, retry_attempts]
+        raise
+      end
+    end
   end
 end