Skip to content
GreenSnippets

Token bucket rate limiter

Allow bursts but cap the sustained request rate with a small thread-safe token bucket.

Python
import threading
import time


class TokenBucket:
    def __init__(self, rate, capacity, clock=time.monotonic):
        self.rate = rate
        self.capacity = capacity
        self._tokens = float(capacity)
        self._clock = clock
        self._updated = clock()
        self._lock = threading.Lock()

    def acquire(self, cost=1):
        """Return 0 if allowed, otherwise seconds to wait before retrying."""
        with self._lock:
            now = self._clock()
            self._tokens = min(self.capacity, self._tokens + (now - self._updated) * self.rate)
            self._updated = now
            if self._tokens >= cost:
                self._tokens -= cost
                return 0.0
            return (cost - self._tokens) / self.rate

Example / expected output

Output
bucket = TokenBucket(rate=5, capacity=10)
wait = bucket.acquire()  # 0.0 when allowed

How it works

Tokens refill continuously at rate per second up to capacity. Each call consumes one; if none is available the request is rejected (or the caller can sleep for the returned wait). The clock is injectable, which keeps the class testable without real sleeping.

Related snippets