Merge branch 'issue-50' into 'master'

rate limiter plugin

See merge request honeyryderchuck/httpx!99

lib/httpx.rb

@@ -5,6 +5,7 @@ require "httpx/version"
 require "httpx/extensions"
 
 require "httpx/errors"
+require "httpx/utils"
 require "httpx/altsvc"
 require "httpx/callbacks"
 require "httpx/loggable"

lib/httpx/plugins/follow_redirects.rb

@@ -59,8 +59,26 @@ module HTTPX
             return ErrorResponse.new(request, error, options)
           end
 
-          connection = find_connection(retry_request, connections, options)
-          connection.send(retry_request)
+          retry_after = response.headers["retry-after"]
+
+          if retry_after
+            # Servers send the "Retry-After" header field to indicate how long the
+            # user agent ought to wait before making a follow-up request.
+            # When sent with any 3xx (Redirection) response, Retry-After indicates
+            # the minimum time that the user agent is asked to wait before issuing
+            # the redirected request.
+            #
+            retry_after = Utils.parse_retry_after(retry_after)
+
+            log { "redirecting after #{retry_after} secs..." }
+            pool.after(retry_after) do
+              connection = find_connection(retry_request, connections, options)
+              connection.send(retry_request)
+            end
+          else
+            connection = find_connection(retry_request, connections, options)
+            connection.send(retry_request)
+          end
           nil
         end
 

lib/httpx/plugins/persistent.rb

@@ -19,7 +19,12 @@ module HTTPX
     #
     module Persistent
       def self.load_dependencies(klass)
-        klass.plugin(:retries, max_retries: 1, retry_change_requests: true)
+        max_retries = if klass.default_options.respond_to?(:max_retries)
+          [klass.default_options.max_retries, 1].max
+        else
+          1
+        end
+        klass.plugin(:retries, max_retries: max_retries, retry_change_requests: true)
       end
 
       def self.extra_options(options)

lib/httpx/plugins/rate_limiter.rb

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+module HTTPX
+  module Plugins
+    #
+    # This plugin adds support for retrying requests when the request:
+    #
+    # * is rate limited;
+    # * when the server is unavailable (503);
+    # * when a 3xx request comes with a "retry-after" value
+    #
+    # https://gitlab.com/honeyryderchuck/httpx/wikis/RateLimiter
+    #
+    module RateLimiter
+      class << self
+        RATE_LIMIT_CODES = [429, 503].freeze
+
+        def load_dependencies(klass)
+          klass.plugin(:retries,
+                       retry_change_requests: true,
+                       retry_on: method(:retry_on_rate_limited_response),
+                       retry_after: method(:retry_after_rate_limit))
+        end
+
+        def retry_on_rate_limited_response(response)
+          status = response.status
+
+          RATE_LIMIT_CODES.include?(status)
+        end
+
+        # Servers send the "Retry-After" header field to indicate how long the
+        # user agent ought to wait before making a follow-up request.  When
+        # sent with a 503 (Service Unavailable) response, Retry-After indicates
+        # how long the service is expected to be unavailable to the client.
+        # When sent with any 3xx (Redirection) response, Retry-After indicates
+        # the minimum time that the user agent is asked to wait before issuing
+        # the redirected request.
+        #
+        def retry_after_rate_limit(_, response)
+          retry_after = response.headers["retry-after"]
+
+          return unless retry_after
+
+          Utils.parse_retry_after(retry_after)
+        end
+      end
+    end
+
+    register_plugin :rate_limiter, RateLimiter
+  end
+end

lib/httpx/plugins/retries.rb

@@ -81,8 +81,9 @@ module HTTPX
             request.transition(:idle)
 
             retry_after = options.retry_after
+            retry_after = retry_after.call(request, response) if retry_after.respond_to?(:call)
+
             if retry_after
-              retry_after = retry_after.call(request) if retry_after.respond_to?(:call)
 
               log { "retrying after #{retry_after} secs..." }
               pool.after(retry_after) do

lib/httpx/pool.rb

@@ -50,7 +50,7 @@ module HTTPX
       @timers.cancel
       connections = connections.reject(&:inflight?)
       connections.each(&:close)
-      next_tick until connections.none? { |c| @connections.include?(c) }
+      next_tick until connections.none? { |c| c.state != :idle && @connections.include?(c) }
       @resolvers.each_value do |resolver|
         resolver.close unless resolver.closed?
       end if @connections.empty?

