fix(ets): carry the sub-token refill remainder across hits

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}

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,