75 lines
3.1 KiB
Python
75 lines
3.1 KiB
Python
"""
|
||
Rate-limiter на основе токен-бакета для внешних API-вызовов.
|
||
|
||
Каждый API получает свой RateLimiter с настройками rate (токенов/сек)
|
||
и burst (максимальный размер бакета). Перед каждым запросом вызывается
|
||
async acquire(), который ждёт освобождения токена при необходимости.
|
||
"""
|
||
|
||
import asyncio
|
||
import logging
|
||
import os
|
||
import time
|
||
from typing import Final
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class RateLimiter:
|
||
"""Токен-бакет: заполняется со скоростью rate токенов/сек, максимум burst."""
|
||
|
||
def __init__(self, rate: float, burst: int) -> None:
|
||
"""
|
||
Args:
|
||
rate: Скорость пополнения токенов (токенов в секунду).
|
||
burst: Максимальный размер бакета.
|
||
"""
|
||
self.rate: float = rate
|
||
self.burst: int = burst
|
||
self.tokens: float = float(burst)
|
||
self.lock: asyncio.Lock = asyncio.Lock()
|
||
self._last_refill: float = time.monotonic()
|
||
|
||
def _refill(self) -> None:
|
||
"""Пополнить токены за прошедшее время."""
|
||
now: float = time.monotonic()
|
||
elapsed: float = now - self._last_refill
|
||
self.tokens = min(self.burst, self.tokens + elapsed * self.rate)
|
||
self._last_refill = now
|
||
|
||
async def acquire(self, token: int = 1) -> None:
|
||
"""Забрать token токенов, ждать если их нет."""
|
||
while True:
|
||
async with self.lock:
|
||
self._refill()
|
||
if self.tokens >= token:
|
||
self.tokens -= token
|
||
return
|
||
# Ждём достаточно времени для восстановления 1 токена
|
||
await asyncio.sleep(token / self.rate)
|
||
|
||
|
||
# --- Готовые лимитеры по API ---
|
||
|
||
# TheCatAPI: бесплатно, 1 req/sec, burst 3
|
||
_CAT_RATE: Final[float] = float(os.getenv("CAT_API_RATE", "1"))
|
||
_CAT_BURST: Final[int] = int(os.getenv("CAT_API_BURST", "3"))
|
||
|
||
# wttr.in: без ключа, 1 req/sec, burst 3
|
||
_WEATHER_RATE: Final[float] = float(os.getenv("WEATHER_API_RATE", "1"))
|
||
_WEATHER_BURST: Final[int] = int(os.getenv("WEATHER_API_BURST", "3"))
|
||
|
||
# Open-Meteo: fallback, 2 req/sec, burst 5
|
||
_OPEN_METEO_RATE: Final[float] = float(os.getenv("OPEN_METEO_API_RATE", "2"))
|
||
_OPEN_METEO_BURST: Final[int] = int(os.getenv("OPEN_METEO_API_BURST", "5"))
|
||
|
||
# Habr RSS: 1 req/sec, burst 2
|
||
_HABR_RSS_RATE: Final[float] = float(os.getenv("HABR_RSS_RATE", "1"))
|
||
_HABR_RSS_BURST: Final[int] = int(os.getenv("HABR_RSS_BURST", "2"))
|
||
|
||
# Экземпляры лимитеров
|
||
cat_limiter: RateLimiter = RateLimiter(_CAT_RATE, _CAT_BURST)
|
||
weather_limiter: RateLimiter = RateLimiter(_WEATHER_RATE, _WEATHER_BURST)
|
||
open_meteo_limiter: RateLimiter = RateLimiter(_OPEN_METEO_RATE, _OPEN_METEO_BURST)
|
||
habr_rss_limiter: RateLimiter = RateLimiter(_HABR_RSS_RATE, _HABR_RSS_BURST)
|