86 lines
2.9 KiB
Python
86 lines
2.9 KiB
Python
from __future__ import annotations
|
|
|
|
import time
|
|
from dataclasses import dataclass
|
|
|
|
from django.conf import settings
|
|
|
|
from jobs.services import emit_progress, renew_lease
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class HandlerContext:
|
|
job_id: str
|
|
worker_id: str
|
|
attempt: int
|
|
|
|
def progress(self, percent: int, message: str = "") -> None:
|
|
emit_progress(self.job_id, worker_id=self.worker_id, attempt=self.attempt, percent=percent, message=message)
|
|
|
|
def renew(self) -> None:
|
|
renew_lease(
|
|
self.job_id,
|
|
worker_id=self.worker_id,
|
|
attempt=self.attempt,
|
|
lease_seconds=settings.JOB_WORKER_LEASE_SECONDS,
|
|
)
|
|
|
|
|
|
def _sleep_with_progress(context: HandlerContext, seconds: int, *, renew: bool) -> None:
|
|
seconds = max(1, seconds)
|
|
for elapsed in range(seconds):
|
|
time.sleep(1)
|
|
percent = min(95, int(((elapsed + 1) / seconds) * 100))
|
|
if renew:
|
|
context.renew()
|
|
context.progress(percent, f"{percent}% complete")
|
|
|
|
|
|
def handle_success(payload: dict, context: HandlerContext) -> dict:
|
|
sleep_seconds = int(payload.get("sleep_seconds", 1))
|
|
_sleep_with_progress(context, sleep_seconds, renew=True)
|
|
return {"ok": True, "mode": "success"}
|
|
|
|
|
|
def handle_fail(payload: dict, context: HandlerContext) -> dict:
|
|
context.progress(25, "Intentional failure started")
|
|
raise RuntimeError(str(payload.get("error", "Intentional demo failure")))
|
|
|
|
|
|
def handle_slow(payload: dict, context: HandlerContext) -> dict:
|
|
sleep_seconds = int(payload.get("sleep_seconds", 8))
|
|
_sleep_with_progress(context, sleep_seconds, renew=True)
|
|
return {"ok": True, "mode": "slow", "slept_seconds": sleep_seconds}
|
|
|
|
|
|
def handle_timeout(payload: dict, context: HandlerContext) -> dict:
|
|
sleep_seconds = int(payload.get("sleep_seconds", settings.JOB_WORKER_LEASE_SECONDS + 10))
|
|
_sleep_with_progress(context, sleep_seconds, renew=False)
|
|
return {"ok": True, "mode": "timeout", "slept_seconds": sleep_seconds}
|
|
|
|
|
|
def handle_flaky(payload: dict, context: HandlerContext) -> dict:
|
|
fail_until_attempt = int(payload.get("fail_until_attempt", 2))
|
|
context.progress(35, "Flaky job evaluated")
|
|
if context.attempt <= fail_until_attempt:
|
|
raise RuntimeError(f"Flaky failure on attempt {context.attempt}")
|
|
context.progress(100, "Flaky job succeeded")
|
|
return {"ok": True, "mode": "flaky", "attempt": context.attempt}
|
|
|
|
|
|
HANDLERS = {
|
|
"demo.success": handle_success,
|
|
"demo.fail": handle_fail,
|
|
"demo.slow": handle_slow,
|
|
"demo.timeout": handle_timeout,
|
|
"demo.flaky": handle_flaky,
|
|
}
|
|
|
|
|
|
def execute_handler(job, *, worker_id: str) -> dict:
|
|
handler = HANDLERS.get(job.type)
|
|
if handler is None:
|
|
raise ValueError(f"Unknown job type: {job.type}")
|
|
context = HandlerContext(job_id=str(job.id), worker_id=worker_id, attempt=job.attempts)
|
|
return handler(job.payload or {}, context)
|