lib/httpx/session.rb

@@ -221,7 +221,7 @@ module HTTPX
       def plugin(pl, options = nil, &block)
         # raise Error, "Cannot add a plugin to a frozen config" if frozen?
         pl = Plugins.load_plugin(pl) if pl.is_a?(Symbol)
-        unless @plugins.include?(pl)
+        if !@plugins.include?(pl)
           @plugins << pl
           pl.load_dependencies(self, &block) if pl.respond_to?(:load_dependencies)
           @default_options = @default_options.dup
@@ -245,6 +245,13 @@ module HTTPX
           opts.connection_class.__send__(:include, pl::ConnectionMethods) if defined?(pl::ConnectionMethods)
           pl.configure(self, &block) if pl.respond_to?(:configure)
 
+          @default_options.freeze
+        elsif options
+          # this can happen when two plugins are loaded, an one of them calls the other under the hood,
+          # albeit changing some default.
+          @default_options = @default_options.dup
+          @default_options = @default_options.merge(options)
+
           @default_options.freeze
         end
         self

lib/httpx/utils.rb

@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+
+module HTTPX
+  module Utils
+    module_function
+
+    # The value of this field can be either an HTTP-date or a number of
+    # seconds to delay after the response is received.
+    def parse_retry_after(retry_after)
+      # first: bet on it being an integer
+      Integer(retry_after)
+    rescue ArgumentError
+      # Then it's a datetime
+      time = Time.httpdate(retry_after)
+      time - Time.now
+    end
+  end
+end

sig/plugins/rate_limiter.rbs

@@ -0,0 +1,11 @@
+module HTTPX
+  module Plugins
+    module RateLimiter
+      # def self.load_dependencies: (singleton(Session)) -> void
+
+      def self.retry_on_rate_limited_response: (_Response) -> bool
+
+      def self.retry_after_rate_limit: (untyped, _Response) -> Numeric?
+    end
+  end
+end

sig/plugins/retries.rbs

@@ -10,9 +10,9 @@ module HTTPX
       end
 
       interface _RetriesOptions
-        def retry_after: () -> Integer?
-        def retry_after=: (int) -> Integer
-        def with_retry_after: (int) -> instance
+        def retry_after: () -> Numeric?
+        def retry_after=: (Numeric) -> Numeric
+        def with_retry_after: (Numeric) -> instance
 
         def max_retries: () -> Integer?
         def max_retries=: (int) -> Integer

test/http_test.rb

@@ -25,6 +25,7 @@ class HTTPTest < Minitest::Test
   include Plugins::Retries
   include Plugins::Multipart
   include Plugins::Expect
+  include Plugins::RateLimiter
 
   def test_verbose_log
     log = StringIO.new

test/https_test.rb

@@ -25,6 +25,8 @@ class HTTPSTest < Minitest::Test
   include Plugins::Retries
   include Plugins::Multipart
   include Plugins::Expect
+  include Plugins::RateLimiter
+  include Plugins::Persistent unless RUBY_ENGINE == "jruby" || RUBY_VERSION < "2.3"
 
   def test_connection_coalescing
     coalesced_origin = "https://#{ENV["HTTPBIN_COALESCING_HOST"]}"

test/support/request_inspector.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+module RequestInspector
+  module InstanceMethods
+    attr_reader :calls, :total_responses
+
+    def initialize(*args)
+      super
+      # we're comparing against max-retries + 1, because the calls increment will happen
+      # also in the last call, where the request is not going to be retried.
+      @calls = -1
+      @total_responses = []
+    end
+
+    def reset
+      @calls = -1
+      @total_responses.clear
+    end
+
+    def fetch_response(*)
+      response = super
+      if response
+        @calls += 1
+        @total_responses << response
+      end
+      response
+    end
+  end
+end

test/support/requests/plugins/follow_redirects.rb

@@ -36,6 +36,19 @@ module Requests
         verify_status(response, 302)
       end
 
