diff --git a/lib/hammer/ets/token_bucket.ex b/lib/hammer/ets/token_bucket.ex index 2e9a11d..c35e976 100644 --- a/lib/hammer/ets/token_bucket.ex +++ b/lib/hammer/ets/token_bucket.ex @@ -130,13 +130,34 @@ defmodule Hammer.ETS.TokenBucket do {capacity, now} end - new_tokens = trunc((now - last_update) * refill_rate / 1000) + elapsed = now - last_update + new_tokens = trunc(elapsed * refill_rate / 1000) current_tokens = min(capacity, current_level + new_tokens) if current_tokens >= cost do final_level = current_tokens - cost - :ets.insert(table, {key, final_level, now}) + + # Advance the clock only by the time whose tokens we actually credited, + # so the sub-token remainder carries into the next hit. Stamping `now` + # unconditionally discards it, and a caller hitting faster than one + # token-period would then never refill at all: at 55 tokens/sec, hits + # every 5ms each credit trunc(0.275) == 0 tokens forever. + # + # The exception is an overflowing refill. When the bucket filled to + # capacity the surplus is legitimately discarded, so the clock has to + # snap to `now` or a long-idle bucket banks unbounded credit. That only + # applies when tokens actually accrued: if the bucket merely sat at + # capacity and `new_tokens` is 0, nothing overflowed, and the elapsed + # time is still owed to the level this hit is about to draw down. + new_last_update = + if current_tokens == capacity and new_tokens > 0 do + now + else + last_update + trunc(new_tokens * 1000 / refill_rate) + end + + :ets.insert(table, {key, final_level, new_last_update}) {:allow, final_level} else {:deny, 1000} diff --git a/test/hammer/ets/token_bucket_test.exs b/test/hammer/ets/token_bucket_test.exs index d8257c4..ec11094 100644 --- a/test/hammer/ets/token_bucket_test.exs +++ b/test/hammer/ets/token_bucket_test.exs @@ -63,6 +63,132 @@ defmodule Hammer.ETS.TokenBucketTest do TokenBucket.hit(table, key, refill_rate, capacity, 1) end + # The three tests below pin the refill clock's bookkeeping directly rather + # than measuring throughput. They seed `last_update` at a known offset from + # `now` and assert the value written back, so they are deterministic: the + # stored timestamp is derived from `last_update` plus the time actually + # credited, never from when the test happened to run. + test "carries the sub-token remainder instead of discarding it", %{table: table} do + key = "key" + # One token every ~18.18ms. + refill_rate = 55 + capacity = 10 + + now = System.system_time(:millisecond) + seeded_at = now - 30 + :ets.insert(table, {key, 5, seeded_at}) + + # 30ms accrues trunc(30 * 55 / 1000) == 1 token, which is worth 18ms. + # The remaining ~12ms must stay on the clock for the next hit. + assert {:allow, 5} = TokenBucket.hit(table, key, refill_rate, capacity, 1) + + assert [{^key, 5, stored_at}] = :ets.lookup(table, key) + assert stored_at == seeded_at + 18 + end + + test "does not reset the clock when a full bucket accrued no new tokens", %{table: table} do + key = "key" + refill_rate = 55 + capacity = 10 + + now = System.system_time(:millisecond) + seeded_at = now - 10 + :ets.insert(table, {key, capacity, seeded_at}) + + # 10ms accrues trunc(10 * 55 / 1000) == 0 tokens. The bucket reads as full + # only because it already was, so nothing overflowed and nothing may be + # discarded -- this hit drops it to 9, and the 10ms is still owed to that + # level. Stamping `now` here would silently throw the accrual away. + assert {:allow, 9} = TokenBucket.hit(table, key, refill_rate, capacity, 1) + + assert [{^key, 9, stored_at}] = :ets.lookup(table, key) + assert stored_at == seeded_at + end + + test "resets the clock when an idle bucket overflows", %{table: table} do + key = "key" + refill_rate = 55 + capacity = 10 + + now = System.system_time(:millisecond) + :ets.insert(table, {key, capacity, now - 5_000}) + + # 5s would accrue 275 tokens into a bucket that holds 10. The surplus is + # legitimately discarded, so the clock must snap forward -- otherwise an + # idle bucket banks unbounded credit and the next burst is unbounded too. + assert {:allow, 9} = TokenBucket.hit(table, key, refill_rate, capacity, 1) + + assert [{^key, 9, stored_at}] = :ets.lookup(table, key) + assert stored_at >= now + assert stored_at <= now + 1_000 + end + + test "a caller paced at the nominal refill rate does not starve", %{table: table} do + key = "key" + refill_rate = 55 + capacity = 10 + + # Walk 60 hits spaced one whole millisecond faster than a token period + # (18ms vs ~18.18ms) by advancing the seeded clock by hand. Each hit + # credits exactly one token, so a lossless bucket holds its level. If the + # ~0.18ms remainder is dropped per hit the level decays and the caller is + # eventually denied at a rate it was entitled to sustain. + start = System.system_time(:millisecond) + :ets.insert(table, {key, capacity, start}) + + for step <- 1..60 do + [{^key, level, last_update}] = :ets.lookup(table, key) + :ets.insert(table, {key, level, last_update - 18}) + + assert {:allow, _} = TokenBucket.hit(table, key, refill_rate, capacity, 1), + "denied at step #{step} -- the bucket drained while paced under its own refill rate" + end + + # The level settles a little below capacity rather than at it: each time a + # refill tops the bucket out, the overflow branch correctly discards the + # surplus. What matters is that it reaches a steady state instead of + # decaying -- before this fix the same loop is denied by step 11. + assert [{^key, final_level, _}] = :ets.lookup(table, key) + assert final_level >= div(capacity, 2) + end + + test "the stored clock stays between last_update and now", %{table: table} do + # Carrying the remainder means the stored timestamp is deliberately + # allowed to lag `now`. Two bounds keep that safe, and both matter to + # clean/1, which reaps rows by comparing this timestamp against + # `now - key_older_than`: + # + # * it can never run AHEAD of now, or a row outlives its real idleness + # * it can never lag by more than one token period, or an actively + # used row could look idle and be reaped out from under its caller + # + # The upper bound holds because the credited time is at most the elapsed + # time: trunc(new_tokens * 1000 / refill_rate) <= elapsed, always. + for {refill_rate, capacity} <- [{1, 5}, {55, 10}, {100, 1000}, {1_000_000, 10}], + offset <- [0, 1, 9, 18, 500, 5_000] do + key = "bounds:#{refill_rate}:#{capacity}:#{offset}" + before = System.system_time(:millisecond) + seeded_at = before - offset + :ets.insert(table, {key, div(capacity, 2) + 1, seeded_at}) + + assert {:allow, _} = TokenBucket.hit(table, key, refill_rate, capacity, 1) + now_after = System.system_time(:millisecond) + + assert [{^key, _, stored_at}] = :ets.lookup(table, key) + + assert stored_at <= now_after, + "clock ran ahead of now for rate=#{refill_rate} cap=#{capacity} offset=#{offset}" + + assert stored_at >= seeded_at, + "clock moved backwards for rate=#{refill_rate} cap=#{capacity} offset=#{offset}" + + token_period = div(1000, refill_rate) + 1 + + assert now_after - stored_at <= offset + token_period, + "lagged more than one token period for rate=#{refill_rate} cap=#{capacity} offset=#{offset}" + end + end + test "refills at sub-second granularity when refill_rate > capacity", %{table: table} do key = "key" # Small burst allowance with a higher sustained rate,