danbooru::http: support automatically retrying 429 errors.

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