feat: add jobs queue

backend/src/userbot/modules/jobs/consumer.py

@@ -0,0 +1,85 @@
+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).[/]")
+        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()