fix(userbot): job queue survives a consumer failure, a flood wait defers its job instead of holding the queue

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

@@ -14,6 +14,7 @@ from utils.logging import logger
 
 JOBS_CHANGED_CHANNEL = "jobs_changed"
 POLL_INTERVAL_SECONDS = 60.0
+RESTART_SECONDS = 10.0
 
 
 class JobConsumer:
@@ -38,25 +39,20 @@ class JobConsumer:
                 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
+        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:
@@ -66,6 +62,17 @@ class JobConsumer:
             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).[/]")
@@ -81,11 +88,15 @@ class JobConsumer:
         logger.info(f"[green]Job consumer running for account {self.account_id}.[/]")
         try:
             while True:
-                await self.drain()
                 self._wake.clear()
+                await self.drain()
                 with contextlib.suppress(TimeoutError):
                     await asyncio.wait_for(
-                        self._wake.wait(), timeout=POLL_INTERVAL_SECONDS
+                        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)

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

@@ -32,6 +32,7 @@ async def claim_next(pool: asyncpg.Pool, account_id: int) -> Job | None:
         "updated_at = now(), attempts = attempts + 1 "
         "WHERE id = ("
         "    SELECT id FROM jobs WHERE account_id = $1 AND status = 'pending' "
+        "    AND (run_after IS NULL OR run_after <= now()) "
         "    ORDER BY id FOR UPDATE SKIP LOCKED LIMIT 1"
         ") RETURNING *",
         account_id,
@@ -76,11 +77,21 @@ async def report_progress(
     )
 
 
-async def bump_flood_wait(pool: asyncpg.Pool, job_id: int) -> None:
+async def defer_flood_wait(pool: asyncpg.Pool, job_id: int, seconds: int) -> None:
     await pool.execute(
-        "UPDATE jobs SET flood_waits = flood_waits + 1, updated_at = now() "
-        "WHERE id = $1",
+        "UPDATE jobs SET status = 'pending', flood_waits = flood_waits + 1, "
+        "run_after = now() + make_interval(secs => $2), updated_at = now() "
+        "WHERE id = $1 AND status = 'running'",
         job_id,
+        seconds,
+    )
+
+
+async def seconds_to_next_due(pool: asyncpg.Pool, account_id: int) -> float | None:
+    return await pool.fetchval(
+        "SELECT extract(epoch FROM min(run_after) - now())::float FROM jobs "
+        "WHERE account_id = $1 AND status = 'pending' AND run_after > now()",
+        account_id,
     )