Spaces:
Running on CPU Upgrade
Running on CPU Upgrade
File size: 1,422 Bytes
3530326 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 | 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
|