63 lines
2.0 KiB
Python
63 lines
2.0 KiB
Python
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import asyncio
|
|||
|
|
import time
|
|||
|
|
|
|||
|
|
|
|||
|
|
class RateLimiter:
|
|||
|
|
"""Twitch 速率限制令牌桶
|
|||
|
|
|
|||
|
|
默认 20msg/30s (known user)，Mod/VIP 切换至 100msg/30s。
|
|||
|
|
mod 模式 60 秒后自动切回 known 模式。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
MOD_TIMEOUT = 60.0
|
|||
|
|
|
|||
|
|
def __init__(self, limit: int = 20, window: float = 30.0, mod_limit: int = 100):
|
|||
|
|
self._default_limit = limit
|
|||
|
|
self._mod_limit = mod_limit
|
|||
|
|
self._current_limit = limit
|
|||
|
|
self._window = window
|
|||
|
|
self._tokens = float(limit)
|
|||
|
|
self._last_refill = time.monotonic()
|
|||
|
|
self._mod_until: float | None = None
|
|||
|
|
self._lock = asyncio.Lock()
|
|||
|
|
|
|||
|
|
async def switch_to_mod(self) -> None:
|
|||
|
|
async with self._lock:
|
|||
|
|
if self._current_limit == self._mod_limit:
|
|||
|
|
return
|
|||
|
|
self._current_limit = self._mod_limit
|
|||
|
|
self._tokens = float(self._mod_limit)
|
|||
|
|
self._last_refill = time.monotonic()
|
|||
|
|
self._mod_until = time.monotonic() + self.MOD_TIMEOUT
|
|||
|
|
|
|||
|
|
async def switch_to_known(self) -> None:
|
|||
|
|
async with self._lock:
|
|||
|
|
self._current_limit = self._default_limit
|
|||
|
|
self._tokens = float(self._default_limit)
|
|||
|
|
self._last_refill = time.monotonic()
|
|||
|
|
self._mod_until = None
|
|||
|
|
|
|||
|
|
async def acquire(self) -> bool:
|
|||
|
|
async with self._lock:
|
|||
|
|
now = time.monotonic()
|
|||
|
|
|
|||
|
|
if self._mod_until is not None and now >= self._mod_until:
|
|||
|
|
self._current_limit = self._default_limit
|
|||
|
|
self._mod_until = None
|
|||
|
|
|
|||
|
|
elapsed = now - self._last_refill
|
|||
|
|
refill = elapsed * (self._current_limit / self._window)
|
|||
|
|
self._tokens = min(float(self._current_limit), self._tokens + refill)
|
|||
|
|
self._last_refill = now
|
|||
|
|
|
|||
|
|
if self._tokens >= 1.0:
|
|||
|
|
self._tokens -= 1.0
|
|||
|
|
return True
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def available_tokens(self) -> float:
|
|||
|
|
return self._tokens
|