更新:

This commit is contained in:
Your Name
2026-07-17 10:21:18 +08:00
parent 530e7f839d
commit 2f2150d2f6
6 changed files with 596 additions and 30 deletions
+117
View File
@@ -0,0 +1,117 @@
"""全局出站发送限流器。
目的:托管账号数量多(如 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