fix(ets): return the real time-to-next-token on token bucket deny

guides/Tutorial.md

@@ -115,8 +115,13 @@ defp rate_limit_videos(conn, _opts) do
       conn
 
     {:deny, retry_after} ->
+      # `retry-after` is in whole seconds, and `retry_after` may be under a
+      # second, so round UP and floor at 1. Truncating with `div/2` alone can
+      # emit `retry-after: 0`, which tells the client to retry immediately.
+      retry_after_seconds = max(ceil(retry_after / 1000), 1)
+
       conn
-      |> put_resp_header("retry-after", Integer.to_string(div(retry_after, 1000)))
+      |> put_resp_header("retry-after", Integer.to_string(retry_after_seconds))
       |> send_resp(429, [])
       |> halt()
   end
@@ -144,7 +149,9 @@ defmodule MyAppWeb.Endpoint do
         conn
 
       {:deny, retry_after} ->
-        retry_after_seconds = div(retry_after, 1000)
+        # Round UP and floor at 1 -- see the note above; a sub-second wait
+        # truncated with `div/2` becomes `retry-after: 0`.
+        retry_after_seconds = max(ceil(retry_after / 1000), 1)
 
         conn
         |> put_resp_header("retry-after", Integer.to_string(retry_after_seconds))

lib/hammer/ets/token_bucket.ex

@@ -139,7 +139,20 @@ defmodule Hammer.ETS.TokenBucket do
       :ets.insert(table, {key, final_level, now})
       {:allow, final_level}
     else
-      {:deny, 1000}
+      # Time until the bucket holds enough tokens to pay `cost`. The other
+      # algorithms already answer with a real wait (`expires_at - now`), and
+      # the callback types this value as a `timeout`; the buckets returned a
+      # flat 1000 because whole-second refill could not express anything
+      # finer. Millisecond refill can: at 55 tokens/sec a caller short one
+      # token waits 19ms, not a second.
+      #
+      # Integer ceiling division, so the answer is exact at every magnitude
+      # and never rounds down into a wait that is still too short. Sleeping
+      # the returned value is always sufficient: it is at least
+      # `deficit * 1000 / refill_rate` ms, which accrues at least `deficit`
+      # whole tokens on top of whatever the current `trunc` already credited.
+      deficit = cost - current_tokens
+      {:deny, max(div(deficit * 1000 + refill_rate - 1, refill_rate), 1)}
     end
   end
 

test/hammer/ets/token_bucket_test.exs

@@ -80,6 +80,67 @@ defmodule Hammer.ETS.TokenBucketTest do
       assert {:allow, _} = TokenBucket.hit(table, key, refill_rate, capacity, 1)
     end
 
+    test "denies with the real time until the next token, not a flat second", %{table: table} do
+      key = "key"
+      # One token every ~18.18ms.
+      refill_rate = 55
+      capacity = 10
+
+      assert {:allow, 0} = TokenBucket.hit(table, key, refill_rate, capacity, capacity)
+
+      # Short exactly one token: ceil(1000 / 55) == 19ms. A flat 1000 here
+      # makes the caller sleep ~53x longer than the limiter actually requires.
+      assert {:deny, retry_after} = TokenBucket.hit(table, key, refill_rate, capacity, 1)
+      assert retry_after == 19
+
+      # The wait scales with how many tokens the cost is short of.
+      assert {:deny, retry_after_5} = TokenBucket.hit(table, key, refill_rate, capacity, 5)
+      assert retry_after_5 == 91
+    end
+
+    test "the deny wait is long enough to actually pay the cost", %{table: table} do
+      key = "key"
+      refill_rate = 55
+      capacity = 10
+
+      assert {:allow, 0} = TokenBucket.hit(table, key, refill_rate, capacity, capacity)
+      assert {:deny, retry_after} = TokenBucket.hit(table, key, refill_rate, capacity, 3)
+
+      # Sleeping the advertised wait must be sufficient -- a wait that rounds
+      # down would hand callers a retry that is still denied.
+      :timer.sleep(retry_after)
+      assert {:allow, _} = TokenBucket.hit(table, key, refill_rate, capacity, 3)
+    end
+
+    test "the advertised wait is sufficient across rates, costs and deficits", %{table: table} do
+      # The refill path truncates DOWN and the deny path rounds UP, so it is
+      # worth pinning that the two cannot conspire to hand a caller a wait
+      # that leaves it still denied. Simulated rather than slept: seeding
+      # `last_update` back by exactly the advertised wait is equivalent to the
+      # caller sleeping it, and keeps the sweep deterministic and instant.
+      for refill_rate <- [1, 2, 3, 7, 55, 100, 999, 1_000_000],
+          capacity <- [1, 10, 1000],
+          cost <- [1, 3, capacity],
+          cost <= capacity,
+          starting_level <- [0, div(cost, 2)],
+          starting_level < cost do
+        key = "sufficient:#{refill_rate}:#{capacity}:#{cost}:#{starting_level}"
+        :ets.insert(table, {key, starting_level, System.system_time(:millisecond)})
+
+        assert {:deny, wait} = TokenBucket.hit(table, key, refill_rate, capacity, cost)
+        assert wait >= 1
+
+        # Rewind the stored clock by exactly the advertised wait -- the state
+        # the caller would find itself in having slept precisely that long.
+        [{^key, level, last_update}] = :ets.lookup(table, key)
+        :ets.insert(table, {key, level, last_update - wait})
+
+        assert {:allow, _} = TokenBucket.hit(table, key, refill_rate, capacity, cost),
+               "wait of #{wait}ms was too short for rate=#{refill_rate} " <>
+                 "cap=#{capacity} cost=#{cost} level=#{starting_level}"
+      end
+    end
+
     test "handles costs greater than 1 correctly", %{table: table} do
       key = "key"
       refill_rate = 2