lark2lark/ratelimit.py

"""限流缓冲队列：单 Worker 匀速消费，保证 < max_qps QPS。

飞书限制：同一群发送消息上限 5 QPS。本模块用 asyncio.Queue 削峰，
独立 Worker 协程以 1/max_qps 秒/条的速率匀速消费，避免瞬时并发触发封禁。
"""
from __future__ import annotations

import asyncio
import logging
from dataclasses import dataclass
from typing import Awaitable, Callable, Optional

logger = logging.getLogger(__name__)


@dataclass
class ForwardTask:
    """待转发的单条任务：源消息 id -> 单个目标群。"""
    message_id: str
    target_chat_id: str
    # 可选：携带源信息用于日志
    source_chat_id: str = ""


class RateLimitedWorker:
    """单消费协程 + asyncio.Queue，匀速向目标群转发。

    - 入队即返回（非阻塞），由 Worker 串行消费
    - 每条间隔 1/max_qps 秒，确保 QPS 不超限
    - 单条失败按 max_retry 退避重试
    """

    def __init__(
        self,
        max_qps: int,
        max_retry: int,
        forward_fn: Optional[Callable[[ForwardTask], Awaitable[bool]]] = None,
        queue: Optional[asyncio.Queue] = None,
    ) -> None:
        if max_qps <= 0:
            raise ValueError("max_qps 必须为正数")
        self._interval = 1.0 / max_qps
        self._max_retry = max_retry
        self._forward_fn = forward_fn
        self._queue: asyncio.Queue = queue or asyncio.Queue()

    def set_forward_fn(self, fn: Callable[[ForwardTask], Awaitable[bool]]) -> None:
        """后置注入转发函数（解决与 Forwarder 的循环依赖）。"""
        self._forward_fn = fn

    async def enqueue(self, task: ForwardTask) -> None:
        await self._queue.put(task)

    async def run(self) -> None:
        logger.info("转发 Worker 启动，目标速率 %.2f QPS（间隔 %.3fs）",
                    1.0 / self._interval, self._interval)
        while True:
            task: ForwardTask = await self._queue.get()
            try:
                await self._process(task)
            except asyncio.CancelledError:
                logger.info("Worker 收到取消信号，退出")
                raise
            except Exception as e:
                logger.exception("Worker 处理异常: %s", e)
            finally:
                self._queue.task_done()

    async def _process(self, task: ForwardTask) -> None:
        if self._forward_fn is None:
            logger.error("forward_fn 未设置，丢弃任务 message_id=%s", task.message_id)
            return
        for attempt in range(self._max_retry + 1):
            ok = await self._forward_fn(task)
            if ok:
                if attempt > 0:
                    logger.info("重试成功 message_id=%s -> %s（第 %d 次）",
                                task.message_id, task.target_chat_id, attempt)
                break
            # 失败：退避后重试
            backoff = self._interval * (attempt + 1) * 2
            logger.warning("转发失败 message_id=%s -> %s，%ds 后重试（%d/%d）",
                           task.message_id, task.target_chat_id,
                           int(backoff), attempt + 1, self._max_retry)
            await asyncio.sleep(backoff)
        else:
            logger.error("转发最终失败 message_id=%s -> %s，已放弃",
                         task.message_id, task.target_chat_id)
        # 成功或失败都遵守速率间隔
        await asyncio.sleep(self._interval)

    async def wait_drained(self, timeout: Optional[float] = None) -> bool:
        """等待队列清空，返回是否在 timeout 内清空。"""
        try:
            await asyncio.wait_for(self._queue.join(), timeout=timeout)
            return True
        except asyncio.TimeoutError:
            return False