Spaces:
Running
Running
| """In-memory token bucket. | |
| Clients are keyed by a salted hash of their address. The salt is random, lives | |
| only in this process, and rotates hourly, so the key can't be reversed, joined | |
| across restarts, or stored. Nothing here is logged or persisted. | |
| """ | |
| from __future__ import annotations | |
| import hashlib | |
| import secrets | |
| import threading | |
| import time | |
| class RateLimiter: | |
| def __init__(self, rate_per_minute: float = 60, burst: int = 30, rotate_seconds: float = 3600): | |
| self.rate = rate_per_minute / 60.0 | |
| self.burst = burst | |
| self.rotate_seconds = rotate_seconds | |
| self._salt = secrets.token_bytes(16) | |
| self._salt_at = time.time() | |
| self._buckets: dict[str, tuple[float, float]] = {} | |
| self._lock = threading.Lock() | |
| def _key(self, client: str) -> str: | |
| if time.time() - self._salt_at > self.rotate_seconds: | |
| self._salt = secrets.token_bytes(16) | |
| self._salt_at = time.time() | |
| self._buckets.clear() | |
| return hashlib.blake2b(client.encode(), key=self._salt, digest_size=12).hexdigest() | |
| def allow(self, client: str, cost: float = 1.0) -> bool: | |
| now = time.time() | |
| with self._lock: | |
| key = self._key(client) | |
| tokens, last = self._buckets.get(key, (float(self.burst), now)) | |
| tokens = min(self.burst, tokens + (now - last) * self.rate) | |
| if tokens < cost: | |
| self._buckets[key] = (tokens, now) | |
| return False | |
| self._buckets[key] = (tokens - cost, now) | |
| if len(self._buckets) > 50_000: | |
| self._buckets.clear() | |
| return True | |