From 87ed882234ea4639d07614eee202f432631af715 Mon Sep 17 00:00:00 2001 From: evazion Date: Sat, 20 Jun 2020 22:38:58 -0500 Subject: [PATCH] danbooru::http: support automatically retrying 429 errors. --- app/logical/danbooru/http.rb | 14 ++--- .../danbooru/http/application_client.rb | 31 +++++++++++ app/logical/danbooru/http/retriable.rb | 52 +++++++++++++++++++ test/unit/danbooru_http_test.rb | 29 +++++++++++ 4 files changed, 120 insertions(+), 6 deletions(-) create mode 100644 app/logical/danbooru/http/application_client.rb create mode 100644 app/logical/danbooru/http/retriable.rb diff --git a/app/logical/danbooru/http.rb b/app/logical/danbooru/http.rb index d6bad229c..c57f85f0b 100644 --- a/app/logical/danbooru/http.rb +++ b/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. - follow(strict: false, max_hops: MAX_REDIRECTS). - timeout(DEFAULT_TIMEOUT). - use(:auto_inflate). - headers(Danbooru.config.http_headers). - headers("Accept-Encoding" => "gzip") + @http ||= + ::Danbooru::Http::ApplicationClient.new + .follow(strict: false, max_hops: MAX_REDIRECTS) + .timeout(DEFAULT_TIMEOUT) + .use(:auto_inflate) + .headers(Danbooru.config.http_headers) + .headers("Accept-Encoding" => "gzip") end end end diff --git a/app/logical/danbooru/http/application_client.rb b/app/logical/danbooru/http/application_client.rb new file mode 100644 index 000000000..27ac6a6d1 --- /dev/null +++ b/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 diff --git a/app/logical/danbooru/http/retriable.rb b/app/logical/danbooru/http/retriable.rb new file mode 100644 index 000000000..681dd8bbe --- /dev/null +++ b/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 diff --git a/test/unit/danbooru_http_test.rb b/test/unit/danbooru_http_test.rb index 50fc67d32..ae347bfd8 100644 --- a/test/unit/danbooru_http_test.rb +++ b/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")