Problem 84 · Concurrency & Python

Thread-Safe Token Bucket Rate Limiter

MediumClassic DSA & systems

01Problem

Implement a token bucket with a refill rate and burst capacity that supports try_acquire(n) and reports how long to wait, safe under concurrent callers.

02Key insight

Refill lazily on each call from elapsed time instead of using a background timer; inject the clock for deterministic tests. Guard refill+consume with one lock so the check-then-act is atomic.

03Solution

solution.pypython33 lines

import threading
import time
from typing import Callable


class TokenBucket:
    def __init__(self, rate: float, capacity: float,
                 clock: Callable[[], float] = time.monotonic) -> None:
        self.rate = rate          # tokens per second
        self.capacity = capacity  # max burst
        self._clock = clock
        self._tokens = capacity
        self._last = clock()
        self._lock = threading.Lock()

    def _refill(self) -> None:
        now = self._clock()
        self._tokens = min(self.capacity, self._tokens + (now - self._last) * self.rate)
        self._last = now

    def try_acquire(self, n: float = 1) -> bool:
        with self._lock:
            self._refill()
            if self._tokens >= n:
                self._tokens -= n
                return True
            return False

    def wait_time(self, n: float = 1) -> float:
        """Seconds until n tokens would be available (0 if available now)."""
        with self._lock:
            self._refill()
            return max(0.0, (n - self._tokens) / self.rate)

Try it first.

TimeO(1) per call

SpaceO(1)

tests.py · 4 passing

tests.pypython16 lines

import threading
now = [0.0]
tb = TokenBucket(rate=2, capacity=5, clock=lambda: now[0])
assert sum(tb.try_acquire() for _ in range(10)) == 5
assert tb.wait_time(1) == 0.5
now[0] = 1.0
assert tb.try_acquire(2) and not tb.try_acquire()
# concurrency: frozen clock, 8 threads hammering -> exactly capacity succeed
tb2 = TokenBucket(rate=1, capacity=10, clock=lambda: 0.0)
wins = []
def worker():
    wins.append(sum(tb2.try_acquire() for _ in range(100)))
ts = [threading.Thread(target=worker) for _ in range(8)]
for t in ts: t.start()
for t in ts: t.join()
assert sum(wins) == 10

04Staff-level follow-up

How do you rate-limit across 50 API servers? Keep the bucket in Redis and do refill+consume atomically in a Lua script keyed by tenant; or give each node a local bucket with rate/N and periodically rebalance to avoid a hot central dependency.