46 lines
1.5 KiB
Python
46 lines
1.5 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import time
|
|
|
|
from yuxi.utils.logging_config import logger
|
|
|
|
|
|
class RateLimiter:
|
|
def __init__(self, limit: int = 30, window: float = 60.0):
|
|
self._limit = limit
|
|
self._window = window
|
|
self._tokens = float(limit)
|
|
self._last_refill = time.monotonic()
|
|
self._lock = asyncio.Lock()
|
|
|
|
async def acquire(self) -> bool:
|
|
async with self._lock:
|
|
now = time.monotonic()
|
|
elapsed = now - self._last_refill
|
|
refill = elapsed * (self._limit / self._window)
|
|
self._tokens = min(float(self._limit), self._tokens + refill)
|
|
self._last_refill = now
|
|
|
|
if self._tokens >= 1.0:
|
|
self._tokens -= 1.0
|
|
return True
|
|
|
|
wait_s = (1.0 - self._tokens) * (self._window / self._limit)
|
|
logger.debug(f"[Urbit] Rate limit: token exhausted, need ~{wait_s:.1f}s")
|
|
return False
|
|
|
|
async def wait_and_acquire(self, timeout: float = 30.0) -> bool:
|
|
deadline = time.monotonic() + timeout
|
|
while time.monotonic() < deadline:
|
|
if await self.acquire():
|
|
return True
|
|
wait_s = (1.0 - max(0, self._tokens)) * (self._window / self._limit)
|
|
await asyncio.sleep(min(wait_s, 1.0))
|
|
logger.warning("[Urbit] Rate limiter wait timed out")
|
|
return False
|
|
|
|
@property
|
|
def available_tokens(self) -> float:
|
|
return max(0.0, self._tokens)
|