65 lines
2.1 KiB
Python
65 lines
2.1 KiB
Python
"""Sliding-window rate limiter with injectable clock and thread safety.
|
|
|
|
A request is admitted only if fewer than `limit` requests have been recorded
|
|
in the trailing `window` seconds. Timestamps are kept in a deque; expired
|
|
entries are pruned on every call. Concurrency is guarded by a lock.
|
|
|
|
Subtleties handled:
|
|
- Clock injection (monotonic-ish) for deterministic tests.
|
|
- Boundary: a request exactly `window` seconds old is expired (half-open interval).
|
|
- limit <= 0 rejects everything; window <= 0 rejects everything.
|
|
- Thread safety.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import threading
|
|
import time
|
|
from collections import deque
|
|
from typing import Callable, Optional
|
|
|
|
|
|
class RateLimiter:
|
|
def __init__(
|
|
self,
|
|
limit: int,
|
|
window: float,
|
|
clock: Optional[Callable[[], float]] = None,
|
|
) -> None:
|
|
if limit < 0:
|
|
raise ValueError("limit must be >= 0")
|
|
if window < 0:
|
|
raise ValueError("window must be >= 0")
|
|
self.limit = limit
|
|
self.window = window
|
|
self._clock = clock if clock is not None else time.monotonic
|
|
self._timestamps: deque[float] = deque()
|
|
self._lock = threading.Lock()
|
|
|
|
def _now(self) -> float:
|
|
return self._clock()
|
|
|
|
def allow(self) -> bool:
|
|
"""Return True and record the request if admitted."""
|
|
with self._lock:
|
|
now = self._now()
|
|
# prune expired: keep ts where now - ts < window (half-open)
|
|
while self._timestamps and (now - self._timestamps[0]) >= self.window:
|
|
self._timestamps.popleft()
|
|
if self.limit <= 0 or self.window <= 0:
|
|
return False
|
|
if len(self._timestamps) < self.limit:
|
|
self._timestamps.append(now)
|
|
return True
|
|
return False
|
|
|
|
def reset(self) -> None:
|
|
with self._lock:
|
|
self._timestamps.clear()
|
|
|
|
def count(self) -> int:
|
|
with self._lock:
|
|
now = self._now()
|
|
while self._timestamps and (now - self._timestamps[0]) >= self.window:
|
|
self._timestamps.popleft()
|
|
return len(self._timestamps)
|