Mark throttler re-checks after a timeout with X-Is-Retry

CHANGELOG.md

@@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- `is_retry` request option, which sends the `X-Is-Retry: true` header with that request
+- `mark_timeout_retries` throttler option, off by default, which sends the `X-Is-Retry: true` header on checks that repeat a check which timed out or got an HTTP 408 response
+
 ## [0.9.0] - 2024-04-18
 
 ### Added

README.md

@@ -209,7 +209,8 @@ module Freno
                     instrumenter: Instrumenter::Noop,
                     circuit_breaker: CircuitBreaker::Noop,
                     wait_seconds: DEFAULT_WAIT_SECONDS,
-                    max_wait_seconds: DEFAULT_MAX_WAIT_SECONDS)
+                    max_wait_seconds: DEFAULT_MAX_WAIT_SECONDS,
+                    mark_timeout_retries: false)
 
 
       @client           = client
@@ -219,6 +220,7 @@ module Freno
       @circuit_breaker  = circuit_breaker
       @wait_seconds     = wait_seconds
       @max_wait_seconds = max_wait_seconds
+      @mark_timeout_retries = mark_timeout_retries
 
       yield self if block_given?
 
@@ -238,6 +240,18 @@ You optionally provide the time you want the throttler to sleep in case the chec
 If replication lags badly, you can control until when you want to keep sleeping
 and retrying the check by setting `max_wait_seconds`. When that times out, the throttle will raise a `Freno::Throttler::WaitedTooLong` error.
 
+#### Marking retries after timeouts
+
+Set `mark_timeout_retries: true` to send the `X-Is-Retry: true` header when the throttler repeats a check that failed by timing out. A check fails by timing out when the client gets an HTTP 408 response, or when the client is configured with `raise_on_timeout: false` and the request times out. Load balancers can use the header to identify retry traffic.
+
+Checks repeated after freno answered that a store is not ok are regular polling, not retries, and never carry the header. The first check in each `throttle` call never carries it either. The option is off by default. When it is on, the throttler calls `client.check` instead of `client.check?`.
+
+The header can also be set on a single client request with the `is_retry` option:
+
+```ruby
+freno.check?(options: { is_retry: true })
+```
+
 #### Instrumenting the throttler
 
 You can also configure the throttler with an `instrumenter` collaborator to subscribe to events happening during the `throttle` call.

lib/freno/client.rb

@@ -54,6 +54,14 @@ class DecorationError < ArgumentError; end
     #  freno.check?(options: { low_priority: true })
     #  ```
     #
+    #  The `is_retry: true` option sends the `X-Is-Retry: true` header with
+    #  that single request. Use it only when the request repeats a previously
+    #  failed request, so load balancers can identify retry traffic:
+    #
+    #  ```ruby
+    #  freno.check?(options: { is_retry: true })
+    #  ```
+    #
     def initialize(faraday)
       @faraday            = faraday
       @default_store_type = :mysql

lib/freno/client/request.rb

@@ -9,6 +9,10 @@ class Client
     class Request
       include Freno::Client::Preconditions
 
+      # Header marking a request as a retry of a previously failed request, so
+      # load balancers can identify retry traffic.
+      RETRY_HEADER = "X-Is-Retry"
+
       attr_reader :faraday, :args, :options, :raise_on_timeout
 
       def self.perform(**kwargs)
@@ -22,10 +26,12 @@ def initialize(**kwargs)
 
         @raise_on_timeout = options.fetch(:raise_on_timeout, true)
         @verb = options.fetch(:verb, :head)
+
+        headers[RETRY_HEADER] = "true" if options[:is_retry]
       end
 
       def perform
-        response = request(verb, path, params)
+        response = headers.empty? ? request(verb, path, params) : request(verb, path, params, headers)
         process_response(response)
       rescue Faraday::TimeoutError => error
         raise Freno::Error, error if raise_on_timeout
@@ -37,8 +43,8 @@ def perform
 
       protected
 
-      def request(verb, path, params)
-        faraday.send(verb, path, params)
+      def request(verb, path, params, headers = nil)
+        faraday.send(verb, path, params, headers)
       end
 
       def path
@@ -57,6 +63,10 @@ def params
         @params ||= {}
       end
 
+      def headers
+        @headers ||= {}
+      end
+
       def process_response(response)
         Result.from_faraday_response(response)
       end

lib/freno/throttler.rb

@@ -1,5 +1,6 @@
 # frozen_string_literal: true
 
+require "set"
 require "freno/client"
 require "freno/throttler/errors"
 require "freno/throttler/mapper"
@@ -51,7 +52,8 @@ class Throttler
                   :instrumenter,
                   :circuit_breaker,
                   :wait_seconds,
