blackbird/pymsg/pymsg/async_queue.py

"""
AsyncRpcQueue - 通用异步批量发送队列

提供内存队列 + 后台线程的异步批量处理机制，适用于日志上报、指标采集等
非关键路径场景。业务线程put操作非阻塞，后台线程攒批处理，失败时通过错误回调
通知，不影响主业务流程。

两种使用模式：

1. 简单RPC模式（指定client + method + serialize）：
    queue = AsyncRpcQueue(
        client=rpc_client,
        method='log_api.push',
        batch_size=50,
        flush_interval=2.0,
        serialize=lambda items: {'logs': json.dumps(items)},
    )

2. 自定义处理模式（指定handler回调）：
    def send_batch(batch):
        for item in batch:
            rpc_client.call('execute', item)
    queue = AsyncRpcQueue(
        handler=send_batch,
        batch_size=50,
        flush_interval=2.0,
    )
"""

import threading
import time
from collections import deque
from typing import Any, Callable, Optional


class AsyncRpcQueue:
    """通用异步批量发送队列"""

    def __init__(
        self,
        client: Any = None,
        method: str = None,
        *,
        handler: Optional[Callable[[list], None]] = None,
        batch_size: int = 100,
        flush_interval: float = 2.0,
        max_queue_size: int = 10000,
        serialize: Optional[Callable[[list], dict]] = None,
        on_error: Optional[Callable[[Exception, list], None]] = None,
        retry_on_refused: bool = False,
        retry_interval: float = 1.0,
    ):
        """
        初始化异步队列

        Args:
            client: RpcClient实例（RPC模式必填）
            method: RPC方法名（RPC模式必填）
            handler: 自定义批量处理回调，接收一批items列表。提供此参数时
                     忽略client/method/serialize/retry_on_refused。
            batch_size: 批量发送的阈值，达到此数量立即发送
            flush_interval: 定时flush间隔（秒）
            max_queue_size: 队列最大长度，超过则丢弃最旧数据
            serialize: 序列化函数，接收items列表返回RPC参数字典
                       默认: lambda items: {'items': items}
            on_error: 错误回调，接收(exception, batch_items)
            retry_on_refused: 是否在连接被拒绝时重试（RPC模式，透传给client.call）
            retry_interval: 重试间隔（秒）（RPC模式，透传给client.call）
        """
        if handler is None:
            if client is None or method is None:
                raise ValueError("必须提供 handler 或 (client + method)")
            self._handler = self._make_rpc_handler(
                client, method,
                serialize or (lambda items: {'items': items}),
                retry_on_refused, retry_interval,
            )
        else:
            self._handler = handler

        self._batch_size = batch_size
        self._flush_interval = flush_interval
        self._max_queue_size = max_queue_size
        self._on_error = on_error

        self._queue: deque = deque()
        self._lock = threading.Lock()
        self._wakeup = threading.Event()
        self._thread: Optional[threading.Thread] = None
        self._closed = False
        self._dropped_count = 0

    @staticmethod
    def _make_rpc_handler(client, method, serialize, retry_on_refused, retry_interval):
        """创建RPC模式的默认handler"""
        def handler(batch: list):
            params = serialize(batch)
            client.call(
                method, params,
                retry_on_refused=retry_on_refused,
                retry_interval=retry_interval,
            )
        return handler

    def start(self) -> None:
        """启动后台发送线程"""
        if self._thread is not None and self._thread.is_alive():
            return
        self._closed = False
        self._thread = threading.Thread(
            target=self._run,
            name=f'AsyncRpcQueue',
            daemon=True,
        )
        self._thread.start()

    def put(self, item: Any) -> bool:
        """
        放入一条数据到队列（非阻塞）

        Args:
            item: 要发送的数据项

        Returns:
            bool: True表示放入成功，False表示队列已关闭
        """
        if self._closed:
            return False
        with self._lock:
            if len(self._queue) >= self._max_queue_size:
                self._queue.popleft()
                self._dropped_count += 1
            self._queue.append(item)
        self._wakeup.set()
        return True

    def flush(self) -> None:
        """立即触发一次发送（会阻塞直到处理完成，主要用于关闭时）"""
        self._send_batch()

    def close(self) -> None:
        """关闭队列，flush剩余数据后停止后台线程"""
        if self._closed:
            return
        self._closed = True
        self._wakeup.set()
        if self._thread is not None:
            self._thread.join(timeout=5.0)
        self._send_batch()

    @property
    def pending_count(self) -> int:
        """当前队列中待发送的数据量"""
        with self._lock:
            return len(self._queue)

    @property
    def dropped_count(self) -> int:
        """因队列满而丢弃的数据量"""
        return self._dropped_count

    def _run(self) -> None:
        """后台线程主循环"""
        last_flush = time.time()
        while not self._closed:
            self._wakeup.wait(timeout=0.1)
            self._wakeup.clear()

            now = time.time()
            should_flush = False
            with self._lock:
                if len(self._queue) >= self._batch_size:
                    should_flush = True
                elif now - last_flush >= self._flush_interval and self._queue:
                    should_flush = True
            if should_flush:
                self._send_batch()
                last_flush = time.time()

    def _send_batch(self) -> None:
        """取一批数据处理"""
        with self._lock:
            if not self._queue:
                return
            batch = list(self._queue)
            self._queue.clear()
        try:
            self._handler(batch)
        except Exception as e:
            if self._on_error:
                try:
                    self._on_error(e, batch)
                except Exception:
                    pass

    def __enter__(self):
        self.start()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()
        return False