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