89 lines
3.2 KiB
Python
89 lines
3.2 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.env import env
|
|
from utils.logging import logger
|
|
|
|
JOBS_CHANGED_CHANNEL = "jobs_changed"
|
|
POLL_INTERVAL_SECONDS = 60.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
|
|
while True:
|
|
try:
|
|
await handler(ctx)
|
|
except (FloodWait, FloodPremiumWait) as exc:
|
|
await repository.bump_flood_wait(self.pool, ctx.job_id)
|
|
wait = exc.value if isinstance(exc.value, int) else 1
|
|
logger.warning(
|
|
f"[yellow]FloodWait {wait}s on job {ctx.job_id} ({job.kind})[/]"
|
|
)
|
|
await asyncio.sleep(wait)
|
|
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)
|
|
)
|
|
return
|
|
else:
|
|
await repository.finish(self.pool, ctx.job_id, JobStatus.DONE)
|
|
return
|
|
|
|
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:
|
|
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).[/]")
|
|
conn = await asyncpg.connect(dsn=env.db.connection_url)
|
|
await conn.add_listener(JOBS_CHANGED_CHANNEL, lambda *_: self._wake.set())
|
|
logger.info(f"[green]Job consumer running for account {self.account_id}.[/]")
|
|
try:
|
|
while True:
|
|
await self.drain()
|
|
self._wake.clear()
|
|
with contextlib.suppress(TimeoutError):
|
|
await asyncio.wait_for(
|
|
self._wake.wait(), timeout=POLL_INTERVAL_SECONDS
|
|
)
|
|
finally:
|
|
with contextlib.suppress(Exception):
|
|
await conn.close()
|