mi
38 lines
1.1 KiB
Python
38 lines
1.1 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import time
|
|
from dataclasses import dataclass
|
|
|
|
|
|
@dataclass
|
|
class Bucket:
|
|
tokens: float
|
|
updated: float
|
|
|
|
|
|
class RateLimited(RuntimeError):
|
|
def __init__(self, retry_after: int = 1) -> None:
|
|
self.retry_after = retry_after
|
|
|
|
|
|
class ConservativeRateLimiter:
|
|
"""Process-local fallback. Redis may accelerate this, never own correctness."""
|
|
|
|
def __init__(self, text_rps: int, file_rps: int) -> None:
|
|
self.rates = {"text": text_rps, "file": file_rps}
|
|
now = time.monotonic()
|
|
self.buckets = {kind: Bucket(float(rate), now) for kind, rate in self.rates.items()}
|
|
self.lock = asyncio.Lock()
|
|
|
|
async def acquire(self, kind: str) -> None:
|
|
async with self.lock:
|
|
now = time.monotonic()
|
|
bucket = self.buckets[kind]
|
|
rate = self.rates[kind]
|
|
bucket.tokens = min(float(rate), bucket.tokens + (now - bucket.updated) * rate)
|
|
bucket.updated = now
|
|
if bucket.tokens < 1:
|
|
raise RateLimited
|
|
bucket.tokens -= 1
|