retry requests when rate limit reached

.github/workflows/ci.yml

@@ -12,7 +12,7 @@ jobs:
       fail-fast: true
       matrix:
         os: [ ubuntu-latest, macos-latest ]
-        ruby: [ 2.5, 2.6, 2.7, "3.0", 3.1, 3.2 ] # "x.0" to workaround: https://github.com/ruby/setup-ruby/issues/252
+        ruby: [ 2.6, 2.7, "3.0", 3.1, 3.2 ] # "x.0" to workaround: https://github.com/ruby/setup-ruby/issues/252
     runs-on: ${{ matrix.os }}
     steps:
     - uses: actions/checkout@v3

README.md

@@ -52,6 +52,19 @@ require 'droplet_kit'
 client = DropletKit::Client.new(access_token: 'YOUR_TOKEN', user_agent: 'custom')
 ```
 
+### Automatically Retry Rate Limited Requests
+
+By default, DropletKit will handle requests that are rate limited by the DigitalOcean API's [burst limit](https://docs.digitalocean.com/reference/api/api-reference/#section/Introduction/Rate-Limit). When the burst rate limit is reached, DropletKit will wait according to the value of the API response's `Retry-After` header. Typically the wait time is less than one minute. When the hourly rate limit is hit, an error is raised.
+
+By default, DropletKit will retry a rate limited request three times before returning an error. If you would like to disable the retry behavior altogether, and instead raise an error when any rate limit is reached, you can set the `retry_max` config value to zero.
+
+DropletKit will also wait zero seconds until retrying a request after the `Retry-After` time has elapsed by default. To change this, set the `retry_wait_min` to a different value.
+
+```ruby
+require 'droplet_kit'
+client = DropletKit::Client.new(access_token: 'YOUR_TOKEN', retry_max: 3, retry_wait_min: 1)
+```
+
 ## Design
 
 DropletKit follows a strict design of resources as methods on your client. For examples, for droplets, you will call your client like this:

droplet_kit.gemspec

@@ -20,6 +20,7 @@ Gem::Specification.new do |spec|
   spec.required_ruby_version = '>= 2.5.0'
 
   spec.add_dependency 'faraday', '>= 0.15'
+  spec.add_dependency 'faraday-retry', '~> 2.2.0'
   spec.add_dependency 'kartograph', '~> 0.2.8'
   spec.add_dependency 'resource_kit', '~> 0.1.5'
   spec.add_dependency 'virtus', '>= 1.0.3', '<= 3'

lib/droplet_kit/client.rb

@@ -1,30 +1,49 @@
 # frozen_string_literal: true
 
 require 'faraday'
+require 'faraday/retry'
 require 'droplet_kit/utils'
 
 module DropletKit
   class Client
     DEFAULT_OPEN_TIMEOUT = 60
     DEFAULT_TIMEOUT = 120
+    DEFAULT_RETRY_MAX = 3
+    DEFAULT_RETRY_WAIT_MIN = 0
+
     DIGITALOCEAN_API = 'https://api.digitalocean.com'
 
-    attr_reader :access_token, :api_url, :open_timeout, :timeout, :user_agent
+    attr_reader :access_token, :api_url, :open_timeout, :timeout, :user_agent, :retry_max, :retry_wait_min
 
     def initialize(options = {})
       options = DropletKit::Utils.transform_keys(options, &:to_sym)
-      @access_token = options[:access_token]
-      @api_url      = options[:api_url] || DIGITALOCEAN_API
-      @open_timeout = options[:open_timeout] || DEFAULT_OPEN_TIMEOUT
-      @timeout      = options[:timeout] || DEFAULT_TIMEOUT
-      @user_agent   = options[:user_agent]
+      @access_token   = options[:access_token]
+      @api_url        = options[:api_url] || DIGITALOCEAN_API
+      @open_timeout   = options[:open_timeout] || DEFAULT_OPEN_TIMEOUT
+      @timeout        = options[:timeout] || DEFAULT_TIMEOUT
+      @user_agent     = options[:user_agent]
+      @retry_max      = options[:retry_max] || DEFAULT_RETRY_MAX
+      @retry_wait_min = options[:retry_wait_min] || DEFAULT_RETRY_WAIT_MIN
     end
 
     def connection
       @faraday ||= Faraday.new connection_options do |req|
         req.adapter :net_http
         req.options.open_timeout = open_timeout
         req.options.timeout = timeout
+        unless retry_max.zero?
+          req.request :retry, {
+            max: @retry_max,
+            interval: @retry_wait_min,
+            retry_statuses: [429],
+            # faraday-retry supports both the Retry-After and RateLimit-Reset
+            # headers, however, it favours the RateLimit-Reset one. To force it
+            # to use the Retry-After header, we override the header that it
+            # expects for the RateLimit-Reset header to something that we know
+            # we don't set.
+            rate_limit_reset_header: 'undefined'
+          }
+        end
       end
     end
 
@@ -91,6 +110,11 @@ def connection_options
           content_type: 'application/json',
           authorization: "Bearer #{access_token}",
           user_agent: "#{user_agent} #{default_user_agent}".strip
+        },
+        request: {
+          context: {
+            retry_max: @retry_max
+          }
         }
       }
     end

lib/droplet_kit/error_handling_resourcable.rb

@@ -8,11 +8,13 @@ def self.included(base)
         when 200...299
           next
         when 429
-          error = DropletKit::RateLimitReached.new("#{response.status}: #{response.body}")
-          error.limit = response.headers['RateLimit-Limit']
-          error.remaining = response.headers['RateLimit-Remaining']
-          error.reset_at = response.headers['RateLimit-Reset']
-          raise error
+          unless response.headers.key?('Retry-After') && !connection.options.context.key?(:retry_max)
+            error = DropletKit::RateLimitReached.new("#{response.status}: #{response.body}")
+            error.limit = response.headers['RateLimit-Limit']
+            error.remaining = response.headers['RateLimit-Remaining']
+            error.reset_at = response.headers['RateLimit-Reset']
+            raise error
+          end
         else
           raise DropletKit::Error, "#{response.status}: #{response.body}"
         end

spec/lib/droplet_kit/client_spec.rb

@@ -51,6 +51,39 @@
 
       expect(client.timeout).to eq(timeout)
     end
+
+    it 'allows retry max to be set' do
+      retry_max = 10
+      client = described_class.new(
+        'access_token' => 'my-token',
+        'retry_max' => retry_max
+      )
+
+      expect(client.retry_max).to eq(retry_max)
+    end
+
+    it 'allows retry wait min to be set' do
+      retry_wait_min = 3
+      client = described_class.new(
+        'access_token' => 'my-token',
+        'retry_wait_min' => retry_wait_min
+      )
+
+      expect(client.retry_wait_min).to eq(retry_wait_min)
+    end
+
+    it 'does not handle rate limited requests when retry max is zero' do
+      client = described_class.new(retry_max: 0)
+
+      stub_do_api('/v2/account', :get).to_return(body: { id: :rate_limit, message: '429' }.to_json, status: 429, headers: {
+                                                   'RateLimit-Limit' => 1200,
+                                                   'RateLimit-Remaining' => 1193,
+                                                   'RateLimit-Reset' => 1_402_425_459,
+                                                   'Retry-After' => 0
+                                                 })
+
+      expect { client.account.send(:info).to_a }.to raise_exception(DropletKit::RateLimitReached)
+    end
   end
 
   describe '#method_missing' do

spec/lib/droplet_kit/resources/account_resource_spec.rb

@@ -31,5 +31,11 @@
       let(:method) { :get }
       let(:action) { :info }
     end
+
+    it_behaves_like 'resource that handles rate limit retries' do
+      let(:path) { '/v2/account' }
+      let(:method) { :get }
+      let(:action) { :info }
+    end
   end
 end

spec/lib/droplet_kit/resources/droplet_resource_spec.rb

@@ -133,6 +133,13 @@ def check_droplet(droplet, tags = [], overrides = {})
       let(:action) { :find }
       let(:arguments) { { id: 123 } }
     end
+
+    it_behaves_like 'resource that handles rate limit retries' do
+      let(:path) { '/v2/droplets/123' }
+      let(:method) { :get }
+      let(:action) { :find }
+      let(:arguments) { { id: 123 } }
+    end
   end
 
   describe '#create' do

spec/spec_helper.rb

@@ -5,6 +5,7 @@
 SimpleCov.start
 
 require 'faraday'
+require 'faraday/retry'
 require 'addressable/uri'
 require 'droplet_kit'
 require 'webmock/rspec'

spec/support/shared_examples/common_errors.rb

@@ -3,7 +3,7 @@
 shared_examples_for 'resource that handles common errors' do
   let(:arguments) { {} }
 
-  it 'handles rate limit' do
+  it 'handles rate limit when retry-after is not present' do
     response_body = { id: :rate_limit, message: 'Too much!!!' }
     stub_do_api(path, method).to_return(body: response_body.to_json, status: 429, headers: {
                                           'RateLimit-Limit' => 1200,

spec/support/shared_examples/rate_limit_retry.rb

@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+shared_examples_for 'resource that handles rate limit retries' do
+  let(:arguments) { {} }
+
+  it 'handles rate limit' do
+    response_body = { id: :rate_limit, message: 'example' }
+    stub_do_api(path, method).to_return(
+      [
+        {
+          body: nil,
+          status: 429,
+          headers: {
+            'RateLimit-Limit' => 1200,
+            'RateLimit-Remaining' => 1193,
+            'RateLimit-Reset' => 1_402_425_459,
+            'Retry-After' => 0 # Retry immediately in tests.
+          }
+        },
+        {
+          body: response_body.to_json,
+          status: 200,
+          headers: {
+            'RateLimit-Limit' => 1200,
+            'RateLimit-Remaining' => 1192,
+            'RateLimit-Reset' => 1_402_425_459
+          }
+        }
+      ]
+    )
+
+    expect { resource.send(action, arguments) }.not_to raise_error
+  end
+end