Spaces:
Running
Running
| """FIFO job queue with a small pool of workers.""" | |
| from __future__ import annotations | |
| import asyncio | |
| import itertools | |
| import logging | |
| import time | |
| from dataclasses import dataclass, field | |
| from typing import Any, Awaitable, Callable, Dict, List, Optional | |
| log = logging.getLogger(__name__) | |
| class Job: | |
| id: int | |
| user_id: int | |
| chat_id: int | |
| msg_id: int | |
| file_name: str | |
| size: int | |
| new_name: str | |
| fmt: Optional[Dict[str, Any]] = None # the saved format (auto or manual) this job uses, a snapshot taken when it was queued; None: no auto format | |
| notes: List[str] = field(default_factory=list) # warnings shown with the queued / finished message | |
| order: int = 0 # id of the message that triggered the job: files are processed oldest message first | |
| status_msg: Any = None | |
| state: str = "queued" # queued | running | done | failed | cancelled | |
| cancel: asyncio.Event = field(default_factory=asyncio.Event) | |
| created: float = field(default_factory=time.time) | |
| started: float = 0.0 | |
| Runner = Callable[[Job], Awaitable[None]] | |
| SETTLE_CAP = 4 # an idle worker never waits longer than this many settle periods in a row | |
| def _order_key(job: Job) -> tuple: | |
| return (job.order, job.id) | |
| class JobQueue: | |
| """Jobs run oldest-message-first by a small pool of workers. | |
| Telethon handles messages that arrive together concurrently, so jobs are not necessarily added in the | |
| order the files were sent. Picking by message id (not by insertion order) and letting an idle worker wait | |
| until nothing new was queued for ``settle`` seconds (at most ``SETTLE_CAP`` times that) keeps a forwarded | |
| batch in its original order. | |
| """ | |
| def __init__(self, runner: Runner, workers: int = 1, on_cancelled: Optional[Runner] = None, settle: float = 0.0): | |
| self._runner = runner | |
| self._on_cancelled = on_cancelled | |
| self._workers = max(1, workers) | |
| self._settle = max(0.0, settle) | |
| self._wake = asyncio.Event() | |
| self._ids = itertools.count(1) | |
| self._tasks: List[asyncio.Task] = [] | |
| self.jobs: Dict[int, Job] = {} | |
| def new_id(self) -> int: | |
| return next(self._ids) | |
| def start(self) -> None: | |
| for _ in range(self._workers): | |
| self._tasks.append(asyncio.ensure_future(self._worker())) | |
| async def stop(self) -> None: | |
| for job in list(self.jobs.values()): | |
| job.cancel.set() | |
| for task in self._tasks: | |
| task.cancel() | |
| await asyncio.gather(*self._tasks, return_exceptions=True) | |
| self._tasks.clear() | |
| def add(self, job: Job) -> int: | |
| """Queue a job; returns how many jobs are ahead of it (running + waiting).""" | |
| ahead = self.running + self.waiting | |
| self.jobs[job.id] = job | |
| self._wake.set() | |
| return ahead | |
| def running(self) -> int: | |
| return sum(1 for j in self.jobs.values() if j.state == "running") | |
| def waiting(self) -> int: | |
| return sum(1 for j in self.jobs.values() if j.state == "queued") | |
| def position(self, job: Job) -> int: | |
| """0 = running now, 1 = next up, ...""" | |
| if job.state == "running": | |
| return 0 | |
| waiting = sorted((j for j in self.jobs.values() if j.state == "queued"), key=_order_key) | |
| try: | |
| return waiting.index(job) + 1 | |
| except ValueError: | |
| return 0 | |
| def for_user(self, user_id: int) -> List[Job]: | |
| return sorted((j for j in self.jobs.values() if j.user_id == user_id), key=_order_key) | |
| def cancel(self, job_id: int, user_id: Optional[int] = None) -> bool: | |
| job = self.jobs.get(job_id) | |
| if job is None or (user_id is not None and job.user_id != user_id): | |
| return False | |
| return self._cancel_job(job) | |
| def cancel_user(self, user_id: int) -> int: | |
| return sum(self._cancel_job(j) for j in self.for_user(user_id)) | |
| def cancel_all(self) -> int: | |
| return sum(self._cancel_job(j) for j in list(self.jobs.values())) | |
| def _cancel_job(self, job: Job) -> bool: | |
| if job.state in ("done", "failed", "cancelled") or job.cancel.is_set(): | |
| return False # finished, or a cancel is already on its way | |
| job.cancel.set() | |
| if job.state == "queued": | |
| job.state = "cancelled" | |
| self.jobs.pop(job.id, None) | |
| if self._on_cancelled is not None: | |
| asyncio.ensure_future(self._safe(self._on_cancelled, job)) | |
| return True | |
| async def _safe(self, fn: Runner, job: Job) -> None: | |
| try: | |
| await fn(job) | |
| except Exception: # noqa: BLE001 | |
| log.exception("job callback failed") | |
| def _claim(self) -> Optional[Job]: | |
| """Take the oldest queued job (no await in between, so two workers never take the same one).""" | |
| queued = [j for j in self.jobs.values() if j.state == "queued"] | |
| if not queued: | |
| return None | |
| job = min(queued, key=_order_key) | |
| job.state = "running" | |
| job.started = time.time() | |
| return job | |
| async def _settle_batch(self) -> None: | |
| """The rest of a forwarded batch may still be on its way: wait until nothing new was queued for a while.""" | |
| if not self._settle: | |
| return | |
| deadline = time.monotonic() + self._settle * SETTLE_CAP | |
| while time.monotonic() < deadline: | |
| self._wake.clear() | |
| try: | |
| await asyncio.wait_for(self._wake.wait(), timeout=self._settle) | |
| except asyncio.TimeoutError: | |
| return | |
| async def _next(self) -> Job: | |
| while True: | |
| job = self._claim() | |
| if job is not None: | |
| return job | |
| self._wake.clear() | |
| await self._wake.wait() | |
| await self._settle_batch() | |
| async def _worker(self) -> None: | |
| while True: | |
| job = await self._next() | |
| try: | |
| await self._runner(job) | |
| except asyncio.CancelledError: | |
| raise | |
| except Exception: # noqa: BLE001 - the runner reports its own failures | |
| log.exception("job %s crashed", job.id) | |
| job.state = "failed" | |
| finally: | |
| if job.state == "running": | |
| job.state = "done" | |
| self.jobs.pop(job.id, None) | |