103 lines
3.7 KiB
Python
103 lines
3.7 KiB
Python
import asyncio
|
|
import contextlib
|
|
|
|
import asyncpg
|
|
from pyrogram import Client
|
|
from pyrogram.errors import FloodPremiumWait, FloodWait
|
|
|
|
from userbot.modules.jobs import repository
|
|
from userbot.modules.jobs.context import JobContext
|
|
from userbot.modules.jobs.registry import JOB_HANDLERS, JobHandler
|
|
from utils.db.models import Job, JobStatus
|
|
from utils.listen import Listener
|
|
from utils.logging import logger
|
|
|
|
JOBS_CHANGED_CHANNEL = "jobs_changed"
|
|
POLL_INTERVAL_SECONDS = 60.0
|
|
RESTART_SECONDS = 10.0
|
|
|
|
|
|
class JobConsumer:
|
|
def __init__(
|
|
self,
|
|
client: Client | None,
|
|
pool: asyncpg.Pool,
|
|
account_id: int,
|
|
registry: dict[str, JobHandler] | None = None,
|
|
) -> None:
|
|
self.client = client
|
|
self.pool = pool
|
|
self.account_id = account_id
|
|
self.registry = registry if registry is not None else JOB_HANDLERS
|
|
self._wake = asyncio.Event()
|
|
|
|
async def execute(self, job: Job) -> None:
|
|
ctx = JobContext(self.client, self.pool, self.account_id, job)
|
|
handler = self.registry.get(job.kind)
|
|
if handler is None:
|
|
await repository.finish(
|
|
self.pool, ctx.job_id, JobStatus.FAILED, f"no handler for {job.kind!r}"
|
|
)
|
|
return
|
|
try:
|
|
await handler(ctx)
|
|
except (FloodWait, FloodPremiumWait) as exc:
|
|
wait = exc.value if isinstance(exc.value, int) else 1
|
|
await repository.defer_flood_wait(self.pool, ctx.job_id, wait)
|
|
logger.warning(
|
|
f"[yellow]FloodWait {wait}s on job {ctx.job_id} ({job.kind}), "
|
|
f"deferred[/]"
|
|
)
|
|
except Exception as exc:
|
|
logger.exception(f"Job {ctx.job_id} ({job.kind}) failed")
|
|
await repository.finish(self.pool, ctx.job_id, JobStatus.FAILED, str(exc))
|
|
else:
|
|
await repository.finish(self.pool, ctx.job_id, JobStatus.DONE)
|
|
|
|
async def drain(self) -> None:
|
|
while True:
|
|
job = await repository.claim_next(self.pool, self.account_id)
|
|
if job is None:
|
|
return
|
|
await self.execute(job)
|
|
|
|
async def run(self) -> None:
|
|
while True:
|
|
try:
|
|
await self._serve()
|
|
except Exception:
|
|
logger.exception(
|
|
f"[red]Job consumer for account {self.account_id} failed,[/] "
|
|
f"restarting in {RESTART_SECONDS:.0f}s"
|
|
)
|
|
await asyncio.sleep(RESTART_SECONDS)
|
|
|
|
async def _serve(self) -> None:
|
|
requeued = await repository.requeue_running(self.pool, self.account_id)
|
|
if requeued:
|
|
logger.info(f"[yellow]Requeued {requeued} stale running job(s).[/]")
|
|
pruned = await repository.prune_finished(self.pool, self.account_id)
|
|
if pruned:
|
|
logger.info(f"[yellow]Pruned {pruned} finished job(s).[/]")
|
|
listener = Listener(
|
|
f"jobs {self.account_id}",
|
|
{JOBS_CHANGED_CHANNEL: lambda _payload: self._wake.set()},
|
|
self._wake.set,
|
|
)
|
|
await listener.start()
|
|
logger.info(f"[green]Job consumer running for account {self.account_id}.[/]")
|
|
try:
|
|
while True:
|
|
self._wake.clear()
|
|
await self.drain()
|
|
with contextlib.suppress(TimeoutError):
|
|
await asyncio.wait_for(
|
|
self._wake.wait(), timeout=await self._idle_seconds()
|
|
)
|
|
finally:
|
|
await listener.stop()
|
|
|
|
async def _idle_seconds(self) -> float:
|
|
due = await repository.seconds_to_next_due(self.pool, self.account_id)
|
|
return POLL_INTERVAL_SECONDS if due is None else min(due, POLL_INTERVAL_SECONDS)
|