Spaces:
Running
Running
File size: 6,384 Bytes
58f00c9 6c3a31d c3c6394 58f00c9 c3c6394 58f00c9 c3c6394 58f00c9 c3c6394 58f00c9 c3c6394 58f00c9 c3c6394 58f00c9 c3c6394 58f00c9 c3c6394 58f00c9 c3c6394 58f00c9 c3c6394 58f00c9 c3c6394 | 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 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 | """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__)
@dataclass
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
@property
def running(self) -> int:
return sum(1 for j in self.jobs.values() if j.state == "running")
@property
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)
|