#450: Use one rate-limit counter source

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

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

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

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

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')

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