danbooru::http: support automatically retrying 429 errors.

app/logical/danbooru/http.rb

@@ -1,5 +1,6 @@
 require "danbooru/http/html_adapter"
 require "danbooru/http/xml_adapter"
+require "danbooru/http/retriable"
 
 module Danbooru
   class Http
@@ -139,12 +140,13 @@ module Danbooru
     end
 
     def http
-      @http ||= ::HTTP.
+      @http ||=
-        follow(strict: false, max_hops: MAX_REDIRECTS).
+        ::Danbooru::Http::ApplicationClient.new
-        timeout(DEFAULT_TIMEOUT).
+        .follow(strict: false, max_hops: MAX_REDIRECTS)
-        use(:auto_inflate).
+        .timeout(DEFAULT_TIMEOUT)
-        headers(Danbooru.config.http_headers).
+        .use(:auto_inflate)
-        headers("Accept-Encoding" => "gzip")
+        .headers(Danbooru.config.http_headers)
+        .headers("Accept-Encoding" => "gzip")
     end
   end
 end

app/logical/danbooru/http/application_client.rb

@@ -0,0 +1,31 @@
+# An extension to HTTP::Client that lets us write Rack-style middlewares that
+# hook into the request/response cycle and override how requests are made. This
+# works by extending http.rb's concept of features (HTTP::Feature) to give them
+# a `perform` method that takes a http request and returns a http response.
+# This can be used to intercept and modify requests and return arbitrary responses.
+
+module Danbooru
+  class Http
+    class ApplicationClient < HTTP::Client
+      # Override `perform` to call the `perform` method on features first.
+      def perform(request, options)
+        features = options.features.values.reverse.select do |feature|
+          feature.respond_to?(:perform)
+        end
+
+        perform = proc { |req| super(req, options) }
+        callback_chain = features.reduce(perform) do |callback_chain, feature|
+          proc { |req| feature.perform(req, &callback_chain) }
+        end
+
+        callback_chain.call(request)
+      end
+
+      # Override `branch` to return an ApplicationClient instead of a
+      # HTTP::Client so that chaining works.
+      def branch(...)
+        ApplicationClient.new(...)
+      end
+    end
+  end
+end

app/logical/danbooru/http/retriable.rb

@@ -0,0 +1,52 @@
+# A HTTP::Feature that automatically retries requests that return a 429 error
+# or a Retry-After header. Usage: `Danbooru::Http.use(:retriable).get(url)`.
+#
+# https://developer.mozilla.org/en-US/docs/Web/HTTP/Status/429
+# https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Retry-After
+
+module Danbooru
+  class Http
+    class Retriable < HTTP::Feature
+      HTTP::Options.register_feature :retriable, self
+
+      attr_reader :max_retries, :max_delay
+
+      def initialize(max_retries: 2, max_delay: 5.seconds)
+        @max_retries = max_retries
+        @max_delay = max_delay
+      end
+
+      def perform(request, &block)
+        response = yield request
+
+        retries = max_retries
+        while retriable?(response) && retries > 0 && retry_delay(response) <= max_delay
+          retries -= 1
+          sleep(retry_delay(response))
+          response = yield request
+        end
+
+        response
+      end
+
+      def retriable?(response)
+        response.status == 429 || response.headers["Retry-After"].present?
+      end
+
+      def retry_delay(response, current_time: Time.zone.now)
+        retry_after = response.headers["Retry-After"]
+
+        if retry_after.blank?
+          0.seconds
+        elsif retry_after =~ /\A\d+\z/
+          retry_after.to_i.seconds
+        else
+          retry_at = Time.zone.parse(retry_after)
+          return 0.seconds if retry_at.blank?
+
+          [retry_at - current_time, 0].max.seconds
+        end
+      end
+    end
+  end
+end

test/unit/danbooru_http_test.rb

@@ -61,6 +61,35 @@ class DanbooruHttpTest < ActiveSupport::TestCase
       end
     end
 
+    context "retriable feature" do
+      should "retry immediately if no Retry-After header is sent" do
+        response_429 = ::HTTP::Response.new(status: 429, version: "1.1", body: "")
+        response_200 = ::HTTP::Response.new(status: 200, version: "1.1", body: "")
+        HTTP::Client.any_instance.expects(:perform).times(2).returns(response_429, response_200)
+
+        response = Danbooru::Http.use(:retriable).get("https://httpbin.org/status/429")
+        assert_equal(200, response.status)
+      end
+
+      should "retry if the Retry-After header is an integer" do
+        response_503 = ::HTTP::Response.new(status: 503, version: "1.1", headers: { "Retry-After": "1" }, body: "")
+        response_200 = ::HTTP::Response.new(status: 200, version: "1.1", body: "")
+        HTTP::Client.any_instance.expects(:perform).times(2).returns(response_503, response_200)
+
+        response = Danbooru::Http.use(:retriable).get("https://httpbin.org/status/503")
+        assert_equal(200, response.status)
+      end
+
+      should "retry if the Retry-After header is a date" do
+        response_503 = ::HTTP::Response.new(status: 503, version: "1.1", headers: { "Retry-After": 2.seconds.from_now.httpdate }, body: "")
+        response_200 = ::HTTP::Response.new(status: 200, version: "1.1", body: "")
+        HTTP::Client.any_instance.expects(:perform).times(2).returns(response_503, response_200)
+
+        response = Danbooru::Http.use(:retriable).get("https://httpbin.org/status/503")
+        assert_equal(200, response.status)
+      end
+    end
+
     context "#download method" do
       should "download files" do
         response, file = Danbooru::Http.download_media("https://httpbin.org/bytes/1000")