110 lines
4.4 KiB
Python
110 lines
4.4 KiB
Python
"""
|
||
Rate-limiter на основе токен-бакета для внешних API-вызовов.
|
||
|
||
Каждый API получает свой RateLimiter с настройками rate (токенов/сек)
|
||
и burst (максимальный размер бакета). Перед каждым запросом вызывается
|
||
async acquire(), который ждёт освобождения токена при необходимости.
|
||
"""
|
||
|
||
import asyncio
|
||
import logging
|
||
import os
|
||
import time
|
||
from typing import Callable, Final
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class RateLimiter:
|
||
"""Токен-бакет: заполняется со скоростью rate токенов/сек, максимум burst."""
|
||
|
||
def __init__(
|
||
self,
|
||
rate: float,
|
||
burst: int,
|
||
_time_func: Callable[[], float] | None = None,
|
||
) -> None:
|
||
"""
|
||
Args:
|
||
rate: Скорость пополнения токенов (токенов в секунду).
|
||
burst: Максимальный размер бакета.
|
||
_time_func: Функция получения времени (для тестов). По умолчанию time.monotonic.
|
||
"""
|
||
self.rate: float = rate
|
||
self.burst: int = burst
|
||
self.tokens: float = float(burst)
|
||
self.lock: asyncio.Lock = asyncio.Lock()
|
||
self._time_func = _time_func or time.monotonic
|
||
self._last_refill: float = self._time_func()
|
||
|
||
def _refill(self) -> None:
|
||
"""Пополнить токены за прошедшее время."""
|
||
now: float = self._time_func()
|
||
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 токенов, ждать если их нет."""
|
||
wait_count = 0
|
||
while True:
|
||
async with self.lock:
|
||
self._refill()
|
||
if self.tokens >= token:
|
||
self.tokens -= token
|
||
if wait_count:
|
||
logger.debug("RateLimiter: ждал %d раз(а)", wait_count)
|
||
return
|
||
wait_count += 1
|
||
# Ждём достаточно времени для восстановления 1 токена
|
||
await asyncio.sleep(token / self.rate)
|
||
|
||
|
||
# --- Конфигурация лимитеров ---
|
||
|
||
# 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"))
|
||
|
||
|
||
def make_cat_limiter() -> RateLimiter:
|
||
"""Создать лимитер для TheCatAPI."""
|
||
return RateLimiter(_CAT_RATE, _CAT_BURST)
|
||
|
||
|
||
def make_weather_limiter() -> RateLimiter:
|
||
"""Создать лимитер для wttr.in."""
|
||
return RateLimiter(_WEATHER_RATE, _WEATHER_BURST)
|
||
|
||
|
||
def make_open_meteo_limiter() -> RateLimiter:
|
||
"""Создать лимитер для Open-Meteo."""
|
||
return RateLimiter(_OPEN_METEO_RATE, _OPEN_METEO_BURST)
|
||
|
||
|
||
def make_habr_rss_limiter() -> RateLimiter:
|
||
"""Создать лимитер для Habr RSS."""
|
||
return RateLimiter(_HABR_RSS_RATE, _HABR_RSS_BURST)
|
||
|
||
|
||
# Глобальные синглтоны для production.
|
||
# Каждый API-модуль импортирует свой лимитер напрямую (один экземпляр на процесс).
|
||
# Factory-функции (make_*_limiter) используются для создания
|
||
# изолированных экземпляров в тестах с контролируемым временем.
|
||
cat_limiter: RateLimiter = make_cat_limiter()
|
||
weather_limiter: RateLimiter = make_weather_limiter()
|
||
open_meteo_limiter: RateLimiter = make_open_meteo_limiter()
|
||
habr_rss_limiter: RateLimiter = make_habr_rss_limiter()
|