88 lines
2.6 KiB
Python
88 lines
2.6 KiB
Python
"""Token-bucket rate limiter with injectable clock and thread safety."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import threading
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
|
|
|
|
class RateLimitError(ValueError):
|
|
pass
|
|
|
|
|
|
@dataclass
|
|
class TokenBucket:
|
|
rate: float
|
|
capacity: float
|
|
_tokens: float = field(init=False, repr=False)
|
|
_last: float = field(init=False, repr=False)
|
|
_lock: threading.Lock = field(init=False, repr=False, default_factory=threading.Lock)
|
|
_clock: object = field(init=False, repr=False, default=None)
|
|
|
|
def __post_init__(self) -> None:
|
|
if self.rate <= 0:
|
|
raise ValueError("rate must be > 0")
|
|
if self.capacity <= 0:
|
|
raise ValueError("capacity must be > 0")
|
|
if self._clock is None:
|
|
self._clock = time.monotonic
|
|
self._tokens = self.capacity
|
|
self._last = self._clock()
|
|
|
|
def tokens(self) -> float:
|
|
with self._lock:
|
|
self._refill()
|
|
return self._tokens
|
|
|
|
def acquire(self, n: float = 1.0, *, blocking: bool = True,
|
|
timeout: float | None = None) -> bool:
|
|
if n < 0:
|
|
raise ValueError("n must be >= 0")
|
|
if n == 0:
|
|
return True
|
|
if n > self.capacity and not blocking:
|
|
return False
|
|
if n > self.capacity:
|
|
raise RateLimitError(f"request {n} exceeds capacity {self.capacity}")
|
|
|
|
deadline = None
|
|
if timeout is not None:
|
|
deadline = self._clock() + timeout
|
|
|
|
while True:
|
|
with self._lock:
|
|
self._refill()
|
|
if self._tokens >= n:
|
|
self._tokens -= n
|
|
return True
|
|
deficit = n - self._tokens
|
|
wait = deficit / self.rate
|
|
now = self._clock()
|
|
if deadline is not None and now + wait > deadline:
|
|
return False
|
|
if not blocking:
|
|
return False
|
|
if deadline is not None:
|
|
wait = min(wait, deadline - now)
|
|
sleep_for = wait
|
|
if sleep_for > 0:
|
|
time.sleep(sleep_for)
|
|
|
|
def _refill(self) -> None:
|
|
now = self._clock()
|
|
elapsed = now - self._last
|
|
if elapsed < 0:
|
|
elapsed = 0.0
|
|
self._tokens = min(self.capacity, self._tokens + elapsed * self.rate)
|
|
self._last = now
|
|
|
|
|
|
def make_bucket(rate, capacity, clock=None):
|
|
b = TokenBucket(rate=rate, capacity=capacity)
|
|
if clock is not None:
|
|
b._clock = clock
|
|
b._tokens = capacity
|
|
b._last = clock()
|
|
return b
|