Add a token-bucket rate limiter
Python · Python · advanced · greenfield
Adds a thread-safe token-bucket limiter for outbound API calls. Refill is computed lazily from a monotonic clock on each acquire, the bucket is capped at capacity, and the whole read-modify-write runs under a lock so concurrent workers can't over-draw. An injectable clock keeps the tests deterministic.
Used by a pool of worker threads that call a partner API with a hard requests-per-second contract; over-granting means partner-side bans.
Requirements
- Implement an in-process `TokenBucket` for the API client: `capacity` is the burst size, `refill_rate` is tokens added per second, both positive numbers.
- `try_acquire(tokens=1.0)` is non-blocking: take the tokens and return `True` if available, otherwise return `False` without waiting. `tokens` must be positive; a request larger than `capacity` can never succeed and simply returns `False`.
- The bucket never holds more than `capacity` tokens, no matter how long it sits idle.
- The limiter is shared across worker threads — concurrent `try_acquire` calls must never over-grant (hand out more tokens than exist).
- Refill must be immune to wall-clock changes (NTP adjustments, DST): use a monotonic time source. The constructor's `clock` parameter exists for deterministic tests only.
Files touched
- app/net/token_bucket.py
--- app/net/token_bucket.py
+import threading
+import time
+
+
+class TokenBucket:
+ """Thread-safe token bucket: ``capacity`` burst, ``refill_rate`` tokens/sec."""
+
+ def __init__(self, capacity, refill_rate, clock=time.monotonic):
+ if capacity <= 0 or refill_rate <= 0:
+ raise ValueError("capacity and refill_rate must be positive")
+ self._capacity = float(capacity)
+ self._refill_rate = float(refill_rate)
+ self._clock = clock
+ self._tokens = self._capacity
+ self._updated = clock()
+ self._lock = threading.Lock()
+
+ def try_acquire(self, tokens=1.0):
+ """Take ``tokens`` if available; never blocks."""
+ if tokens <= 0:
+ raise ValueError("tokens must be positive")
+ with self._lock:
+ now = self._clock()
+ elapsed = now - self._updated
+ self._updated = now
+ self._tokens = min(self._capacity, self._tokens + elapsed * self._refill_rate)
+ if self._tokens >= tokens:
+ self._tokens -= tokens
+ return True