-                  :max_wait_seconds
+                  :max_wait_seconds,
+                  :mark_timeout_retries
 
     # Initializes a new instance of the throttler
     #
@@ -99,14 +101,23 @@ class Throttler
     #     seconds the throttler will wait in total for replicas to catch-up
     #     before raising a `WaitedTooLong` error.
     #
+    #  - `:mark_timeout_retries`: A boolean, `false` by default. When `true`,
+    #     a check that repeats a check for the same store which failed by
+    #     timing out (a client-side timeout with `raise_on_timeout: false`, or
+    #     an HTTP 408 response) is sent with the `X-Is-Retry: true` header.
+    #     Checks that repeat after freno answered that a store is not ok are
+    #     regular polling and are never marked. Enabling this calls
+    #     `client.check` instead of `client.check?`.
+    #
     def initialize(
       client: nil,
       app: nil,
       mapper: Mapper::Identity,
       instrumenter: Instrumenter::Noop,
       circuit_breaker: CircuitBreaker::Noop,
       wait_seconds: DEFAULT_WAIT_SECONDS,
-      max_wait_seconds: DEFAULT_MAX_WAIT_SECONDS
+      max_wait_seconds: DEFAULT_MAX_WAIT_SECONDS,
+      mark_timeout_retries: false
     )
       @client           = client
       @app              = app
