Kris
1f78c44b03
这是一个批量整理提交，包含以下主要改动： 1. 删除多处冗余的空行和未使用的导入 2. 修复文件末尾缺少换行符的问题 3. 调整部分模块的导入顺序与代码排版 4. 修复部分配置默认值与策略逻辑 5. 新增多个功能模块与辅助工具 6. 完善异常处理与日志记录 7. 修复速率限制、消息缓存、权限校验等逻辑bug 8. 废弃部分旧有API与配置项并添加警告提示
117 lines
3.8 KiB
Python
117 lines
3.8 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
|
|
from yuxi.channels.models import ChannelMessage
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class MessageQueue:
|
|
def __init__(
|
|
self,
|
|
global_limit: int = 1000,
|
|
per_user_limit: int = 20,
|
|
per_group_limit: int = 50,
|
|
max_concurrency: int = 10,
|
|
):
|
|
self._queue: asyncio.Queue[ChannelMessage] = asyncio.Queue(maxsize=global_limit)
|
|
self._global_limit = global_limit
|
|
self._per_user_limit = per_user_limit
|
|
self._per_group_limit = per_group_limit
|
|
self._max_concurrency = max_concurrency
|
|
|
|
self._per_user_counts: dict[str, int] = {}
|
|
self._per_group_counts: dict[str, int] = {}
|
|
|
|
self._semaphore = asyncio.Semaphore(max_concurrency)
|
|
self._lock = asyncio.Lock()
|
|
|
|
async def enqueue(self, msg: ChannelMessage) -> bool:
|
|
user_id = msg.identity.channel_user_id
|
|
chat_id = msg.identity.channel_chat_id
|
|
|
|
async with self._lock:
|
|
if self._queue.qsize() >= self._global_limit:
|
|
try:
|
|
self._queue.get_nowait()
|
|
self._queue.task_done()
|
|
except asyncio.QueueEmpty:
|
|
pass
|
|
|
|
if user_id:
|
|
user_count = self._per_user_counts.get(user_id, 0)
|
|
if user_count >= self._per_user_limit:
|
|
logger.debug("MessageQueue: per-user limit reached for %s", user_id)
|
|
return False
|
|
|
|
if chat_id and chat_id.startswith("group_"):
|
|
group_count = self._per_group_counts.get(chat_id, 0)
|
|
if group_count >= self._per_group_limit:
|
|
logger.debug("MessageQueue: per-group limit reached for %s", chat_id)
|
|
return False
|
|
|
|
try:
|
|
self._queue.put_nowait(msg)
|
|
if user_id:
|
|
self._per_user_counts[user_id] = self._per_user_counts.get(user_id, 0) + 1
|
|
if chat_id and chat_id.startswith("group_"):
|
|
self._per_group_counts[chat_id] = self._per_group_counts.get(chat_id, 0) + 1
|
|
return True
|
|
except asyncio.QueueFull:
|
|
return False
|
|
|
|
async def dequeue(self) -> ChannelMessage:
|
|
msg = await self._queue.get()
|
|
user_id = msg.identity.channel_user_id
|
|
chat_id = msg.identity.channel_chat_id
|
|
|
|
async with self._lock:
|
|
if user_id:
|
|
self._per_user_counts[user_id] = max(0, self._per_user_counts.get(user_id, 1) - 1)
|
|
if chat_id and chat_id.startswith("group_"):
|
|
self._per_group_counts[chat_id] = max(0, self._per_group_counts.get(chat_id, 1) - 1)
|
|
|
|
return msg
|
|
|
|
def task_done(self) -> None:
|
|
self._queue.task_done()
|
|
|
|
@property
|
|
def size(self) -> int:
|
|
return self._queue.qsize()
|
|
|
|
@property
|
|
def snapshot(self) -> dict:
|
|
return {
|
|
"queue_size": self._queue.qsize(),
|
|
"global_limit": self._global_limit,
|
|
"per_user_counts": dict(self._per_user_counts),
|
|
"per_group_counts": dict(self._per_group_counts),
|
|
}
|
|
|
|
async def acquire(self) -> bool:
|
|
try:
|
|
await self._semaphore.acquire()
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
def release(self) -> None:
|
|
try:
|
|
self._semaphore.release()
|
|
except ValueError:
|
|
pass
|
|
|
|
async def drain(self) -> list[ChannelMessage]:
|
|
messages: list[ChannelMessage] = []
|
|
async with self._lock:
|
|
while not self._queue.empty():
|
|
try:
|
|
messages.append(self._queue.get_nowait())
|
|
except asyncio.QueueEmpty:
|
|
break
|
|
self._per_user_counts.clear()
|
|
self._per_group_counts.clear()
|
|
return messages |