104 lines
3.6 KiB
Python
104 lines
3.6 KiB
Python
"""Thread-safe token bucket rate limiter.
|
|
|
|
Correctness notes:
|
|
- Refill is computed from elapsed wall-clock time using a monotonic clock.
|
|
- Tokens are capped at capacity.
|
|
- acquire(n) blocks until n tokens available or timeout; partial consumes never happen.
|
|
- A single lock guards state; condition variable for waiters.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import threading
|
|
import time
|
|
from typing import Optional
|
|
|
|
|
|
class TokenBucket:
|
|
def __init__(self, rate: float, capacity: float):
|
|
if rate <= 0:
|
|
raise ValueError("rate must be > 0")
|
|
if capacity <= 0:
|
|
raise ValueError("capacity must be > 0")
|
|
if rate > capacity:
|
|
# allow; but typically rate <= capacity. No restriction.
|
|
pass
|
|
self._rate = float(rate)
|
|
self._capacity = float(capacity)
|
|
self._tokens = float(capacity)
|
|
self._last = time.monotonic()
|
|
self._lock = threading.Lock()
|
|
self._cond = threading.Condition(self._lock)
|
|
|
|
@property
|
|
def rate(self) -> float:
|
|
return self._rate
|
|
|
|
@property
|
|
def capacity(self) -> float:
|
|
return self._capacity
|
|
|
|
def _refill(self, now: float) -> None:
|
|
elapsed = now - self._last
|
|
if elapsed <= 0:
|
|
return
|
|
self._tokens = min(self._capacity, self._tokens + elapsed * self._rate)
|
|
self._last = now
|
|
|
|
def tokens(self) -> float:
|
|
with self._lock:
|
|
self._refill(time.monotonic())
|
|
return self._tokens
|
|
|
|
def try_acquire(self, n: float = 1.0) -> bool:
|
|
if n <= 0:
|
|
raise ValueError("n must be > 0")
|
|
if n > self._capacity:
|
|
raise ValueError("n exceeds capacity")
|
|
with self._lock:
|
|
now = time.monotonic()
|
|
self._refill(now)
|
|
if self._tokens >= n - 1e-12:
|
|
self._tokens -= n
|
|
return True
|
|
return False
|
|
|
|
def acquire(self, n: float = 1.0, timeout: Optional[float] = None) -> bool:
|
|
if n <= 0:
|
|
raise ValueError("n must be > 0")
|
|
if n > self._capacity:
|
|
raise ValueError("n exceeds capacity")
|
|
deadline = None if timeout is None else time.monotonic() + timeout
|
|
with self._lock:
|
|
while True:
|
|
now = time.monotonic()
|
|
self._refill(now)
|
|
if self._tokens >= n - 1e-12:
|
|
self._tokens -= n
|
|
return True
|
|
# time until enough tokens
|
|
deficit = n - self._tokens
|
|
wait = deficit / self._rate
|
|
if deadline is not None:
|
|
remaining = deadline - now
|
|
if remaining <= 0:
|
|
return False
|
|
wait = min(wait, remaining)
|
|
# Guard against sub-millisecond waits that degenerate into a
|
|
# busy spin on spurious wakeups (cond.wait(1e-15) returns at
|
|
# once). Never wait longer than the allowed remaining time.
|
|
if wait < 1e-3:
|
|
if deadline is not None:
|
|
if remaining <= 1e-3:
|
|
return False
|
|
wait = min(remaining, 1e-3)
|
|
else:
|
|
wait = 1e-3
|
|
self._cond.wait(timeout=wait)
|
|
# loop rechecks
|
|
|
|
def reset(self, tokens: Optional[float] = None) -> None:
|
|
with self._lock:
|
|
self._tokens = self._capacity if tokens is None else float(tokens)
|
|
self._last = time.monotonic()
|
|
self._cond.notify_all()
|