from __future__ import annotations import threading import time from dataclasses import dataclass @dataclass class Bucket: count: int reset_at: float class RateLimiter: def __init__( self, *, limit: int = 5, window_seconds: int = 10 * 60, max_buckets: int = 10_000, ) -> None: self.limit = limit self.window_seconds = window_seconds self.max_buckets = max_buckets self._buckets: dict[str, Bucket] = {} self._lock = threading.Lock() def take(self, key: str, *, now: float | None = None) -> tuple[bool, int]: current = time.monotonic() if now is None else now with self._lock: if len(self._buckets) >= self.max_buckets: self._buckets = { bucket_key: bucket for bucket_key, bucket in self._buckets.items() if bucket.reset_at > current } bucket = self._buckets.get(key) if bucket is None or bucket.reset_at <= current: self._buckets[key] = Bucket( count=1, reset_at=current + self.window_seconds, ) return True, 0 if bucket.count >= self.limit: return False, max(1, round(bucket.reset_at - current)) bucket.count += 1 return True, 0