fix(userbot): job queue survives a consumer failure, a flood wait defers its job instead of holding the queue
This commit is contained in:
1 parent
3373c31a02
commit
42469eb9f1
4 files changed
+66
-14
No files matched your search
@@ -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")
|
||||||
@@ -14,6 +14,7 @@ from utils.logging import logger
|
|||||||
|
|
||||||
JOBS_CHANGED_CHANNEL = "jobs_changed"
|
JOBS_CHANGED_CHANNEL = "jobs_changed"
|
||||||
POLL_INTERVAL_SECONDS = 60.0
|
POLL_INTERVAL_SECONDS = 60.0
|
||||||
|
RESTART_SECONDS = 10.0
|
||||||
|
|
||||||
|
|
||||||
class JobConsumer:
|
class JobConsumer:
|
||||||
@@ -38,25 +39,20 @@ class JobConsumer:
|
|||||||
self.pool, ctx.job_id, JobStatus.FAILED, f"no handler for {job.kind!r}"
|
self.pool, ctx.job_id, JobStatus.FAILED, f"no handler for {job.kind!r}"
|
||||||
)
|
)
|
||||||
return
|
return
|
||||||
while True:
|
|
||||||
try:
|
try:
|
||||||
await handler(ctx)
|
await handler(ctx)
|
||||||
except (FloodWait, FloodPremiumWait) as exc:
|
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
|
wait = exc.value if isinstance(exc.value, int) else 1
|
||||||
|
await repository.defer_flood_wait(self.pool, ctx.job_id, wait)
|
||||||
logger.warning(
|
logger.warning(
|
||||||
f"[yellow]FloodWait {wait}s on job {ctx.job_id} ({job.kind})[/]"
|
f"[yellow]FloodWait {wait}s on job {ctx.job_id} ({job.kind}), "
|
||||||
|
f"deferred[/]"
|
||||||
)
|
)
|
||||||
await asyncio.sleep(wait)
|
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
logger.exception(f"Job {ctx.job_id} ({job.kind}) failed")
|
logger.exception(f"Job {ctx.job_id} ({job.kind}) failed")
|
||||||
await repository.finish(
|
await repository.finish(self.pool, ctx.job_id, JobStatus.FAILED, str(exc))
|
||||||
self.pool, ctx.job_id, JobStatus.FAILED, str(exc)
|
|
||||||
)
|
|
||||||
return
|
|
||||||
else:
|
else:
|
||||||
await repository.finish(self.pool, ctx.job_id, JobStatus.DONE)
|
await repository.finish(self.pool, ctx.job_id, JobStatus.DONE)
|
||||||
return
|
|
||||||
|
|
||||||
async def drain(self) -> None:
|
async def drain(self) -> None:
|
||||||
while True:
|
while True:
|
||||||
@@ -66,6 +62,17 @@ class JobConsumer:
|
|||||||
await self.execute(job)
|
await self.execute(job)
|
||||||
|
|
||||||
async def run(self) -> None:
|
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)
|
requeued = await repository.requeue_running(self.pool, self.account_id)
|
||||||
if requeued:
|
if requeued:
|
||||||
logger.info(f"[yellow]Requeued {requeued} stale running job(s).[/]")
|
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}.[/]")
|
logger.info(f"[green]Job consumer running for account {self.account_id}.[/]")
|
||||||
try:
|
try:
|
||||||
while True:
|
while True:
|
||||||
await self.drain()
|
|
||||||
self._wake.clear()
|
self._wake.clear()
|
||||||
|
await self.drain()
|
||||||
with contextlib.suppress(TimeoutError):
|
with contextlib.suppress(TimeoutError):
|
||||||
await asyncio.wait_for(
|
await asyncio.wait_for(
|
||||||
self._wake.wait(), timeout=POLL_INTERVAL_SECONDS
|
self._wake.wait(), timeout=await self._idle_seconds()
|
||||||
)
|
)
|
||||||
finally:
|
finally:
|
||||||
await listener.stop()
|
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)
|
||||||
@@ -32,6 +32,7 @@ async def claim_next(pool: asyncpg.Pool, account_id: int) -> Job | None:
|
|||||||
"updated_at = now(), attempts = attempts + 1 "
|
"updated_at = now(), attempts = attempts + 1 "
|
||||||
"WHERE id = ("
|
"WHERE id = ("
|
||||||
" SELECT id FROM jobs WHERE account_id = $1 AND status = 'pending' "
|
" 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"
|
" ORDER BY id FOR UPDATE SKIP LOCKED LIMIT 1"
|
||||||
") RETURNING *",
|
") RETURNING *",
|
||||||
account_id,
|
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(
|
await pool.execute(
|
||||||
"UPDATE jobs SET flood_waits = flood_waits + 1, updated_at = now() "
|
"UPDATE jobs SET status = 'pending', flood_waits = flood_waits + 1, "
|
||||||
"WHERE id = $1",
|
"run_after = now() + make_interval(secs => $2), updated_at = now() "
|
||||||
|
"WHERE id = $1 AND status = 'running'",
|
||||||
job_id,
|
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,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -203,6 +203,9 @@ class Job(SQLModel, table=True):
|
|||||||
finished_at: datetime | None = Field(
|
finished_at: datetime | None = Field(
|
||||||
default=None, sa_column=Column(DateTime(timezone=True))
|
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):
|
class CapturePolicy(SQLModel, table=True):
|
||||||
|
|||||||
Reference in new issue
Block a user