feat(common): add custom predicate matching to RetryPolicy

gapic-common/lib/gapic/common/retry_policy.rb

@@ -45,9 +45,17 @@ class RetryPolicy
       # @param retry_codes [Array<String|Integer>] List of retry codes.
       # @param timeout [Numeric] Timeout threshold value in seconds.
       # @param jitter [Numeric] Random jitter added to the delay in seconds.
+      # @param retry_predicate [Proc, nil] The predicate to evaluate whether to retry on a given error. Optional.
+      #   If the predicate is specified, it is run first. If it returns nil, the decision on
+      #   whether to retry is made on the basis of retry_codes. Otherwise, the truthiness of
+      #   the return value determines whether to retry.
       #
-      def initialize initial_delay: nil, max_delay: nil, multiplier: nil, retry_codes: nil, timeout: nil, jitter: nil
+      def initialize initial_delay: nil, max_delay: nil, multiplier: nil, retry_codes: nil, timeout: nil,
+                     jitter: nil, retry_predicate: nil
         raise ArgumentError, "jitter cannot be negative" if jitter&.negative?
+        if retry_predicate && !retry_predicate.respond_to?(:call)
+          raise ArgumentError, "retry_predicate must respond to :call"
+        end
 
         # Instance values are set as `nil` to determine whether values are overriden from default.
         @initial_delay = initial_delay
@@ -56,6 +64,7 @@ def initialize initial_delay: nil, max_delay: nil, multiplier: nil, retry_codes:
         @retry_codes = convert_codes retry_codes
         @timeout = timeout
         @jitter = jitter
+        @retry_predicate = retry_predicate
         start!
       end
 
@@ -89,6 +98,19 @@ def jitter
         @jitter || DEFAULT_JITTER
       end
 
+      ##
+      # The predicate to evaluate whether to retry on a given error. Optional.
+      #
+      # When a predicate is specified:
+      # 1. The predicate is executed first on the error.
+      # 2. If it returns nil, the decision on whether to retry is made on the basis of retry_codes.
+      # 3. Otherwise, the truthiness of the return value determines whether to retry.
+      #
+      # @return [Proc, nil]
+      def retry_predicate
+        @retry_predicate
+      end
+
       ##
       # Returns a duplicate in a non-executing state, i.e. with the deadline
       # and current delay reset.
@@ -101,7 +123,8 @@ def dup
                         multiplier: @multiplier,
                         retry_codes: @retry_codes,
                         timeout: @timeout,
-                        jitter: @jitter
+                        jitter: @jitter,
+                        retry_predicate: @retry_predicate
       end
 
       ##
@@ -176,6 +199,13 @@ def start! mock_delay: false
       # @return [Boolean] Whether this error should be retried.
       #
       def retry_error? error
+        # Run predicate first. If it returns nil, decision is made on the basis of retry_codes.
+        # Otherwise return truthiness.
+        if @retry_predicate.respond_to? :call
+          result = @retry_predicate.call error
+          return !!result unless result.nil?
+        end
+
         (defined?(::GRPC) && error.is_a?(::GRPC::BadStatus) && retry_codes.include?(error.code)) ||
           (error.respond_to?(:response_status) &&
             retry_codes.include?(ErrorCodes.grpc_error_for(error.response_status)))
@@ -197,11 +227,12 @@ def retry_with_deadline?
       #
       def apply_defaults retry_policy
         return unless retry_policy.is_a? Hash
-        @retry_codes   ||= convert_codes retry_policy[:retry_codes]
-        @initial_delay ||= retry_policy[:initial_delay]
-        @multiplier    ||= retry_policy[:multiplier]
-        @max_delay     ||= retry_policy[:max_delay]
-        @jitter        ||= retry_policy[:jitter]
+        @retry_codes     ||= convert_codes retry_policy[:retry_codes]
+        @initial_delay   ||= retry_policy[:initial_delay]
+        @multiplier      ||= retry_policy[:multiplier]
+        @max_delay       ||= retry_policy[:max_delay]
+        @jitter          ||= retry_policy[:jitter]
+        @retry_predicate ||= retry_policy[:retry_predicate]
 
         self
       end
@@ -214,13 +245,14 @@ def eql? other
           other.multiplier == multiplier &&
           other.retry_codes == retry_codes &&
           other.timeout == timeout &&
-          other.jitter == jitter
+          other.jitter == jitter &&
+          other.retry_predicate == retry_predicate
       end
       alias == eql?
 
       # @private Hash code
       def hash
-        [initial_delay, max_delay, multiplier, retry_codes, timeout, jitter].hash
+        [initial_delay, max_delay, multiplier, retry_codes, timeout, jitter, retry_predicate].hash
       end
 
       private

gapic-common/test/gapic/common/retry_policy_test.rb

@@ -79,6 +79,12 @@ def test_negative_jitter_raises_error
     end
   end
 
+  def test_uncallable_retry_predicate_raises_error
+    assert_raises ArgumentError do
+      Gapic::Common::RetryPolicy.new retry_predicate: "not_a_proc"
+    end
+  end
+
   def test_jitter_is_added
     retry_policy = Gapic::Common::RetryPolicy.new initial_delay: 5, max_delay: 10, multiplier: 2, jitter: 2.0
 
@@ -115,4 +121,42 @@ def test_jitter_bounds
     # internal state is now 10, rand is 2.0 = 12.0, capped at max_delay of 6
     assert_equal 6, delays_performed.last
   end
+
+  class CustomRetryableError < StandardError; end
+
+  def test_retry_predicate_returns_truthy
+    predicate = ->(error) { error.is_a? CustomRetryableError }
+    retry_policy = Gapic::Common::RetryPolicy.new retry_codes: [], retry_predicate: predicate, jitter: 0
+    retry_policy.start! mock_delay: true
+
+    assert retry_policy.call(CustomRetryableError.new("transient"))
+    refute retry_policy.call(StandardError.new("permanent failure"))
+  end
+
+  def test_retry_predicate_returns_false
+    predicate = ->(_error) { false }
+    retry_policy = Gapic::Common::RetryPolicy.new(
+      retry_codes: [GRPC::Core::StatusCodes::UNAVAILABLE],
+      retry_predicate: predicate,
+      jitter: 0
+    )
+    retry_policy.start! mock_delay: true
+
+    error = GRPC::BadStatus.new GRPC::Core::StatusCodes::UNAVAILABLE, "unavailable"
+    refute retry_policy.call(error)
+  end
+
+  def test_retry_predicate_returns_nil_fallback
+    predicate = ->(_error) { nil }
+    retry_policy = Gapic::Common::RetryPolicy.new(
+      retry_codes: [GRPC::Core::StatusCodes::UNAVAILABLE],
+      retry_predicate: predicate,
+      jitter: 0
+    )
+    retry_policy.start! mock_delay: true
+
+    retriable_error = GRPC::BadStatus.new GRPC::Core::StatusCodes::UNAVAILABLE, "unavailable"
+    assert retry_policy.call(retriable_error)
+  end
 end
+