diff --git a/lib/fbe/middleware/rate_limit.rb b/lib/fbe/middleware/rate_limit.rb index 8474975..2088be1 100644 --- a/lib/fbe/middleware/rate_limit.rb +++ b/lib/fbe/middleware/rate_limit.rb @@ -27,13 +27,14 @@ class Fbe::Middleware::RateLimit < Faraday::Middleware # Initializes the rate limit middleware. # # @param [Object] app The next middleware in the stack - def initialize(app) - super + def initialize(app, tracker = nil) + super(app) @cached = nil - @remaining = 0 - @searchleft = 0 + @remaining = nil + @searchleft = nil @counter = 0 @lock = Mutex.new + tracker[:rate_limit] = self unless tracker.nil? end # Processes the HTTP request and handles rate limit caching. @@ -51,6 +52,16 @@ def call(env) end end + # Returns the remaining requests count tracked by this middleware. + # + # @param [Symbol] resource The GitHub API resource (:core or :search) + # @return [Integer, nil] The remaining count, or nil when the resource is absent + def remaining(resource = :core) + @lock.synchronize do + resource == :search ? @searchleft : @remaining + end + end + private # Handles requests to the rate_limit endpoint. @@ -74,8 +85,8 @@ def handle_rate_limit_request(env) def track_request(path = nil) @counter += 1 if path&.start_with?('/search/') - @searchleft -= 1 if @searchleft.positive? - elsif @remaining.positive? + @searchleft -= 1 if @searchleft&.positive? + elsif @remaining&.positive? @remaining -= 1 end end @@ -86,7 +97,8 @@ def track_request(path = nil) # (indicated by +http_cache_trace+ containing +:fresh+), the # +x-ratelimit-remaining+ header is stale, so we keep our # decremented count. When the API was actually contacted, - # we prefer the authoritative header value. + # we seed unknown counters from headers, but avoid raising + # a counter already decremented by this middleware. # # @param [Faraday::Env] response_env The response environment def sync(response_env, path = nil) @@ -95,10 +107,11 @@ def sync(response_env, path = nil) return unless headers remaining = headers['x-ratelimit-remaining'] return unless remaining + count = Integer(remaining) if path&.start_with?('/search/') - @searchleft = Integer(remaining) + @searchleft = @searchleft.nil? ? count : [@searchleft, count].min else - @remaining = Integer(remaining) + @remaining = @remaining.nil? ? count : [@remaining, count].min end end @@ -109,19 +122,21 @@ def sync(response_env, path = nil) def extract_remaining_count(response) body = response.body body = JSON.parse(body) if body.is_a?(String) - return 0 unless body.is_a?(Hash) - body.dig('rate', 'remaining') || 0 + value = body.dig('rate', 'remaining') if body.is_a?(Hash) + value ||= response.headers['x-ratelimit-remaining'] + Integer(value || 0) end # Extracts the search-resource remaining count from the response body. # # @param [Faraday::Response] response The API response - # @return [Integer] The remaining search-API requests count + # @return [Integer, nil] The remaining search-API requests count def extract_search_remaining_count(response) body = response.body body = JSON.parse(body) if body.is_a?(String) - return 0 unless body.is_a?(Hash) - body.dig('resources', 'search', 'remaining') || 0 + return nil unless body.is_a?(Hash) + value = body.dig('resources', 'search', 'remaining') + value.nil? ? nil : Integer(value) end # Builds a fresh body with the current remaining counts written in, @@ -140,9 +155,9 @@ def patched_body(original) else return original end - body['rate']['remaining'] = @remaining if body['rate'] - body.dig('resources', 'core')&.[]=('remaining', @remaining) - body.dig('resources', 'search')&.[]=('remaining', @searchleft) + body['rate']['remaining'] = @remaining if body['rate'] && !@remaining.nil? + body.dig('resources', 'core')&.[]=('remaining', @remaining) unless @remaining.nil? + body.dig('resources', 'search')&.[]=('remaining', @searchleft) unless @searchleft.nil? stringed ? body.to_json : body end @@ -158,7 +173,7 @@ def response_env(env, response) served.request_headers = env.request_headers served.body = patched_body(response.body) served.response_headers = response.headers.dup - served.response_headers['x-ratelimit-remaining'] = @remaining.to_s + served.response_headers['x-ratelimit-remaining'] = @remaining.to_s unless @remaining.nil? served end end diff --git a/lib/fbe/octo.rb b/lib/fbe/octo.rb index e9b9a6f..f2458f4 100644 --- a/lib/fbe/octo.rb +++ b/lib/fbe/octo.rb @@ -56,6 +56,7 @@ def Fbe.octo(options: $options, global: $global, loog: $loog) # rubocop:disable begin loog.info("Fbe version is #{Fbe::VERSION}") trace = [] + limits = {} if options.testing.nil? o = Octokit::Client.new token = options.github_token @@ -94,7 +95,7 @@ def Fbe.octo(options: $options, global: $global, loog: $loog) # rubocop:disable ) builder.use(Octokit::Response::RaiseError) builder.use(Faraday::Response::Logger, loog, formatter: Fbe::Middleware::Formatter) - builder.use(Fbe::Middleware::RateLimit) + builder.use(Fbe::Middleware::RateLimit, limits) builder.use(Fbe::Middleware::Trace, trace, ignores: [:fresh]) if options.sqlite_cache maxsize = Integer(Filesize.from(options.sqlite_cache_maxsize || '100M')) @@ -128,7 +129,7 @@ def Fbe.octo(options: $options, global: $global, loog: $loog) # rubocop:disable o = Fbe::FakeOctokit.new end o = - decoor(o, loog:, trace:) do # rubocop:disable Metrics/BlockLength + decoor(o, loog:, trace:, limits:) do # rubocop:disable Metrics/BlockLength def print_trace!(all: false, max: 5) if @trace.empty? @loog.debug('GitHub API trace is empty') @@ -160,27 +161,17 @@ def print_trace!(all: false, max: 5) @trace.clear end end - def off_quota?(threshold: nil, resource: :core) # rubocop:disable Layout/EmptyLineBetweenDefs, Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity + def off_quota?(threshold: nil, resource: :core) # rubocop:disable Layout/EmptyLineBetweenDefs threshold ||= resource == :search ? 5 : 50 label = resource == :search ? 'GitHub Search API' : 'GitHub API' - left = @origin.rate_limit!.remaining - got = false - if resource == :search && @origin.respond_to?(:last_response) - body = @origin.last_response&.body - body = JSON.parse(body) if body.is_a?(String) - if body.is_a?(Hash) - fresh = body.dig('resources', 'search', 'remaining') || body.dig(:resources, :search, :remaining) - if fresh - left = Integer(fresh) - got = true - end - end - end + rate = @origin.rate_limit! + left = @limits[:rate_limit]&.remaining(resource) + got = !left.nil? + left = rate.remaining unless got if resource == :search && !got - klass = @origin.respond_to?(:last_response) ? @origin.last_response&.body&.class : nil @loog.warn( "Search-quota check fell back to core remaining (#{left}); " \ - "search count unavailable (last_response body class: #{klass.inspect})" + 'search count unavailable in rate-limit middleware' ) end if left < threshold diff --git a/lib/fbe/over.rb b/lib/fbe/over.rb index a8e3171..80cdf77 100644 --- a/lib/fbe/over.rb +++ b/lib/fbe/over.rb @@ -24,15 +24,12 @@ def Fbe.over?( ) if quota_aware octo = Fbe.octo(loog:, options:, global:) - if octo.off_quota?(threshold: 100, resource: :core) || octo.off_quota?(resource: :search) + if octo.off_quota?(threshold: 100, resource: :core) || + octo.off_quota?(resource: :search, threshold: 10) loog.info('We are off GitHub quota, time to stop') return true end end - if quota_aware && Fbe.octo(loog:, options:, global:).off_quota?(resource: :search, threshold: 10) - loog.info('We are off GitHub Search API quota, time to stop') - return true - end if lifetime_aware && options.lifetime && Time.now - epoch > options.lifetime * 0.9 loog.info("We ran out of lifetime (#{epoch.ago} already), must stop here") return true diff --git a/test/fbe/middleware/test_rate_limit.rb b/test/fbe/middleware/test_rate_limit.rb index 8563ed7..659a796 100644 --- a/test/fbe/middleware/test_rate_limit.rb +++ b/test/fbe/middleware/test_rate_limit.rb @@ -57,6 +57,22 @@ def test_decrements_remaining_count_for_non_rate_limit_requests assert_equal('4998', response.headers['x-ratelimit-remaining']) end + def test_updates_remaining_count_from_non_rate_limit_response_header + tracker = {} + stub_request(:get, 'https://api.github.com/user') + .to_return( + status: 200, + body: '{"login": "test"}', + headers: { + 'Content-Type' => 'application/json', + 'X-RateLimit-Remaining' => '1' + } + ) + conn = create_connection(tracker) + conn.get('/user') + assert_equal(1, tracker[:rate_limit].remaining) + end + def test_refreshes_cache_after_hundred_requests payload = { 'rate' => { 'limit' => 5000, 'remaining' => 4999, 'reset' => 1_672_531_200 } } refreshed = { 'rate' => { 'limit' => 5000, 'remaining' => 4950, 'reset' => 1_672_531_200 } } @@ -354,9 +370,9 @@ def test_cached_body_is_not_leaked_to_callers private - def create_connection + def create_connection(tracker = nil) Faraday.new(url: 'https://api.github.com') do |f| - f.use(Fbe::Middleware::RateLimit) + f.use(Fbe::Middleware::RateLimit, tracker) f.response(:json) f.adapter(:net_http) end diff --git a/test/fbe/test_octo.rb b/test/fbe/test_octo.rb index b3b947d..5fe7fac 100644 --- a/test/fbe/test_octo.rb +++ b/test/fbe/test_octo.rb @@ -198,6 +198,34 @@ def test_off_quota_search_when_search_resource_exhausted ) end + def test_off_quota_search_uses_middleware_count_when_last_response_loses_search_resource + WebMock.disable_net_connect! + stub_request(:get, 'https://api.github.com/rate_limit').to_return( + body: { rate: { remaining: 4_999 }, resources: { core: { remaining: 4_999 }, search: { remaining: 5 } } }.to_json, + headers: { 'Content-Type' => 'application/json', 'X-RateLimit-Remaining' => '4999' } + ) + stub_request(:get, %r{https://api.github.com/search/issues}).to_return( + body: { total_count: 0, incomplete_results: false, items: [] }.to_json, + headers: { 'Content-Type' => 'application/json', 'X-RateLimit-Remaining' => '4999' } + ) + o = Fbe.octo(loog: Loog::NULL, global: {}, options: Judges::Options.new) + assert_equal(0, o.search_issues('repo:foo/bar type:issue')[:total_count]) + client = o.instance_variable_get(:@origin) + client.define_singleton_method(:rate_limit!) do + result = @origin.rate_limit! + @lost = Struct.new(:body).new({ 'rate' => { 'remaining' => result.remaining } }) + result + end + client.define_singleton_method(:last_response) do + @lost || @origin.last_response + end + assert( + o.off_quota?(resource: :search), + 'search quota check must use middleware remaining after a search request ' \ + 'even when last_response loses resources.search' + ) + end + def test_search_issues_blocked_when_search_quota_exhausted WebMock.disable_net_connect! stub_request(:get, 'https://api.github.com/rate_limit').to_return( @@ -731,7 +759,7 @@ def test_print_trace octo.print_trace!(all: true, max: 9_999) output = loog.to_s assert_includes(output, '3 URLs vs 4 requests') - assert_includes(output, '222 quota left') + assert_includes(output, '219 quota left') assert_includes(output, '/rate_limit: 1') assert_includes(output, '/user/123: 1') assert_includes(output, '/repos/foo/bar: 2') diff --git a/test/fbe/test_over.rb b/test/fbe/test_over.rb index 352013e..4f16a30 100644 --- a/test/fbe/test_over.rb +++ b/test/fbe/test_over.rb @@ -38,7 +38,7 @@ def test_check_search_off_quota_enabled end assert(Fbe.over?(global:, options:, loog:, quota_aware: true)) assert_includes(calls, { threshold: 100, resource: :core }) - assert_includes(calls, { threshold: nil, resource: :search }) + assert_includes(calls, { threshold: 10, resource: :search }) end def test_check_off_quota_disabled