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)