+      def test_plugin_follow_redirects_retry_after
+        session = HTTPX.plugin(SessionWithMockResponse[302, "retry-after" => "2"]).plugin(:follow_redirects)
+
+        before_time = Process.clock_gettime(Process::CLOCK_MONOTONIC, :second)
+        response = session.get(max_redirect_uri(1))
+        after_time = Process.clock_gettime(Process::CLOCK_MONOTONIC, :second)
+
+        verify_status(response, 200)
+
+        total_time = after_time - before_time
+        assert total_time >= 2, "request didn't take as expected to redirect (#{total_time} secs)"
+      end
+
       def test_plugin_follow_insecure_no_insecure_downgrade
         return unless origin.start_with?("https")
 

test/support/requests/plugins/persistent.rb

@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+module Requests
+  module Plugins
+    module Persistent
+      def test_persistent
+        uri = build_uri("/get")
+
+        non_persistent_session = HTTPX.plugin(SessionWithPool)
+        response = non_persistent_session.get(uri)
+        verify_status(response, 200)
+        assert non_persistent_session.pool.connections.empty?, "unexpected connections ()"
+
+        persistent_session = non_persistent_session.plugin(:persistent)
+        response = persistent_session.get(uri)
+        verify_status(response, 200)
+        assert persistent_session.pool.connections.size == 1, "unexpected connections ()"
+
+        persistent_session.close
+        assert persistent_session.pool.connections.empty?, "unexpected connections ()"
+      end
+
+      def test_persistent_options
+        retry_persistent_session = HTTPX.plugin(:persistent).plugin(:retries, max_retries: 4)
+        options = retry_persistent_session.send(:default_options)
+        assert options.max_retries == 4
+        assert options.retry_change_requests
+        assert options.persistent
+
+        persistent_retry_session = HTTPX.plugin(:retries, max_retries: 4).plugin(:persistent)
+        options = persistent_retry_session.send(:default_options)
+        assert options.max_retries == 4
+        assert options.retry_change_requests
+        assert options.persistent
+      end
+    end
+  end
+end

test/support/requests/plugins/rate_limiter.rb

@@ -0,0 +1,73 @@
+# frozen_string_literal: true
+
+module Requests
+  module Plugins
+    module RateLimiter
+      def test_plugin_rate_limiter_429
+        rate_limiter_session = HTTPX.plugin(RequestInspector)
+                                    .plugin(SessionWithMockResponse[429])
+                                    .plugin(:rate_limiter)
+
+        uri = build_uri("/get")
+
+        rate_limiter_session.get(uri)
+
+        verify_rated_responses(rate_limiter_session, 429)
+      end
+
+      def test_plugin_rate_limiter_503
+        rate_limiter_session = HTTPX.plugin(RequestInspector)
+                                    .plugin(SessionWithMockResponse[503])
+                                    .plugin(:rate_limiter)
+
+        uri = build_uri("/get")
+
+        rate_limiter_session.get(uri)
+
+        verify_rated_responses(rate_limiter_session, 503)
+      end
+
+      def test_plugin_rate_limiter_retry_after_integer
+        rate_limiter_session = HTTPX.plugin(RequestInspector)
+                                    .plugin(SessionWithMockResponse[429, "retry-after" => "2"])
+                                    .plugin(:rate_limiter)
+
+        uri = build_uri("/get")
+
+        before_time = Process.clock_gettime(Process::CLOCK_MONOTONIC, :second)
+        rate_limiter_session.get(uri)
+        after_time = Process.clock_gettime(Process::CLOCK_MONOTONIC, :second)
+
+        verify_rated_responses(rate_limiter_session, 429)
+
+        total_time = after_time - before_time
+        assert total_time >= 2, "request didn't take as expected to retry (#{total_time} secs)"
+      end
+
+      def test_plugin_rate_limiter_retry_after_date
+        rate_limiter_session = HTTPX.plugin(RequestInspector)
+                                    .plugin(SessionWithMockResponse[429, "retry-after" => (Time.now + 3).httpdate])
+                                    .plugin(:rate_limiter)
+
+        uri = build_uri("/get")
+
+        before_time = Process.clock_gettime(Process::CLOCK_MONOTONIC, :second)
+        rate_limiter_session.get(uri)
+        after_time = Process.clock_gettime(Process::CLOCK_MONOTONIC, :second)
+
+        verify_rated_responses(rate_limiter_session, 429)
+        total_time = after_time - before_time
+        assert total_time >= 2, "request didn't take as expected to retry (#{total_time} secs)"
+      end
+
+      private
+
+      def verify_rated_responses(session, rated_status)
+        assert session.total_responses.size == 2, "expected 2 responses(was #{session.total_responses.size})"
+        rated_response, response = session.total_responses
+        verify_status(rated_response, rated_status)
+        verify_status(response, 200)
+      end
+    end
+  end
+end

