Token-Bucket Rate Limiter

Build a token-bucket rate limiter — the workhorse for throttling calls to an LLM API, a downstream tool, or any endpoint with a quota. It is the mechanism that lets you absorb a short burst of traffic while still holding the long-run average to a fixed rate.

How a token bucket works

Picture a bucket that holds up to capacity tokens and is topped up continuously at refill_rate tokens per second, never overflowing past capacity. Every request must take cost tokens out to proceed. If the bucket has enough, the request is allowed and the tokens are removed; if not, the request is throttled and nothing is spent.

You do not need a background timer to add tokens. Refill lazily: on each allow() call, look at how much time has passed since the previous call and credit elapsed * refill_rate tokens (clamped to capacity). This is exact, thread-free, and driven entirely by the injected now() clock.

Burst vs sustained rate

Two numbers describe the limiter, and interviewers want you to separate them:

• Sustained rate R (refill_rate) is the long-run ceiling. Over a long window you can make at most about R requests per second, because that is how fast tokens come back.
• Burst B (capacity) is the slack. After an idle period the bucket is full, so up to B requests can fire back-to-back in an instant before the bucket empties and you fall back to the drip rate R.

So capacity=10, refill_rate=1 means “1 request/second sustained, but up to 10 in a sudden burst.” Raising capacity makes the limiter more forgiving of spikes; raising the rate raises the steady throughput.

Task

Implement the TokenBucket class:

1. __init__(capacity, refill_rate, now) stores the parameters, starts the bucket full (tokens = capacity), and records the current time from now() as the last-refill timestamp. If now is None, default it to time.monotonic.
2. allow(cost=1.0):
   • Read the current time t = now(), compute elapsed = t - last_time, add elapsed * refill_rate tokens capped at capacity, and set last_time = t.
   • If tokens >= cost, subtract cost and return True.
   • Otherwise return False without changing the token count.

Example

clock = [0.0]                      # mutable clock the test advances by hand
bucket = TokenBucket(capacity=2, refill_rate=1.0, now=lambda: clock[0])

bucket.allow()   # -> True   (tokens 2 -> 1)
bucket.allow()   # -> True   (tokens 1 -> 0)
bucket.allow()   # -> False  (empty, no time has passed)

clock[0] = 1.0                     # one second later: +1 token
bucket.allow()   # -> True   (refilled to 1, then 1 -> 0)

Constraints

• Do not call time.sleep, spawn a thread, or use a background timer — refill is computed from elapsed time on each call.
• Never let the token count exceed capacity after a refill, and never let it go below zero.
• A throttled call (returns False) must not consume tokens; only the refill portion may change the count.
• Read the clock exactly through the injected now() so tests are deterministic; default to time.monotonic when now is not supplied.

Approach

Keep two pieces of state: the current token count and the timestamp of the last refill. On every allow() call, first do a lazy refill — credit elapsed * refill_rate tokens for the time that has passed since the last call, clamped to capacity — then either spend cost tokens and return True, or leave the bucket untouched and return False. Injecting now() makes time an argument, so refill is exact and the whole thing is deterministic and thread-free.

Solution

import time
from typing import Callable

class TokenBucket:
    def __init__(
        self,
        capacity: float,
        refill_rate: float,
        now: Callable[[], float] = None,
    ) -> None:
        self.capacity = capacity
        self.refill_rate = refill_rate
        self.now = now if now is not None else time.monotonic
        self.tokens = capacity            # start full: a full burst is available
        self.last_time = self.now()

    def _refill(self) -> None:
        t = self.now()
        elapsed = t - self.last_time
        self.tokens = min(self.capacity, self.tokens + elapsed * self.refill_rate)
        self.last_time = t

    def allow(self, cost: float = 1.0) -> bool:
        self._refill()
        if self.tokens >= cost:
            self.tokens -= cost           # spend only on success
            return True
        return False                      # throttled: refill happened, no spend

Walkthrough

On the example (capacity=2, refill_rate=1.0, clock starts at 0.0):

1. allow() at t=0: refill credits 0 * 1.0 = 0; tokens = 2 >= 1, spend one → tokens = 1, return True.
2. allow() at t=0: no elapsed time, tokens = 1 >= 1, spend one → tokens = 0, return True.
3. allow() at t=0: tokens = 0 < 1, return False; the bucket is unchanged.
4. Advance clock[0] = 1.0. allow() at t=1: elapsed 1.0, credit 1.0 * 1.0 = 1 token → tokens = min(2, 1) = 1; 1 >= 1, spend one → tokens = 0, return True.

The first two calls exhaust the burst, the third is throttled, and one second buys exactly one more token — the sustained rate of 1/sec.

Complexity & notes

• Time is O(1) per allow() call and space is O(1) — two floats of state, no history and no timer thread.
• The min(capacity, ...) clamp is what caps the burst: tokens accrued during a long idle period cannot pile up past capacity, so an idle client cannot later fire an unbounded flood.
• Spending only inside the tokens >= cost branch is the correctness crux — a throttled request must leave the count alone (apart from the refill), otherwise repeated blocked calls would drive tokens negative and distort the next decision.
• Burst vs sustained falls straight out of the two knobs: after idle, up to capacity requests pass immediately (burst B); once drained, throughput is pinned to refill_rate because that is the only source of new tokens (sustained R).
• Using a monotonic clock (not time.time) matters in production: wall-clock adjustments or NTP steps can jump backwards, producing negative elapsed and a spurious token loss or gain. With a monotonic source elapsed is always ≥ 0. For strict wall-clock safety you could also guard with max(0.0, elapsed).
• Real limiters extend this with per-key buckets (one per user or API key), a retry_after hint computed as (cost - tokens) / refill_rate, and shared state in Redis for multi-process fleets — but the single-bucket lazy-refill core here is the piece interviewers grade.