118 lines
4.2 KiB
Python
118 lines
4.2 KiB
Python
"""全局出站发送限流器。
|
||
|
||
目的：托管账号数量多（如 1000 个）时，避免自动回复/手动发送在同一时刻
|
||
大量并发打向抖音接口，占满服务器出站带宽或触发平台级风控。
|
||
|
||
机制（两层，先到先过）：
|
||
1. 并发上限（Semaphore）：同时在途的发送请求数不超过 KEFU_SEND_MAX_CONCURRENCY；
|
||
2. 令牌桶（QPS）：平均发送速率不超过 KEFU_SEND_RATE_PER_SEC，允许 KEFU_SEND_BURST 的突发。
|
||
|
||
限流器按事件循环惰性创建（backend 主循环一个实例；独立脚本各自一个），
|
||
所有账号共享同一个实例，因此是跨账号的全局闸门。
|
||
|
||
环境变量：
|
||
KEFU_SEND_MAX_CONCURRENCY 同时在途发送数上限，默认 20，<=0 表示不限
|
||
KEFU_SEND_RATE_PER_SEC 平均每秒发送数上限，默认 10，<=0 表示不限
|
||
KEFU_SEND_BURST 令牌桶容量（突发上限），默认与并发上限相同
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import logging
|
||
import os
|
||
import time
|
||
|
||
logger = logging.getLogger("douyin_im.rate_limit")
|
||
|
||
|
||
def _env_float(name: str, default: float) -> float:
|
||
try:
|
||
return float(os.getenv(name, "") or default)
|
||
except ValueError:
|
||
return default
|
||
|
||
|
||
class GlobalSendLimiter:
|
||
"""并发上限 + 令牌桶。用法：async with limiter: await 发送。"""
|
||
|
||
def __init__(
|
||
self,
|
||
max_concurrency: int = 20,
|
||
rate_per_sec: float = 10.0,
|
||
burst: float | None = None,
|
||
):
|
||
self.max_concurrency = int(max_concurrency)
|
||
self.rate_per_sec = float(rate_per_sec)
|
||
self.burst = float(burst if burst is not None else max(1, max_concurrency))
|
||
self._sem = (
|
||
asyncio.Semaphore(self.max_concurrency) if self.max_concurrency > 0 else None
|
||
)
|
||
self._tokens = self.burst
|
||
self._last_refill = time.monotonic()
|
||
self._token_lock = asyncio.Lock()
|
||
# 统计：便于日志观察限流是否生效
|
||
self._waited_total = 0.0
|
||
self._acquired_count = 0
|
||
|
||
async def _take_token(self) -> None:
|
||
if self.rate_per_sec <= 0:
|
||
return
|
||
async with self._token_lock:
|
||
while True:
|
||
now = time.monotonic()
|
||
self._tokens = min(
|
||
self.burst, self._tokens + (now - self._last_refill) * self.rate_per_sec
|
||
)
|
||
self._last_refill = now
|
||
if self._tokens >= 1.0:
|
||
self._tokens -= 1.0
|
||
return
|
||
await asyncio.sleep((1.0 - self._tokens) / self.rate_per_sec)
|
||
|
||
async def __aenter__(self) -> "GlobalSendLimiter":
|
||
start = time.monotonic()
|
||
if self._sem is not None:
|
||
await self._sem.acquire()
|
||
try:
|
||
await self._take_token()
|
||
except BaseException:
|
||
if self._sem is not None:
|
||
self._sem.release()
|
||
raise
|
||
waited = time.monotonic() - start
|
||
self._waited_total += waited
|
||
self._acquired_count += 1
|
||
if waited > 1.0:
|
||
logger.info(
|
||
"send throttled: waited %.1fs (in-flight cap=%s, rate=%s/s)",
|
||
waited, self.max_concurrency or "∞", self.rate_per_sec or "∞",
|
||
)
|
||
return self
|
||
|
||
async def __aexit__(self, *exc) -> None:
|
||
if self._sem is not None:
|
||
self._sem.release()
|
||
|
||
|
||
# 每个事件循环一个实例（backend 主循环即全局唯一；脚本自用循环互不影响）
|
||
_limiters: dict[int, GlobalSendLimiter] = {}
|
||
|
||
|
||
def get_send_limiter() -> GlobalSendLimiter:
|
||
loop = asyncio.get_running_loop()
|
||
key = id(loop)
|
||
limiter = _limiters.get(key)
|
||
if limiter is None:
|
||
max_conc = int(_env_float("KEFU_SEND_MAX_CONCURRENCY", 20))
|
||
rate = _env_float("KEFU_SEND_RATE_PER_SEC", 10.0)
|
||
burst = _env_float("KEFU_SEND_BURST", max(1, max_conc))
|
||
limiter = GlobalSendLimiter(max_conc, rate, burst)
|
||
_limiters[key] = limiter
|
||
logger.info(
|
||
"global send limiter ready: concurrency=%s, rate=%s/s, burst=%s",
|
||
max_conc if max_conc > 0 else "unlimited",
|
||
rate if rate > 0 else "unlimited",
|
||
burst,
|
||
)
|
||
return limiter
|