First product

app/auth/rate_limit.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+from collections import defaultdict, deque
+from collections.abc import Callable
+from dataclasses import dataclass
+from math import ceil
+from threading import Lock
+from time import monotonic
+
+
+@dataclass(frozen=True)
+class RateLimitDecision:
+    allowed: bool
+    retry_after_seconds: int = 0
+
+
+class SlidingWindowRateLimiter:
+    def __init__(
+        self,
+        *,
+        max_attempts: int,
+        window_seconds: int,
+        now: Callable[[], float] = monotonic,
+    ) -> None:
+        self._max_attempts = max_attempts
+        self._window_seconds = window_seconds
+        self._now = now
+        self._attempts: dict[str, deque[float]] = defaultdict(deque)
+        self._lock = Lock()
+
+    def check(self, key: str) -> RateLimitDecision:
+        with self._lock:
+            now_value = self._now()
+            attempts = self._attempts[key]
+            self._trim_expired(attempts, now_value)
+            if len(attempts) >= self._max_attempts:
+                retry_after = self._window_seconds - (now_value - attempts[0])
+                return RateLimitDecision(
+                    allowed=False,
+                    retry_after_seconds=max(1, ceil(retry_after)),
+                )
+            return RateLimitDecision(allowed=True)
+
+    def record_failure(self, key: str) -> None:
+        with self._lock:
+            now_value = self._now()
+            attempts = self._attempts[key]
+            self._trim_expired(attempts, now_value)
+            attempts.append(now_value)
+
+    def reset(self, key: str) -> None:
+        with self._lock:
+            self._attempts.pop(key, None)
+
+    def clear_all(self) -> None:
+        with self._lock:
+            self._attempts.clear()
+
+    def _trim_expired(self, attempts: deque[float], now_value: float) -> None:
+        while attempts and now_value - attempts[0] >= self._window_seconds:
+            attempts.popleft()
+