46 lines
1.5 KiB
Python
46 lines
1.5 KiB
Python
"""In-process bounded sliding-window rate limiter with isolated compound keys."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import time
|
|
from collections import defaultdict, deque
|
|
from threading import Lock
|
|
from uuid import UUID
|
|
|
|
|
|
class RateLimitExceeded(RuntimeError):
|
|
code = "RATE_LIMITED"
|
|
|
|
|
|
class SlidingWindowRateLimiter:
|
|
def __init__(
|
|
self, *, limit: int = 60, window_seconds: float = 60.0, clock: object = time.monotonic
|
|
) -> None:
|
|
if limit < 1 or window_seconds <= 0:
|
|
raise ValueError("rate limit must be positive")
|
|
self.limit = limit
|
|
self.window_seconds = window_seconds
|
|
self.clock = clock
|
|
self._entries: dict[tuple[str, str, str], deque[float]] = defaultdict(deque)
|
|
self._lock = Lock()
|
|
|
|
def allow(self, *, client_id: UUID | str, project: str, tool: str) -> bool:
|
|
key = (str(client_id), project, tool)
|
|
now = self.clock() # type: ignore[operator]
|
|
with self._lock:
|
|
bucket = self._entries[key]
|
|
while bucket and bucket[0] <= now - self.window_seconds:
|
|
bucket.popleft()
|
|
if len(bucket) >= self.limit:
|
|
return False
|
|
bucket.append(now)
|
|
return True
|
|
|
|
def check(self, *, client_id: UUID | str, project: str, tool: str) -> None:
|
|
if not self.allow(client_id=client_id, project=project, tool=tool):
|
|
raise RateLimitExceeded("invocation rate limit exceeded")
|
|
|
|
def clear(self) -> None:
|
|
with self._lock:
|
|
self._entries.clear()
|