rate_limit.py

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