retries.rb in httpx-0.6.6

- old
+ new

@@ -23,18 +23,35 @@
                           Errno::EINVAL,
                           Errno::ETIMEDOUT].freeze
 
       def self.extra_options(options)
         Class.new(options.class) do
+          # number of seconds after which one can retry the request
+          def_option(:retry_after) do |num|
+            # return early if callable
+            return num if num.respond_to?(:call)
+
+            num = Integer(num)
+            raise Error, ":retry_after must be positive" unless num.positive?
+
+            num
+          end
+
           def_option(:max_retries) do |num|
             num = Integer(num)
             raise Error, ":max_retries must be positive" unless num.positive?
 
             num
           end
 
           def_option(:retry_change_requests)
+
+          def_option(:retry_on) do |callback|
+            raise ":retry_on must be called with the response" unless callback.respond_to?(:call) && callback.method(:call).arity == 1
+
+            callback
+          end
         end.new(options).merge(max_retries: MAX_RETRIES)
       end
 
       module InstanceMethods
         def max_retries(n)
@@ -43,20 +60,23 @@
 
         private
 
         def fetch_response(request, connections, options)
           response = super
+
+          retry_on = options.retry_on
+
           if response.is_a?(ErrorResponse) &&
              request.retries.positive? &&
              __repeatable_request?(request, options) &&
-             __retryable_error?(response.error)
+             __retryable_error?(response.error) &&
+             (!retry_on || retry_on.call(response))
             request.retries -= 1
             log { "failed to get response, #{request.retries} tries to go..." }
             request.transition(:idle)
             connection = find_connection(request, connections, options)
-            connection.send(request)
-            set_request_timeout(connection, request, options)
+            __retry_request(connection, request, options)
             return
           end
           response
         end
 
@@ -64,9 +84,26 @@
           IDEMPOTENT_METHODS.include?(request.verb) || options.retry_change_requests
         end
 
         def __retryable_error?(ex)
           RETRYABLE_ERRORS.any? { |klass| ex.is_a?(klass) }
+        end
+
+        def __retry_request(connection, request, options)
+          retry_after = options.retry_after
+          unless retry_after
+            connection.send(request)
+            set_request_timeout(connection, request, options)
+            return
+          end
+
+          retry_after = retry_after.call(request) if retry_after.respond_to?(:call)
+          log { "retrying after #{retry_after} secs..." }
+          pool.after(retry_after) do
+            log { "retrying!!" }
+            connection.send(request)
+            set_request_timeout(connection, request, options)
+          end
         end
       end
 
       module RequestMethods
         attr_accessor :retries