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
|