44 lines
1.8 KiB
Python
44 lines
1.8 KiB
Python
|
"""A tiny in-process sliding-window rate limiter.
|
||
|
|
|
||
|
|
Generic groundwork: keyed by an arbitrary string (e.g. a client IP), so it can throttle any
|
||
|
|
endpoint. Single-worker uvicorn → in-memory state is sufficient; it resets on restart, which
|
||
|
|
is fine for abuse throttling (not for anything that must survive a deploy). Not shared across
|
||
|
|
processes — if we ever run multiple workers, swap the backing store for Redis behind the same
|
||
|
|
``allow()`` interface.
|
||
|
|
"""
|
||
|
|
import threading
|
||
|
|
import time
|
||
|
|
|
||
|
|
|
||
|
|
class RateLimiter:
|
||
|
|
def __init__(self, max_events: int, window_seconds: float):
|
||
|
|
self.max_events = max_events
|
||
|
|
self.window = window_seconds
|
||
|
|
self._hits: dict[str, list[float]] = {}
|
||
|
|
self._lock = threading.Lock()
|
||
|
|
|
||
|
|
def allow(self, key: str) -> bool:
|
||
|
|
"""Record an attempt for ``key`` and return whether it is within the limit.
|
||
|
|
|
||
|
|
True -> under the cap (the attempt is counted).
|
||
|
|
False -> the cap for the current window is already reached (attempt NOT counted, so a
|
||
|
|
blocked caller can't keep pushing the window forward)."""
|
||
|
|
now = time.monotonic()
|
||
|
|
cutoff = now - self.window
|
||
|
|
with self._lock:
|
||
|
|
hits = [t for t in self._hits.get(key, []) if t > cutoff]
|
||
|
|
if len(hits) >= self.max_events:
|
||
|
|
self._hits[key] = hits
|
||
|
|
return False
|
||
|
|
hits.append(now)
|
||
|
|
self._hits[key] = hits
|
||
|
|
# Opportunistic cleanup so abandoned keys don't accumulate unboundedly.
|
||
|
|
if len(self._hits) > 4096:
|
||
|
|
for k in list(self._hits):
|
||
|
|
fresh = [t for t in self._hits[k] if t > cutoff]
|
||
|
|
if fresh:
|
||
|
|
self._hits[k] = fresh
|
||
|
|
else:
|
||
|
|
del self._hits[k]
|
||
|
|
return True
|