Spaces:
Running
Running
| """Async task queue for file processing — multiple workers, producer-consumer pattern.""" | |
| from __future__ import annotations | |
| import asyncio | |
| from dataclasses import dataclass, field | |
| from typing import Optional | |
| from loguru import logger | |
| class QueueItem: | |
| """A single file-processing request with optional priority.""" | |
| user_id: int | |
| file_id: str | |
| file_name: str | |
| caption: str | |
| file_size: int | |
| priority: int = 0 | |
| submitted_at: float = field(default_factory=asyncio.get_event_loop().time) | |
| class UploadQueue: | |
| """Bounded async FIFO queue with multiple concurrent workers.""" | |
| def __init__(self, maxsize: int = 5, num_workers: int = 2) -> None: | |
| self._queue: asyncio.Queue[QueueItem] = asyncio.Queue(maxsize=maxsize) | |
| self._workers: list[asyncio.Task[None]] = [] | |
| self._active: bool = False | |
| self._shutdown: bool = False | |
| self._num_workers = num_workers | |
| async def submit( | |
| self, | |
| user_id: int, | |
| file_id: str, | |
| file_name: str, | |
| caption: str, | |
| file_size: int, | |
| priority: int = 0, | |
| ) -> None: | |
| """Add a task to the queue. Raises ``asyncio.QueueFull`` if full.""" | |
| if self._shutdown: | |
| raise RuntimeError("Queue is shut down — cannot accept new tasks") | |
| item = QueueItem( | |
| user_id=user_id, | |
| file_id=file_id, | |
| file_name=file_name, | |
| caption=caption, | |
| file_size=file_size, | |
| priority=priority, | |
| ) | |
| await self._queue.put(item) | |
| logger.info(f"Queued task for user {user_id} ({file_name})") | |
| self._ensure_workers() | |
| async def shutdown(self) -> None: | |
| """Cancel pending workers gracefully and stop accepting tasks.""" | |
| self._shutdown = True | |
| for w in self._workers: | |
| w.cancel() | |
| await asyncio.gather(*self._workers, return_exceptions=True) | |
| self._active = False | |
| def status(self) -> dict: | |
| """Return current queue state.""" | |
| return { | |
| "queued": self._queue.qsize(), | |
| "active": self._active, | |
| "maxsize": self._queue.maxsize, | |
| "shutdown": self._shutdown, | |
| "workers": len(self._workers), | |
| } | |
| def _ensure_workers(self) -> None: | |
| """Start workers until we reach ``_num_workers``.""" | |
| target = self._num_workers | |
| while len(self._workers) < target: | |
| worker = asyncio.create_task(self._worker_loop()) | |
| worker.add_done_callback(self._on_worker_done) | |
| self._workers.append(worker) | |
| def _on_worker_done(self, _task: asyncio.Task[None]) -> None: | |
| """Clean up finished workers and restart if work remains.""" | |
| self._workers = [w for w in self._workers if not w.done()] | |
| if not self._queue.empty() and not self._shutdown: | |
| self._ensure_workers() | |
| async def _worker_loop(self) -> None: | |
| """Worker coroutine — pull items from the queue and process them.""" | |
| from bot.app import get_client | |
| from services.pipeline import Pipeline | |
| client = get_client() | |
| self._active = True | |
| while not self._shutdown: | |
| try: | |
| item = await asyncio.wait_for(self._queue.get(), timeout=30) | |
| except asyncio.TimeoutError: | |
| continue | |
| except asyncio.CancelledError: | |
| break | |
| try: | |
| await Pipeline(client, item.user_id).run() | |
| except Exception: | |
| logger.exception(f"Pipeline failed for user {item.user_id}") | |
| finally: | |
| self._queue.task_done() | |
| _upload_queue: Optional[UploadQueue] = None | |
| def get_queue() -> UploadQueue: | |
| global _upload_queue | |
| if _upload_queue is None: | |
| from config import settings | |
| _upload_queue = UploadQueue( | |
| maxsize=5, | |
| num_workers=settings.num_workers, | |
| ) | |
| return _upload_queue | |