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)