From 7c6e32e4b6216527b500485625a001bbf3e2048a Mon Sep 17 00:00:00 2001
From: Steve Richert <laserlemon@github.com>
Date: Thu, 18 Apr 2024 13:28:18 +0000
Subject: [PATCH] Change wait time behavior to respect max wait time as a hard
 cap

---
 lib/freno/throttler.rb       | 10 ++--
 test/freno/throttler_test.rb | 89 +++++++++++++++++++++++++++++++++---
 2 files changed, 88 insertions(+), 11 deletions(-)

diff --git a/lib/freno/throttler.rb b/lib/freno/throttler.rb
index 51263ac..aec22bb 100644
--- a/lib/freno/throttler.rb
+++ b/lib/freno/throttler.rb
@@ -178,14 +178,14 @@ module Freno
           break
         end
 
-        wait
-        waited += wait_seconds
-        instrument(:waited, store_names: store_names, waited: waited, max: max_wait_seconds)
-
-        if waited > max_wait_seconds
+        if waited + wait_seconds > max_wait_seconds
           instrument(:waited_too_long, store_names: store_names, waited: waited, max: max_wait_seconds)
           circuit_breaker.failure
           raise WaitedTooLong.new(waited_seconds: waited, max_wait_seconds: max_wait_seconds)
+        else
+          wait
+          waited += wait_seconds
+          instrument(:waited, store_names: store_names, waited: waited, max: max_wait_seconds)
         end
       end
 
diff --git a/test/freno/throttler_test.rb b/test/freno/throttler_test.rb
index cbc3477..7053983 100644
--- a/test/freno/throttler_test.rb
+++ b/test/freno/throttler_test.rb
@@ -90,7 +90,7 @@ class FrenoThrottlerTest < ThrottlerTest
       t.mapper = ->(_context) { [:mysqla] }
       t.instrumenter = MemoryInstrumenter.new
     end
-    throttler.expects(:wait).once.returns(0.1)
+    throttler.expects(:wait).once
 
     throttler.throttle do
       block_called = true
@@ -128,11 +128,11 @@ class FrenoThrottlerTest < ThrottlerTest
       t.app = :github
       t.mapper = ->(_context) { [:mysqla] }
       t.instrumenter = MemoryInstrumenter.new
-      t.wait_seconds = 0.1
-      t.max_wait_seconds = 0.3
+      t.wait_seconds = 1
+      t.max_wait_seconds = 3
     end
 
-    throttler.expects(:wait).times(3).returns(0.1)
+    throttler.expects(:wait).times(3)
 
     assert_raises(Freno::Throttler::WaitedTooLong) do
       throttler.throttle do
@@ -151,8 +151,8 @@ class FrenoThrottlerTest < ThrottlerTest
 
     assert_equal 1, waited_too_long_events.count
     assert_equal [:mysqla], waited_too_long_events.first[:store_names]
-    assert_in_delta 0.3, waited_too_long_events.first[:max], 0.01
-    assert_operator waited_too_long_events.first[:waited], :>=, 0.3
+    assert_equal 3, waited_too_long_events.first[:max]
+    assert_equal 3, waited_too_long_events.first[:waited]
 
     assert_equal 0, throttler.instrumenter.count("throttler.freno_errored")
     assert_equal 0, throttler.instrumenter.count("throttler.circuit_open")
@@ -272,4 +272,81 @@ class FrenoThrottlerTest < ThrottlerTest
 
     assert_equal array, result
   end
+
+  # This test ensures that a throttle call will not wait if that wait would
+  # not be followed by another check, making that wait time useless. For
+  # example, consider a throttler with 1s wait time and 3s max wait time:
+  #
+  # C = check all stores
+  # - = 100ms of wait time
+  # X = raise waited too long
+  #
+  #       0s         1s         2s         3s         4s
+  #       ├──────────┼──────────┼──────────┼──────────┤
+  # v0.8: C----------C----------C----------C----------X
+  # v0.9: C----------C----------C----------CX
+  #
+  def test_does_not_wait_longer_than_needed
+    block_called = false
+    client = sample_client
+
+    throttler = Freno::Throttler.new(
+      client: client,
+      app: :github,
+      wait_seconds: 1,
+      max_wait_seconds: 3
+    )
+
+    # We expect to check four times with three
+    # one-second waits between the attempts.
+    client.stubs(:check?).times(4).returns(false)
+    throttler.expects(:wait).times(3)
+
+    assert_raises(Freno::Throttler::WaitedTooLong) do
+      throttler.throttle(:mysqla) do
+        block_called = true
+      end
+    end
+
+    refute block_called, "block should not have been called"
+  end
+
+  # This test ensures that a throttle call will not wait longer than the
+  # configured maximum wait time, even when that maximum doesn't divide by
+  # the configured wait time evenly. For example, consider a throttler with
+  # 2s wait time and 5s max wait time:
+  #
+  # C = check all stores
+  # - = 100ms of wait time
+  # X = raise waited too long
+  #                                             max_wait_seconds ↴
+  #       0s         1s         2s         3s         4s         5s         6s
+  #       ├──────────┼──────────┼──────────┼──────────┼──────────┼──────────┤
+  # v0.8: C---------------------C---------------------C---------------------CX
+  # v0.9: C---------------------C---------------------CX
+  #
+  def test_does_not_exceed_max_wait_time
+    block_called = false
+    client = sample_client
+
+    throttler = Freno::Throttler.new(
+      client: client,
+      app: :github,
+      wait_seconds: 2,
+      max_wait_seconds: 5
+    )
+
+    # We expect to check four times with three
+    # one-second waits between the attempts.
+    client.stubs(:check?).times(3).returns(false)
+    throttler.expects(:wait).times(2)
+
+    assert_raises(Freno::Throttler::WaitedTooLong) do
+      throttler.throttle(:mysqla) do
+        block_called = true
+      end
+    end
+
+    refute block_called, "block should not have been called"
+  end
 end