Adds optional auto-retry on Buildkite API rate limit (429)

lib/buildkit/client.rb

@@ -42,11 +42,16 @@ def build_middleware
 
     def initialize(endpoint: ENV.fetch('BUILDKITE_API_ENDPOINT', DEFAULT_ENDPOINT),
                    token: ENV.fetch('BUILDKITE_API_TOKEN'),
-                   middleware: self.class.build_middleware, auto_paginate: false)
+                   middleware: self.class.build_middleware,
+                   auto_paginate: false,
+                   auto_retry_rate_limit: false,
+                   rate_limit_retry_count: 3)
       @middleware = middleware
       @endpoint = endpoint
       @token = token
       @auto_paginate = auto_paginate
+      @auto_retry_rate_limit = auto_retry_rate_limit
+      @rate_limit_retry_count = rate_limit_retry_count
     end
 
     # Make a HTTP GET request
@@ -119,15 +124,28 @@ def root
     private
 
     def request(method, path, data, options = {})
+      attempts = 0
       if data.is_a?(Hash)
         options = extract_query_and_headers_from data
         if accept = data.delete(:accept)
           options[:headers][:accept] = accept
         end
       end
 
-      @last_response = response = sawyer_agent.call(method, URI::DEFAULT_PARSER.escape(path.to_s), data, options)
-      response.data
+      begin
+        @last_response = response = sawyer_agent.call(method, URI::DEFAULT_PARSER.escape(path.to_s), data, options)
+        return response.data
+      rescue Buildkit::RateLimitExceeded => e
+        raise unless @auto_retry_rate_limit
+        attempts += 1
+        raise if attempts > @rate_limit_retry_count
+
+        headers = e.instance_variable_get('@response')[:response_headers]
+        wait_seconds = headers && headers[:rate_limit_reset]&.to_i
+        wait_seconds = 60 if wait_seconds.nil? || wait_seconds.zero?
+        sleep wait_seconds
+        retry
+      end
     end
 
     def extract_query_and_headers_from(data)

lib/buildkit/error.rb

@@ -21,6 +21,7 @@ def self.from_response(response)
                   when 415      then Buildkit::UnsupportedMediaType
                   when 422      then Buildkit::UnprocessableEntity
                   when 400..499 then Buildkit::ClientError
+                  when 429      then Buildkit::RateLimitExceeded
                   when 500      then Buildkit::InternalServerError
                   when 501      then Buildkit::NotImplemented
                   when 502      then Buildkit::BadGateway
@@ -146,6 +147,9 @@ class Conflict < ClientError; end
   # Raised when Buildkite returns a 414 HTTP status code
   class UnsupportedMediaType < ClientError; end
 
+  # Raised when Buildkite returns a 429 HTTP status code (rate limit exceeded)
+  class RateLimitExceeded < ClientError; end
+
   # Raised when Buildkite returns a 422 HTTP status code
   class UnprocessableEntity < ClientError; end
 

spec/client/rate_limit_spec.rb

@@ -0,0 +1,60 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+require 'ostruct'
+
+RSpec.describe 'Rate limit handling' do
+  def stub_response(http_code, body = '', headers = {})
+    {
+      method: 'GET',
+      url: 'https://example.com',
+      status: http_code,
+      body: body,
+      response_headers: headers,
+    }
+  end
+
+  context 'error mapping' do
+    it 'maps 429 to RateLimitExceeded' do
+      response = stub_response(429, 'Rate limit exceeded')
+      error = Buildkit::Error.from_response(response)
+      expect(error).to be_kind_of(Buildkit::RateLimitExceeded)
+    end
+  end
+
+  context 'client auto retry' do
+    it 'retries and succeeds when enabled' do
+      client = Buildkit::Client.new(token: 'abc', auto_retry_rate_limit: true, rate_limit_retry_count: 2)
+
+      # Build fake Sawyer agent
+      agent = double('Sawyer::Agent')
+      call_count = 0
+      success_response = OpenStruct.new(data: 'ok')
+      # First call raises rate limit, second succeeds
+      allow(agent).to receive(:call) do |_method, _path, _data, _options|
+        if (call_count += 1) == 1
+          raise Buildkit::RateLimitExceeded.new(stub_response(429, '', { rate_limit_reset: '0' }))
+        else
+          success_response
+        end
+      end
+      allow(client).to receive(:sawyer_agent).and_return(agent)
+      allow(client).to receive(:sleep) # prevent actual wait
+
+      result = client.send(:request, :get, '/path', {})
+      expect(result).to eq('ok')
+      expect(call_count).to eq(2)
+    end
+
+    it 'does not retry when disabled' do
+      client = Buildkit::Client.new(token: 'abc', auto_retry_rate_limit: false)
+      agent = double('Sawyer::Agent')
+      allow(agent).to receive(:call).and_raise(Buildkit::RateLimitExceeded.new(stub_response(429)))
+      allow(client).to receive(:sawyer_agent).and_return(agent)
+
+      expect do
+        client.send(:request, :get, '/path', {})
+      end.to raise_error(Buildkit::RateLimitExceeded)
+    end
+  end
+end