ForcePilot/backend/package/yuxi/channel/message/rate_limiter.py

import asyncio
import logging
import time
from collections.abc import AsyncIterator, Callable
from dataclasses import dataclass

logger = logging.getLogger(__name__)

_global_semaphore = asyncio.Semaphore(50)


def get_global_semaphore() -> asyncio.Semaphore:
    return _global_semaphore

DEFAULT_RATE = 10.0
DEFAULT_CAPACITY = 20
DEFAULT_PER_ACCOUNT_CONCURRENCY = 5
DEFAULT_QUEUE_DEPTH = 20
DEFAULT_MAX_LIMITERS = 10_000


@dataclass(slots=True)
class AcquireResult:
    allowed: bool
    retry_after_sec: float = 0.0

    def __bool__(self) -> bool:
        return self.allowed


class TokenBucket:
    """令牌桶算法 — 控制单账户消息速率。

    参数：
    - rate: 每秒补充的令牌数（默认 10）
    - capacity: 桶容量上限（默认 20，允许短时突发）
    - time_func: 可选时钟注入，便于测试（默认 time.monotonic）
    """

    def __init__(
        self,
        rate: float = DEFAULT_RATE,
        capacity: int = DEFAULT_CAPACITY,
        time_func: Callable[[], float] = time.monotonic,
    ):
        self._rate = rate
        self._capacity = capacity
        self._tokens = float(capacity)
        self._time = time_func
        self._last_refill = self._time()
        self._lock = asyncio.Lock()

    async def acquire(self) -> AcquireResult:
        async with self._lock:
            now = self._time()
            elapsed = now - self._last_refill
            self._tokens = min(self._capacity, self._tokens + elapsed * self._rate)
            self._last_refill = now

            if self._tokens >= 1.0:
                self._tokens -= 1.0
                return AcquireResult(allowed=True)

            if self._rate <= 0.0:
                return AcquireResult(allowed=False, retry_after_sec=float("inf"))

            retry_after = (1.0 - self._tokens) / self._rate
            return AcquireResult(allowed=False, retry_after_sec=retry_after)

    def reset(self) -> None:
        self._tokens = float(self._capacity)
        self._last_refill = self._time()


class ChannelRateLimiter:
    """单渠道账户的限流控制器，组合令牌桶 + 并发信号量。

    线程安全：所有方法在 asyncio 上下文中调用。
    """

    def __init__(
        self,
        rate: float = DEFAULT_RATE,
        per_account_concurrency: int = DEFAULT_PER_ACCOUNT_CONCURRENCY,
        queue_depth: int = DEFAULT_QUEUE_DEPTH,
    ):
        self._rate = rate
        self._per_account_concurrency = per_account_concurrency
        self._bucket = TokenBucket(rate=rate, capacity=DEFAULT_CAPACITY)
        self._semaphore = asyncio.Semaphore(per_account_concurrency)
        self._queue_depth = queue_depth
        self._waiting: int = 0

    async def try_acquire(self) -> AcquireResult:
        """尝试获取消息处理许可。

        返回 AcquireResult(allowed=..., retry_after_sec=...)。
        """
        result = await self._bucket.acquire()
        if result.allowed:
            return result

        if self._waiting >= self._queue_depth:
            logger.warning(
                "Rate limit queue full (depth=%d), rejecting message",
                self._queue_depth,
            )
            return AcquireResult(allowed=False, retry_after_sec=result.retry_after_sec)

        self._waiting += 1
        try:
            await asyncio.sleep(result.retry_after_sec)
            retry = await self._bucket.acquire()
            if not retry.allowed:
                return AcquireResult(allowed=False, retry_after_sec=retry.retry_after_sec)
            return retry
        finally:
            self._waiting -= 1

    async def run_with_limit(self, coro) -> AsyncIterator:
        """在并发限制内执行 Agent 调用。

        先获取账户级信号量，再获取全局信号量。
        """
        async with self._semaphore:
            async with _global_semaphore:
                async for chunk in coro:
                    yield chunk

    def reset(self) -> None:
        self._bucket.reset()

    @property
    def waiting_count(self) -> int:
        return self._waiting

    @property
    def config(self) -> dict:
        return {
            "rate": self._rate,
            "per_account_concurrency": self._per_account_concurrency,
            "queue_depth": self._queue_depth,
        }


class RateLimitManager:
    """限流管理器 — 按账户维度管理限流器实例。

    每个 (channel_type, account_id) 组合共享同一个 ChannelRateLimiter。
    内置自动清理机制防止无限增长。
    """

    def __init__(self, max_limiters: int = DEFAULT_MAX_LIMITERS):
        self._limiters: dict[str, tuple[float, ChannelRateLimiter]] = {}
        self._max_limiters = max_limiters
        self._last_cleanup = time.monotonic()
        self._lock = asyncio.Lock()

    def _key(self, channel_type: str, account_id: str) -> str:
        return f"{channel_type}:{account_id}"

    async def get_limiter(
        self,
        channel_type: str,
        account_id: str,
        rate: float = DEFAULT_RATE,
        per_account_concurrency: int = DEFAULT_PER_ACCOUNT_CONCURRENCY,
    ) -> ChannelRateLimiter:
        async with self._lock:
            self._maybe_cleanup()
            key = self._key(channel_type, account_id)
            if key not in self._limiters:
                self._limiters[key] = (
                    time.monotonic(),
                    ChannelRateLimiter(
                        rate=rate,
                        per_account_concurrency=per_account_concurrency,
                    ),
                )
            else:
                self._limiters[key] = (time.monotonic(), self._limiters[key][1])
            return self._limiters[key][1]

    async def remove_limiter(self, channel_type: str, account_id: str) -> None:
        async with self._lock:
            key = self._key(channel_type, account_id)
            self._limiters.pop(key, None)

    async def clear(self) -> None:
        async with self._lock:
            self._limiters.clear()

    async def get_stats(self) -> dict[str, dict]:
        async with self._lock:
            return {
                key: {
                    "waiting": limiter.waiting_count,
                    **limiter.config,
                }
                for key, (_, limiter) in self._limiters.items()
            }

    def _maybe_cleanup(self) -> None:
        if len(self._limiters) <= self._max_limiters:
            return
        now = time.monotonic()
        if now - self._last_cleanup < 300:
            return
        self._last_cleanup = now
        expired = [key for key, (last_touch, _) in self._limiters.items() if now - last_touch > 3600]
        for key in expired:
            self._limiters.pop(key, None)
        if expired:
            logger.info("Cleaned up %d stale rate limiters", len(expired))


rate_limit_manager = RateLimitManager()