diff --git a/lib/buildkit/client.rb b/lib/buildkit/client.rb index a01dc2d..f532c4a 100644 --- a/lib/buildkit/client.rb +++ b/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,6 +124,7 @@ 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) @@ -126,8 +132,20 @@ def request(method, path, data, options = {}) 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) diff --git a/lib/buildkit/error.rb b/lib/buildkit/error.rb index 45bb284..858aa2c 100644 --- a/lib/buildkit/error.rb +++ b/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 diff --git a/spec/client/rate_limit_spec.rb b/spec/client/rate_limit_spec.rb new file mode 100644 index 0000000..ad93463 --- /dev/null +++ b/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