diff --git a/backend/migrations/versions/f8c2d6e1a4b7_job_run_after.py b/backend/migrations/versions/f8c2d6e1a4b7_job_run_after.py new file mode 100644 index 0000000..14cef95 --- /dev/null +++ b/backend/migrations/versions/f8c2d6e1a4b7_job_run_after.py @@ -0,0 +1,27 @@ +"""job run after + +Revision ID: f8c2d6e1a4b7 +Revises: c9e2a7f4b1d8 +Create Date: 2026-10-19 03:00:00.000000 + +""" + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op + +revision: str = "f8c2d6e1a4b7" +down_revision: str | None = "c9e2a7f4b1d8" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + op.add_column( + "jobs", sa.Column("run_after", sa.DateTime(timezone=True), nullable=True) + ) + + +def downgrade() -> None: + op.drop_column("jobs", "run_after") diff --git a/backend/src/userbot/modules/jobs/consumer.py b/backend/src/userbot/modules/jobs/consumer.py index 68f598a..8840029 100644 --- a/backend/src/userbot/modules/jobs/consumer.py +++ b/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) diff --git a/backend/src/userbot/modules/jobs/repository.py b/backend/src/userbot/modules/jobs/repository.py index 72b8eae..9fe84ce 100644 --- a/backend/src/userbot/modules/jobs/repository.py +++ b/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, ) diff --git a/backend/src/utils/db/models.py b/backend/src/utils/db/models.py index 5706525..1adb50c 100644 --- a/backend/src/utils/db/models.py +++ b/backend/src/utils/db/models.py @@ -203,6 +203,9 @@ class Job(SQLModel, table=True): finished_at: datetime | None = Field( default=None, sa_column=Column(DateTime(timezone=True)) ) + run_after: datetime | None = Field( + default=None, sa_column=Column(DateTime(timezone=True)) + ) class CapturePolicy(SQLModel, table=True):