Spaces:
Running
Running
File size: 3,995 Bytes
f913c73 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 | """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
@dataclass
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
|