#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.
@@ -45,7 +46,19 @@ def call(env)
       @lock.synchronize { handle_rate_limit_request(env) }
     else
       @lock.synchronize { track_request(env.url.path) }
-      @app.call(env)
+      response = @app.call(env)
+      @lock.synchronize { update(response, env.url.path) }
+      response
+    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
 
@@ -72,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
@@ -85,19 +98,36 @@ def track_request(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
+
+  # Updates the tracked remaining count from a non-rate_limit response header.
+  #
+  # @param [Faraday::Response] response The API response
+  # @param [String] path The requested path
+  def update(response, path)
+    value = response.headers['x-ratelimit-remaining']
+    return if value.nil?
+    count = Integer(value)
+    if path&.start_with?('/search/')
+      @searchleft = @searchleft.nil? ? count : [@searchleft, count].min
+    else
+      @remaining = @remaining.nil? ? count : [@remaining, count].min
+    end
   end
 
   # Builds a fresh body with the current remaining counts written in,
@@ -116,8 +146,8 @@ def patched_body(original)
       else
         return original
       end
-    body['rate']['remaining'] = @remaining if body['rate']
-    body.dig('resources', 'search')&.[]=('remaining', @searchleft)
+    body['rate']['remaining'] = @remaining if body['rate'] && !@remaining.nil?
+    body.dig('resources', 'search')&.[]=('remaining', @searchleft) unless @searchleft.nil?
     stringed ? body.to_json : body
   end
 
@@ -133,7 +163,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

@@ -54,6 +54,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
@@ -92,7 +93,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'))
@@ -126,7 +127,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')
@@ -158,27 +159,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

test/fbe/middleware/test_rate_limit.rb

@@ -53,6 +53,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 } }
@@ -309,9 +325,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(