feat(backend): implement postgres job queue
This commit is contained in:
349
backend/jobs/services.py
Normal file
349
backend/jobs/services.py
Normal file
@@ -0,0 +1,349 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from datetime import timedelta
|
||||
|
||||
from django.conf import settings
|
||||
from django.db import IntegrityError, connection, transaction
|
||||
from django.db.models import Count, F
|
||||
from django.utils import timezone
|
||||
|
||||
from jobs.models import Job, JobEvent
|
||||
|
||||
|
||||
class JobOwnershipLost(Exception):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class JobCreateResult:
|
||||
job: Job
|
||||
created: bool
|
||||
|
||||
|
||||
def default_max_attempts() -> int:
|
||||
return max(1, int(settings.JOB_WORKER_MAX_ATTEMPTS_DEFAULT))
|
||||
|
||||
|
||||
def backoff_delay(attempt: int) -> timedelta:
|
||||
exponent = max(0, attempt - 1)
|
||||
seconds = min(
|
||||
int(settings.JOB_WORKER_BACKOFF_BASE_SECONDS) * (2**exponent),
|
||||
int(settings.JOB_WORKER_BACKOFF_MAX_SECONDS),
|
||||
)
|
||||
return timedelta(seconds=seconds)
|
||||
|
||||
|
||||
def emit_event(
|
||||
job: Job,
|
||||
event_type: str,
|
||||
*,
|
||||
attempt: int | None = None,
|
||||
worker_id: str | None = None,
|
||||
message: str = "",
|
||||
data: dict | None = None,
|
||||
) -> JobEvent:
|
||||
return JobEvent.objects.create(
|
||||
job=job,
|
||||
type=event_type,
|
||||
attempt=job.attempts if attempt is None else attempt,
|
||||
worker_id=worker_id,
|
||||
message=message,
|
||||
data=data or {},
|
||||
)
|
||||
|
||||
|
||||
@transaction.atomic
|
||||
def create_job(
|
||||
*,
|
||||
job_type: str,
|
||||
payload: dict | None = None,
|
||||
priority: int = 50,
|
||||
available_at=None,
|
||||
max_attempts: int | None = None,
|
||||
idempotency_key: str | None = None,
|
||||
) -> JobCreateResult:
|
||||
normalized_key = idempotency_key.strip() if idempotency_key else None
|
||||
if normalized_key:
|
||||
existing = Job.objects.filter(idempotency_key=normalized_key).first()
|
||||
if existing is not None:
|
||||
return JobCreateResult(existing, created=False)
|
||||
|
||||
try:
|
||||
job = Job.objects.create(
|
||||
type=job_type,
|
||||
payload=payload or {},
|
||||
priority=priority,
|
||||
available_at=available_at or timezone.now(),
|
||||
max_attempts=max_attempts or default_max_attempts(),
|
||||
idempotency_key=normalized_key,
|
||||
)
|
||||
except IntegrityError:
|
||||
if not normalized_key:
|
||||
raise
|
||||
return JobCreateResult(Job.objects.get(idempotency_key=normalized_key), created=False)
|
||||
|
||||
emit_event(
|
||||
job,
|
||||
JobEvent.Type.CREATED,
|
||||
message="Job created",
|
||||
data={
|
||||
"type": job.type,
|
||||
"priority": job.priority,
|
||||
"available_at": job.available_at.isoformat(),
|
||||
"idempotency_key": normalized_key,
|
||||
},
|
||||
)
|
||||
return JobCreateResult(job, created=True)
|
||||
|
||||
|
||||
@transaction.atomic
|
||||
def claim_next_job(*, worker_id: str, lease_seconds: int | None = None) -> Job | None:
|
||||
now = timezone.now()
|
||||
lease_seconds = lease_seconds or settings.JOB_WORKER_LEASE_SECONDS
|
||||
queryset = Job.objects.filter(
|
||||
status=Job.Status.QUEUED,
|
||||
available_at__lte=now,
|
||||
attempts__lt=F("max_attempts"),
|
||||
).order_by("-priority", "available_at", "created_at", "id")
|
||||
|
||||
if connection.features.has_select_for_update_skip_locked:
|
||||
queryset = queryset.select_for_update(skip_locked=True)
|
||||
else:
|
||||
queryset = queryset.select_for_update()
|
||||
|
||||
job = queryset.first()
|
||||
if job is None:
|
||||
return None
|
||||
|
||||
job.status = Job.Status.RUNNING
|
||||
job.attempts += 1
|
||||
job.locked_by = worker_id
|
||||
job.locked_until = now + timedelta(seconds=lease_seconds)
|
||||
job.last_error = ""
|
||||
job.result = None
|
||||
job.save(update_fields=["status", "attempts", "locked_by", "locked_until", "last_error", "result", "updated_at"])
|
||||
emit_event(
|
||||
job,
|
||||
JobEvent.Type.CLAIMED,
|
||||
worker_id=worker_id,
|
||||
message="Job claimed",
|
||||
data={"locked_until": job.locked_until.isoformat()},
|
||||
)
|
||||
return job
|
||||
|
||||
|
||||
@transaction.atomic
|
||||
def renew_lease(job_id, *, worker_id: str, attempt: int, lease_seconds: int | None = None) -> Job | None:
|
||||
now = timezone.now()
|
||||
lease_seconds = lease_seconds or settings.JOB_WORKER_LEASE_SECONDS
|
||||
locked_until = now + timedelta(seconds=lease_seconds)
|
||||
updated = Job.objects.filter(
|
||||
id=job_id,
|
||||
status=Job.Status.RUNNING,
|
||||
locked_by=worker_id,
|
||||
attempts=attempt,
|
||||
).update(locked_until=locked_until, updated_at=now)
|
||||
if not updated:
|
||||
return None
|
||||
|
||||
job = Job.objects.get(id=job_id)
|
||||
emit_event(
|
||||
job,
|
||||
JobEvent.Type.LEASE_RENEWED,
|
||||
attempt=attempt,
|
||||
worker_id=worker_id,
|
||||
message="Lease renewed",
|
||||
data={"locked_until": locked_until.isoformat()},
|
||||
)
|
||||
return job
|
||||
|
||||
|
||||
@transaction.atomic
|
||||
def emit_progress(job_id, *, worker_id: str, attempt: int, percent: int, message: str = "") -> bool:
|
||||
job = (
|
||||
Job.objects.filter(
|
||||
id=job_id,
|
||||
status=Job.Status.RUNNING,
|
||||
locked_by=worker_id,
|
||||
attempts=attempt,
|
||||
)
|
||||
.select_for_update()
|
||||
.first()
|
||||
)
|
||||
if job is None:
|
||||
return False
|
||||
|
||||
emit_event(
|
||||
job,
|
||||
JobEvent.Type.PROGRESS,
|
||||
attempt=attempt,
|
||||
worker_id=worker_id,
|
||||
message=message or f"{percent}% complete",
|
||||
data={"percent": percent},
|
||||
)
|
||||
return True
|
||||
|
||||
|
||||
@transaction.atomic
|
||||
def complete_job(job_id, *, worker_id: str, attempt: int, result: dict | None = None) -> Job | None:
|
||||
now = timezone.now()
|
||||
updated = Job.objects.filter(
|
||||
id=job_id,
|
||||
status=Job.Status.RUNNING,
|
||||
locked_by=worker_id,
|
||||
attempts=attempt,
|
||||
).update(
|
||||
status=Job.Status.SUCCEEDED,
|
||||
result=result or {},
|
||||
locked_by=None,
|
||||
locked_until=None,
|
||||
finished_at=now,
|
||||
updated_at=now,
|
||||
)
|
||||
if not updated:
|
||||
return None
|
||||
|
||||
job = Job.objects.get(id=job_id)
|
||||
emit_event(
|
||||
job,
|
||||
JobEvent.Type.SUCCEEDED,
|
||||
attempt=attempt,
|
||||
worker_id=worker_id,
|
||||
message="Job succeeded",
|
||||
data={"result": job.result},
|
||||
)
|
||||
return job
|
||||
|
||||
|
||||
@transaction.atomic
|
||||
def fail_job(job_id, *, worker_id: str, attempt: int, error: str) -> Job | None:
|
||||
now = timezone.now()
|
||||
job = (
|
||||
Job.objects.select_for_update()
|
||||
.filter(id=job_id, status=Job.Status.RUNNING, locked_by=worker_id, attempts=attempt)
|
||||
.first()
|
||||
)
|
||||
if job is None:
|
||||
return None
|
||||
|
||||
if job.attempts < job.max_attempts:
|
||||
delay = backoff_delay(job.attempts)
|
||||
job.status = Job.Status.QUEUED
|
||||
job.available_at = now + delay
|
||||
job.locked_by = None
|
||||
job.locked_until = None
|
||||
job.last_error = error
|
||||
job.save(update_fields=["status", "available_at", "locked_by", "locked_until", "last_error", "updated_at"])
|
||||
emit_event(
|
||||
job,
|
||||
JobEvent.Type.RETRY_SCHEDULED,
|
||||
attempt=attempt,
|
||||
worker_id=worker_id,
|
||||
message=error,
|
||||
data={"delay_seconds": int(delay.total_seconds()), "available_at": job.available_at.isoformat()},
|
||||
)
|
||||
return job
|
||||
|
||||
job.status = Job.Status.FAILED
|
||||
job.locked_by = None
|
||||
job.locked_until = None
|
||||
job.last_error = error
|
||||
job.finished_at = now
|
||||
job.save(update_fields=["status", "locked_by", "locked_until", "last_error", "finished_at", "updated_at"])
|
||||
emit_event(job, JobEvent.Type.FAILED, attempt=attempt, worker_id=worker_id, message=error, data={"error": error})
|
||||
return job
|
||||
|
||||
|
||||
@transaction.atomic
|
||||
def cleanup_expired_jobs(*, batch_size: int | None = None) -> int:
|
||||
now = timezone.now()
|
||||
batch_size = batch_size or settings.JOB_WORKER_CLEANUP_BATCH_SIZE
|
||||
queryset = Job.objects.filter(status=Job.Status.RUNNING, locked_until__lt=now).order_by("locked_until")
|
||||
if connection.features.has_select_for_update_skip_locked:
|
||||
queryset = queryset.select_for_update(skip_locked=True)
|
||||
else:
|
||||
queryset = queryset.select_for_update()
|
||||
|
||||
expired_jobs = list(queryset[:batch_size])
|
||||
for job in expired_jobs:
|
||||
worker_id = job.locked_by
|
||||
attempt = job.attempts
|
||||
if job.attempts < job.max_attempts:
|
||||
delay = backoff_delay(job.attempts)
|
||||
job.status = Job.Status.QUEUED
|
||||
job.available_at = now + delay
|
||||
job.locked_by = None
|
||||
job.locked_until = None
|
||||
job.last_error = "Worker lease expired"
|
||||
job.save(update_fields=["status", "available_at", "locked_by", "locked_until", "last_error", "updated_at"])
|
||||
emit_event(
|
||||
job,
|
||||
JobEvent.Type.TIMEOUT_REQUEUED,
|
||||
attempt=attempt,
|
||||
worker_id=worker_id,
|
||||
message="Worker lease expired; job requeued",
|
||||
data={"delay_seconds": int(delay.total_seconds()), "available_at": job.available_at.isoformat()},
|
||||
)
|
||||
continue
|
||||
|
||||
job.status = Job.Status.FAILED
|
||||
job.locked_by = None
|
||||
job.locked_until = None
|
||||
job.last_error = "Worker lease expired"
|
||||
job.finished_at = now
|
||||
job.save(update_fields=["status", "locked_by", "locked_until", "last_error", "finished_at", "updated_at"])
|
||||
emit_event(
|
||||
job,
|
||||
JobEvent.Type.FAILED,
|
||||
attempt=attempt,
|
||||
worker_id=worker_id,
|
||||
message="Worker lease expired; attempts exhausted",
|
||||
data={"error": "Worker lease expired"},
|
||||
)
|
||||
|
||||
return len(expired_jobs)
|
||||
|
||||
|
||||
@transaction.atomic
|
||||
def retry_failed_job(job_id) -> Job:
|
||||
job = Job.objects.select_for_update().get(id=job_id)
|
||||
if job.status != Job.Status.FAILED:
|
||||
raise ValueError("Only failed jobs can be retried.")
|
||||
|
||||
job.status = Job.Status.QUEUED
|
||||
job.available_at = timezone.now()
|
||||
job.attempts = 0
|
||||
job.locked_by = None
|
||||
job.locked_until = None
|
||||
job.last_error = ""
|
||||
job.result = None
|
||||
job.finished_at = None
|
||||
job.save(
|
||||
update_fields=[
|
||||
"status",
|
||||
"available_at",
|
||||
"attempts",
|
||||
"locked_by",
|
||||
"locked_until",
|
||||
"last_error",
|
||||
"result",
|
||||
"finished_at",
|
||||
"updated_at",
|
||||
]
|
||||
)
|
||||
emit_event(job, JobEvent.Type.MANUAL_RETRY, message="Failed job manually retried")
|
||||
return job
|
||||
|
||||
|
||||
def job_stats() -> dict:
|
||||
status_counts = dict(Job.objects.order_by().values_list("status").annotate(count=Count("id")))
|
||||
now = timezone.now()
|
||||
oldest_queued = Job.objects.filter(status=Job.Status.QUEUED).order_by("created_at").first()
|
||||
return {
|
||||
"total": Job.objects.count(),
|
||||
"by_status": {status: status_counts.get(status, 0) for status in Job.Status.values},
|
||||
"overdue_running": Job.objects.filter(status=Job.Status.RUNNING, locked_until__lt=now).count(),
|
||||
"retries_pending": Job.objects.filter(status=Job.Status.QUEUED, attempts__gt=0).count(),
|
||||
"oldest_queued_at": oldest_queued.created_at.isoformat() if oldest_queued else None,
|
||||
}
|
||||
Reference in New Issue
Block a user