78 lines
2.5 KiB
Python
78 lines
2.5 KiB
Python
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import asyncio
|
|||
|
|
from collections.abc import Awaitable, Callable
|
|||
|
|
from typing import Any
|
|||
|
|
|
|||
|
|
from yuxi.utils.logging_config import logger
|
|||
|
|
|
|||
|
|
|
|||
|
|
class QueueDebounce:
|
|||
|
|
"""消息队列防抖器
|
|||
|
|
|
|||
|
|
同一 session_key 在 debounce_ms 内的多条连续消息合并为一次批量处理。
|
|||
|
|
适用于群聊中同一用户连续发送多条消息的场景。
|
|||
|
|
|
|||
|
|
用法：
|
|||
|
|
debouncer = QueueDebounce(debounce_ms=500)
|
|||
|
|
await debouncer.enqueue(session_key, message, handler)
|
|||
|
|
|
|||
|
|
工作原理：
|
|||
|
|
1. enqueue(msg) → 将消息加入队列，重置定时器
|
|||
|
|
2. debounce_ms 到期后 → 批量调用 handler(msgs)
|
|||
|
|
3. 期间又来新消息 → 重新重置定时器（重新等待 debounce_ms）
|
|||
|
|
|
|||
|
|
性能：使用 asyncio.Task，内存占用 O(活跃会话数)
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(self, debounce_ms: int = 500):
|
|||
|
|
self._debounce_s: float = debounce_ms / 1000.0
|
|||
|
|
self._queues: dict[str, list[Any]] = {}
|
|||
|
|
self._timers: dict[str, asyncio.Task] = {}
|
|||
|
|
self._lock = asyncio.Lock()
|
|||
|
|
|
|||
|
|
async def enqueue(
|
|||
|
|
self,
|
|||
|
|
session_key: str,
|
|||
|
|
message: Any,
|
|||
|
|
handler: Callable[[list[Any]], Awaitable[None]],
|
|||
|
|
) -> None:
|
|||
|
|
async with self._lock:
|
|||
|
|
if session_key not in self._queues:
|
|||
|
|
self._queues[session_key] = []
|
|||
|
|
self._queues[session_key].append(message)
|
|||
|
|
|
|||
|
|
old_timer = self._timers.pop(session_key, None)
|
|||
|
|
if old_timer:
|
|||
|
|
old_timer.cancel()
|
|||
|
|
|
|||
|
|
self._timers[session_key] = asyncio.create_task(self._flush_after_delay(session_key, handler))
|
|||
|
|
|
|||
|
|
async def _flush_after_delay(
|
|||
|
|
self,
|
|||
|
|
session_key: str,
|
|||
|
|
handler: Callable[[list[Any]], Awaitable[None]],
|
|||
|
|
) -> None:
|
|||
|
|
await asyncio.sleep(self._debounce_s)
|
|||
|
|
async with self._lock:
|
|||
|
|
msgs = self._queues.pop(session_key, [])
|
|||
|
|
self._timers.pop(session_key, None)
|
|||
|
|
|
|||
|
|
if msgs:
|
|||
|
|
try:
|
|||
|
|
await handler(msgs)
|
|||
|
|
logger.debug(f"[Debounce] Flushed {len(msgs)} messages for {session_key}")
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"[Debounce] Handler error for {session_key}: {e}")
|
|||
|
|
|
|||
|
|
async def cancel(self, session_key: str) -> None:
|
|||
|
|
async with self._lock:
|
|||
|
|
timer = self._timers.pop(session_key, None)
|
|||
|
|
if timer:
|
|||
|
|
timer.cancel()
|
|||
|
|
self._queues.pop(session_key, None)
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def pending_count(self) -> int:
|
|||
|
|
return len(self._queues)
|