@@ -115,6 +126,7 @@ def initialize(
       @circuit_breaker  = circuit_breaker
       @wait_seconds     = wait_seconds
       @max_wait_seconds = max_wait_seconds
+      @mark_timeout_retries = mark_timeout_retries
 
       yield self if block_given?
 
@@ -165,14 +177,15 @@ def throttle(context = nil, **options)
       store_names = mapper.call(context)
       instrument(:called, store_names: store_names)
       waited = 0
+      timed_out_store_names = Set.new
 
       while true
         unless circuit_breaker.allow_request?
           instrument(:circuit_open, store_names: store_names, waited: waited)
           raise CircuitOpen
         end
 
-        if all_stores_ok?(store_names, **options)
+        if all_stores_ok?(store_names, timed_out_store_names, **options)
           instrument(:succeeded, store_names: store_names, waited: waited)
           circuit_breaker.success
           break
@@ -208,16 +221,31 @@ def validate_args
       raise ArgumentError, errors.join("\n") if errors.any?
     end
 
-    def all_stores_ok?(store_names, **options)
+    def all_stores_ok?(store_names, timed_out_store_names, **options)
       store_names.all? do |store_name|
-        client.check?(app: app, store_name: store_name, options: options)
+        store_ok?(store_name, timed_out_store_names, **options)
       end
     rescue Freno::Error => error
       instrument(:freno_errored, store_names: store_names, error: error)
       circuit_breaker.failure
       raise ClientError, error
     end
 
+    def store_ok?(store_name, timed_out_store_names, **options)
+      return client.check?(app: app, store_name: store_name, options: options) unless mark_timeout_retries
+
+      check_options = timed_out_store_names.include?(store_name) ? options.merge(is_retry: true) : options
+      result = client.check(app: app, store_name: store_name, options: check_options)
+
+      if result.meaning == :request_timeout
+        timed_out_store_names << store_name
+      else
+        timed_out_store_names.delete(store_name)
+      end
+
+      result.ok?
+    end
+
     def wait
       sleep wait_seconds
     end

test/freno/client/requests/check_test.rb

@@ -79,4 +79,38 @@ def test_perform_calls_the_proper_service_endpoint_with_low_priority_and_succeed
     assert_equal 200, response.code
     assert_equal({ "StatusCode" => 200, "Value" => 0.025075, "Threshold" => 1, "Message" => "" }, response.body)
   end
+
+  def test_perform_sends_the_retry_header_when_marked_as_retry
+    sent_header = :not_called
+    faraday = stubbed_faraday do |stub|
+      stub.head("/check/github/mysql/main") do |env|
+        sent_header = env.request_headers["X-Is-Retry"]
+        [200, {}, ""]
+      end
+    end
+
+    Check.new(
+      faraday: faraday,
+      app: "github",
+      store_type: "mysql",
+      store_name: "main",
+      options: { is_retry: true }
+    ).perform
+
+    assert_equal "true", sent_header
+  end
+
+  def test_perform_does_not_send_the_retry_header_by_default
+    sent_header = :not_called
+    faraday = stubbed_faraday do |stub|
+      stub.head("/check/github/mysql/main") do |env|
+        sent_header = env.request_headers["X-Is-Retry"]
+        [200, {}, ""]
+      end
+    end
+
+    Check.new(faraday: faraday, app: "github", store_type: "mysql", store_name: "main").perform
+
+    assert_nil sent_header
+  end
 end

test/freno/throttler_retry_header_test.rb

@@ -0,0 +1,120 @@
+# frozen_string_literal: true
+
+require "test_helper"
+
+class FrenoThrottlerRetryHeaderTest < ThrottlerTest
+  CHECK_PATH = "/check/github/mysql/mysqla"
+  OTHER_CHECK_PATH = "/check/github/mysql/mysqlb"
+
+  def test_first_check_never_has_the_retry_header
+    responses = { CHECK_PATH => [200] }
+    sent_headers = throttle_with(responses)
+
+    assert_equal [nil], sent_headers[CHECK_PATH]
+  end
+
+  def test_recheck_after_throttled_answer_has_no_retry_header
+    responses = { CHECK_PATH => [417, 429, 200] }
+    sent_headers = throttle_with(responses)
+
+    assert_equal [nil, nil, nil], sent_headers[CHECK_PATH]
+  end
+
+  def test_recheck_after_client_timeout_has_the_retry_header
+    responses = { CHECK_PATH => [:timeout, 200] }
+    sent_headers = throttle_with(responses)
+
+    assert_equal [nil, "true"], sent_headers[CHECK_PATH]
+  end
+
+  def test_recheck_after_408_has_the_retry_header
+    responses = { CHECK_PATH => [408, 200] }
+    sent_headers = throttle_with(responses)
+
+    assert_equal [nil, "true"], sent_headers[CHECK_PATH]
+  end
+
+  def test_retry_header_is_dropped_once_a_check_gets_an_answer
+    responses = { CHECK_PATH => [:timeout, 417, 200] }
+    sent_headers = throttle_with(responses)
+
+    assert_equal [nil, "true", nil], sent_headers[CHECK_PATH]
+  end
+
+  def test_retry_header_applies_only_to_the_store_that_timed_out
+    responses = { CHECK_PATH => [:timeout, 200], OTHER_CHECK_PATH => [200] }
+    sent_headers = throttle_with(responses, store_names: %i[mysqla mysqlb])
+
+    assert_equal [nil, "true"], sent_headers[CHECK_PATH]
+    assert_equal [nil], sent_headers[OTHER_CHECK_PATH]
+  end
+
+  def test_retry_header_does_not_leak_into_the_next_throttle_call
+    responses = { CHECK_PATH => [:timeout, 200, 200] }
+    sent_headers = {}
+    throttler = build_throttler(responses, sent_headers)
+
+    throttler.throttle(:mysqla) { nil }
+    throttler.throttle(:mysqla) { nil }
+
+    assert_equal [nil, "true", nil], sent_headers[CHECK_PATH]
+  end
+
+  def test_retry_header_is_not_sent_unless_enabled
+    responses = { CHECK_PATH => [:timeout, 408, 200] }
+    sent_headers = throttle_with(responses, mark_timeout_retries: false)
+
+    assert_equal [nil, nil, nil], sent_headers[CHECK_PATH]
+  end
+
+  def test_mark_timeout_retries_defaults_to_false
+    throttler = Freno::Throttler.new(client: sample_client, app: :github)
+
+    refute throttler.mark_timeout_retries
+  end
+
+  private
+
+  def throttle_with(responses, store_names: :mysqla, mark_timeout_retries: true)
+    sent_headers = {}
+    throttler = build_throttler(responses, sent_headers, mark_timeout_retries: mark_timeout_retries)
+    block_called = false
+
+    throttler.throttle(store_names) do
+      block_called = true
+    end
+
+    assert block_called, "block should have been called"
+    sent_headers
+  end
+
+  def build_throttler(responses, sent_headers, mark_timeout_retries: true)
+    stubs = Faraday::Adapter::Test::Stubs.new do |stub|
+      responses.each do |path, outcomes|
+        remaining = outcomes.dup
+        stub.head(path) do |env|
+          (sent_headers[path] ||= []) << env.request_headers[Freno::Client::Request::RETRY_HEADER]
+          outcome = remaining.shift
+          raise Faraday::TimeoutError, "timed out" if outcome == :timeout
+
+          [outcome, {}, ""]
+        end
+      end
+    end
+    faraday = Faraday.new { |builder| builder.adapter :test, stubs }
+    client = Freno::Client.new(faraday) do |freno|
+      freno.default_store_type = :mysql
+      freno.options[:raise_on_timeout] = false
+    end
+
+    throttler = Freno::Throttler.new(
+      client: client,
+      app: :github,
+      wait_seconds: 1,
+      max_wait_seconds: 30,
+      mark_timeout_retries: mark_timeout_retries
+    )
+    throttler.stubs(:wait)
+    throttler
+  end
+end