test/support/requests/plugins/retries.rb

@@ -82,7 +82,7 @@ module Requests
 
       def test_plugin_retries_retry_after_callable
         retries = 0
-        exponential = ->(_) { (retries += 1) * 2 }
+        exponential = ->(*) { (retries += 1) * 2 }
         before_time = Process.clock_gettime(Process::CLOCK_MONOTONIC, :second)
         retries_session = HTTPX
                           .plugin(RequestInspector)
@@ -102,28 +102,6 @@ module Requests
       def verify_error_response(response)
         assert response.is_a?(HTTPX::ErrorResponse), "expected an error response, instead got #{response.inspect}"
       end
-
-      module RequestInspector
-        module InstanceMethods
-          attr_reader :calls
-          def initialize(*args)
-            super
-            # we're comparing against max-retries + 1, because the calls increment will happen
-            # also in the last call, where the request is not going to be retried.
-            @calls = -1
-          end
-
-          def reset
-            @calls = -1
-          end
-
-          def fetch_response(*)
-            response = super
-            @calls += 1 if response
-            response
-          end
-        end
-      end
     end
   end
 end

test/support/session_with_mock_response.rb

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+module SessionWithMockResponse
+  def self.[](status, headers = {})
+    Thread.current[:httpx_mock_response_status] = status
+    Thread.current[:httpx_mock_response_headers] = headers
+    self
+  end
+
+  module ResponseMethods
+    attr_writer :status
+  end
+
+  module InstanceMethods
+    def initialize(*)
+      super
+      @mock_responses_counter = 1
+    end
+
+    def on_response(request, response)
+      return super unless response && @mock_responses_counter.positive?
+
+      response.close
+      @mock_responses_counter -= 1
+
+      response.status = Thread.current[:httpx_mock_response_status]
+      response.merge_headers(Thread.current[:httpx_mock_response_headers])
+      super(request, response)
+    end
+  end
+end

test/support/session_with_pool.rb

@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+module SessionWithPool
+  ConnectionPool = Class.new(HTTPX::Pool) do
+    attr_reader :connections
+    attr_reader :connection_count
+    attr_reader :ping_count
+
+    def initialize(*)
+      super
+      @connection_count = 0
+      @ping_count = 0
+    end
+
+    def init_connection(connection, _)
+      super
+      connection.on(:open) { @connection_count += 1 }
+      connection.on(:pong) { @ping_count += 1 }
+    end
+  end
+
+  module InstanceMethods
+    def pool
+      @pool ||= ConnectionPool.new
+    end
+  end
+
+  module ConnectionMethods
+    def set_parser_callbacks(parser)
+      super
+      parser.on(:pong) { emit(:pong) }
+    end
+  end
+end

test/test_helper.rb

@@ -27,39 +27,6 @@ Dir[File.join(".", "test", "support", "**", "*.rb")].sort.each { |f| require f }
 # This adds it manually.
 OpenSSL::SSL::SSLContext::DEFAULT_CERT_STORE.add_file(ENV["SSL_CERT_FILE"]) if RUBY_VERSION.start_with?("2.3") && ENV.key?("SSL_CERT_FILE")
 
-module SessionWithPool
-  ConnectionPool = Class.new(HTTPX::Pool) do
-    attr_reader :connections
-    attr_reader :connection_count
-    attr_reader :ping_count
-
-    def initialize(*)
-      super
-      @connection_count = 0
-      @ping_count = 0
-    end
-
-    def init_connection(connection, _)
-      super
-      connection.on(:open) { @connection_count += 1 }
-      connection.on(:pong) { @ping_count += 1 }
-    end
-  end
-
-  module InstanceMethods
-    def pool
-      @pool ||= ConnectionPool.new
-    end
-  end
-
-  module ConnectionMethods
-    def set_parser_callbacks(parser)
-      super
-      parser.on(:pong) { emit(:pong) }
-    end
-  end
-end
-
 # 9090 drops SYN packets for connect timeout tests, make sure there's a server binding there.
 CONNECT_TIMEOUT_PORT = ENV.fetch("CONNECT_TIMEOUT_PORT", 9090).to_i