Build a token bucket rate limiter for an LLM gateway

The question

You’re building the gateway every internal team uses to call LLMs. Implement a rate limiter so each team gets a fair share. It must support limits in requests per second and in tokens per minute.

The second sentence is what makes this an LLM question. Requests are not equal: one call can be 200 tokens and the next 20,000. A limiter that only counts requests lets one team with long prompts use up everyone’s provider quota.

Clarify first

• Per what key? (Per team, possibly per team and model.)
• On rejection, block or fail fast? (Fail fast, and return how long to wait so the caller can back off.)
• Are bursts allowed? (Yes, up to a capacity. That’s the point of a token bucket over a fixed window.)
• One process or many gateway instances? (Start with one; distributed is the follow-up.)
• How do we know the token cost before the call? (Not exactly. Estimate it from the prompt length plus max_tokens, then reconcile afterwards. See the follow-ups.)

Approach

Keep one token bucket per key. Tokens refill continuously at rate per second, up to capacity. A request with cost c is allowed if the bucket holds at least c. The refill is computed lazily on each call from the elapsed time, so there are no background timers. One lock guards the dictionary, because the read-modify-write must be atomic.

import threading
import time


class TokenBucketLimiter:
    """Per-key token bucket: `rate` tokens are added per second, up to `capacity`."""

    def __init__(self, rate: float, capacity: float, clock=time.monotonic):
        if rate <= 0 or capacity <= 0:
            raise ValueError("rate and capacity must be positive")
        self.rate = rate
        self.capacity = capacity
        self._clock = clock
        self._buckets: dict = {}  # key -> [tokens, last_refill]
        self._lock = threading.Lock()

    def allow(self, key, cost: float = 1.0) -> bool:
        """Take `cost` tokens for `key` if available. Never blocks."""
        if cost > self.capacity:
            return False  # can never succeed; don't let it wait forever
        with self._lock:
            now = self._clock()
            tokens, last = self._buckets.get(key, (self.capacity, now))
            tokens = min(self.capacity, tokens + (now - last) * self.rate)
            allowed = tokens >= cost
            if allowed:
                tokens -= cost
            self._buckets[key] = (tokens, now)
            return allowed

    def retry_after(self, key, cost: float = 1.0) -> float:
        """Seconds until `cost` tokens will be available for `key`."""
        with self._lock:
            now = self._clock()
            tokens, last = self._buckets.get(key, (self.capacity, now))
            tokens = min(self.capacity, tokens + (now - last) * self.rate)
            return max(0.0, (cost - tokens) / self.rate)

For tokens per minute, use the same class with rate = TPM / 60 and capacity = TPM (or a smaller burst), and pass the token estimate as cost.

Tests

import threading

import pytest

from token_bucket import TokenBucketLimiter


class FakeClock:
    def __init__(self):
        self.t = 0.0

    def __call__(self):
        return self.t


def test_allows_burst_up_to_capacity_then_blocks():
    clock = FakeClock()
    rl = TokenBucketLimiter(rate=1, capacity=3, clock=clock)
    assert [rl.allow("u") for _ in range(4)] == [True, True, True, False]


def test_refills_over_time_without_exceeding_capacity():
    clock = FakeClock()
    rl = TokenBucketLimiter(rate=2, capacity=3, clock=clock)
    for _ in range(3):
        rl.allow("u")
    clock.t = 0.5  # +1 token
    assert rl.allow("u") is True
    assert rl.allow("u") is False
    clock.t = 100  # refills to capacity, not 200 tokens
    assert sum(rl.allow("u") for _ in range(10)) == 3


def test_keys_are_independent():
    clock = FakeClock()
    rl = TokenBucketLimiter(rate=1, capacity=1, clock=clock)
    assert rl.allow("a") and rl.allow("b")
    assert not rl.allow("a")


def test_weighted_cost_for_tokens_per_minute_limits():
    clock = FakeClock()
    rl = TokenBucketLimiter(rate=1000, capacity=1000, clock=clock)
    assert rl.allow("team", cost=800)
    assert not rl.allow("team", cost=300)
    assert rl.retry_after("team", cost=300) == pytest.approx(0.1)


def test_cost_larger_than_capacity_is_rejected():
    rl = TokenBucketLimiter(rate=1, capacity=5)
    assert rl.allow("u", cost=6) is False


def test_thread_safety_never_over_admits():
    clock = FakeClock()  # frozen time: no refill during the test
    rl = TokenBucketLimiter(rate=1, capacity=100, clock=clock)
    results = []
    lock = threading.Lock()

    def worker():
        for _ in range(50):
            ok = rl.allow("shared")
            with lock:
                results.append(ok)

    threads = [threading.Thread(target=worker) for _ in range(8)]
    for t in threads:
        t.start()
    for t in threads:
        t.join()
    assert sum(results) == 100

The last test matters most. With time frozen, eight threads making 400 attempts against a capacity of 100 must admit exactly 100. Without the lock it over-admits, but only intermittently, which is exactly why you write the test.

Complexity

O(1) per call and O(number of keys) space. Mention that idle keys should be expired, or the dict grows forever.

Follow-ups the interviewer will ask

1. Many gateway instances. In-process buckets give each instance its own quota. Move the bucket state to Redis and do the refill-and-take in one Lua script, so it’s atomic across instances. Keep the token count and the last refill time in one hash per key. The cost is one network round trip per request, around a millisecond inside a data center, which is negligible next to an LLM call.
2. The token cost is only known after the call. Reserve an estimate up front (prompt tokens plus max_tokens), then refund or charge the difference when the response arrives. Callers who set a huge max_tokens pay for it in quota, which nudges them to set it honestly.
3. Fairness under contention. When the provider quota is the real bottleneck, per-team limits that add up to more than the quota will still overload it. Put a global bucket in front of the per-team buckets, and decide which low-priority traffic gets shed first.
4. Priorities. Interactive traffic should beat batch jobs. Give them separate buckets, or let batch borrow only unused interactive capacity.
5. Why not a sliding window log? It’s exact, but uses O(requests) memory per key. A sliding window counter is a cheaper approximation. The token bucket allows controlled bursts with O(1) memory, so it’s the usual choice.

Common mistakes

• Refilling past capacity after a long idle period.
• Not handling a cost larger than capacity, so the request can never succeed and the caller retries forever.
• Taking the lock only for the write, so two threads both see “enough tokens”.
• Using wall-clock time, which can jump backwards. Use a monotonic clock.