feat(backend): implement postgres job queue
This commit is contained in:
1
backend/jobs/__init__.py
Normal file
1
backend/jobs/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
|
||||
18
backend/jobs/admin.py
Normal file
18
backend/jobs/admin.py
Normal file
@@ -0,0 +1,18 @@
|
||||
from django.contrib import admin
|
||||
|
||||
from jobs.models import Job, JobEvent
|
||||
|
||||
|
||||
@admin.register(Job)
|
||||
class JobAdmin(admin.ModelAdmin):
|
||||
list_display = ("id", "type", "status", "priority", "attempts", "max_attempts", "locked_by", "available_at")
|
||||
list_filter = ("status", "type")
|
||||
search_fields = ("id", "type", "idempotency_key", "locked_by")
|
||||
readonly_fields = ("created_at", "updated_at", "finished_at")
|
||||
|
||||
|
||||
@admin.register(JobEvent)
|
||||
class JobEventAdmin(admin.ModelAdmin):
|
||||
list_display = ("id", "job", "type", "attempt", "worker_id", "created_at")
|
||||
list_filter = ("type",)
|
||||
search_fields = ("job__id", "worker_id", "message")
|
||||
6
backend/jobs/apps.py
Normal file
6
backend/jobs/apps.py
Normal file
@@ -0,0 +1,6 @@
|
||||
from django.apps import AppConfig
|
||||
|
||||
|
||||
class JobsConfig(AppConfig):
|
||||
default_auto_field = "django.db.models.BigAutoField"
|
||||
name = "jobs"
|
||||
85
backend/jobs/handlers.py
Normal file
85
backend/jobs/handlers.py
Normal file
@@ -0,0 +1,85 @@
|
||||
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)
|
||||
1
backend/jobs/management/__init__.py
Normal file
1
backend/jobs/management/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
|
||||
1
backend/jobs/management/commands/__init__.py
Normal file
1
backend/jobs/management/commands/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
|
||||
21
backend/jobs/management/commands/run_job_workers.py
Normal file
21
backend/jobs/management/commands/run_job_workers.py
Normal file
@@ -0,0 +1,21 @@
|
||||
import logging
|
||||
import signal
|
||||
|
||||
from django.core.management.base import BaseCommand
|
||||
|
||||
from jobs.worker import JobWorkerRunner
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Command(BaseCommand):
|
||||
help = "Run the configured threaded job worker."
|
||||
|
||||
def handle(self, *args, **options):
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s %(message)s")
|
||||
runner = JobWorkerRunner()
|
||||
signal.signal(signal.SIGINT, runner.request_stop)
|
||||
signal.signal(signal.SIGTERM, runner.request_stop)
|
||||
self.stdout.write(self.style.SUCCESS(f"Starting {runner.thread_count} job worker thread(s)."))
|
||||
runner.run_forever()
|
||||
self.stdout.write(self.style.SUCCESS("Job worker stopped."))
|
||||
128
backend/jobs/migrations/0001_initial.py
Normal file
128
backend/jobs/migrations/0001_initial.py
Normal file
@@ -0,0 +1,128 @@
|
||||
# Generated manually for the interview project.
|
||||
|
||||
import uuid
|
||||
|
||||
import django.db.models.deletion
|
||||
from django.db import migrations, models
|
||||
from django.db.models import Q
|
||||
|
||||
|
||||
class Migration(migrations.Migration):
|
||||
initial = True
|
||||
|
||||
dependencies = []
|
||||
|
||||
operations = [
|
||||
migrations.CreateModel(
|
||||
name="Job",
|
||||
fields=[
|
||||
("id", models.UUIDField(default=uuid.uuid4, editable=False, primary_key=True, serialize=False)),
|
||||
("type", models.CharField(max_length=120)),
|
||||
("payload", models.JSONField(blank=True, default=dict)),
|
||||
(
|
||||
"status",
|
||||
models.CharField(
|
||||
choices=[
|
||||
("queued", "Queued"),
|
||||
("running", "Running"),
|
||||
("succeeded", "Succeeded"),
|
||||
("failed", "Failed"),
|
||||
],
|
||||
db_index=True,
|
||||
default="queued",
|
||||
max_length=20,
|
||||
),
|
||||
),
|
||||
("priority", models.SmallIntegerField(default=50)),
|
||||
("available_at", models.DateTimeField(db_index=True)),
|
||||
("attempts", models.PositiveIntegerField(default=0)),
|
||||
("max_attempts", models.PositiveIntegerField(default=3)),
|
||||
("idempotency_key", models.CharField(blank=True, max_length=255, null=True)),
|
||||
("locked_by", models.CharField(blank=True, max_length=255, null=True)),
|
||||
("locked_until", models.DateTimeField(blank=True, null=True)),
|
||||
("last_error", models.TextField(blank=True, default="")),
|
||||
("result", models.JSONField(blank=True, null=True)),
|
||||
("created_at", models.DateTimeField(auto_now_add=True, db_index=True)),
|
||||
("updated_at", models.DateTimeField(auto_now=True)),
|
||||
("finished_at", models.DateTimeField(blank=True, null=True)),
|
||||
],
|
||||
options={"ordering": ("-created_at",)},
|
||||
),
|
||||
migrations.CreateModel(
|
||||
name="JobEvent",
|
||||
fields=[
|
||||
("id", models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name="ID")),
|
||||
(
|
||||
"type",
|
||||
models.CharField(
|
||||
choices=[
|
||||
("created", "Created"),
|
||||
("claimed", "Claimed"),
|
||||
("progress", "Progress"),
|
||||
("lease_renewed", "Lease renewed"),
|
||||
("succeeded", "Succeeded"),
|
||||
("retry_scheduled", "Retry scheduled"),
|
||||
("failed", "Failed"),
|
||||
("timeout_requeued", "Timeout requeued"),
|
||||
("manual_retry", "Manual retry"),
|
||||
],
|
||||
db_index=True,
|
||||
max_length=40,
|
||||
),
|
||||
),
|
||||
("attempt", models.PositiveIntegerField(default=0)),
|
||||
("worker_id", models.CharField(blank=True, max_length=255, null=True)),
|
||||
("message", models.TextField(blank=True, default="")),
|
||||
("data", models.JSONField(blank=True, default=dict)),
|
||||
("created_at", models.DateTimeField(auto_now_add=True, db_index=True)),
|
||||
("job", models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name="events", to="jobs.job")),
|
||||
],
|
||||
options={"ordering": ("id",)},
|
||||
),
|
||||
migrations.AddConstraint(
|
||||
model_name="job",
|
||||
constraint=models.CheckConstraint(
|
||||
name="job_status_valid_shape",
|
||||
condition=(
|
||||
Q(finished_at__isnull=True, locked_by__isnull=True, locked_until__isnull=True, status="queued")
|
||||
| Q(finished_at__isnull=True, locked_by__isnull=False, locked_until__isnull=False, status="running")
|
||||
| Q(finished_at__isnull=False, locked_by__isnull=True, locked_until__isnull=True, status__in=["succeeded", "failed"])
|
||||
),
|
||||
),
|
||||
),
|
||||
migrations.AddConstraint(
|
||||
model_name="job",
|
||||
constraint=models.CheckConstraint(
|
||||
name="job_attempts_valid",
|
||||
condition=Q(attempts__gte=0, attempts__lte=models.F("max_attempts"), max_attempts__gte=1),
|
||||
),
|
||||
),
|
||||
migrations.AddConstraint(
|
||||
model_name="job",
|
||||
constraint=models.UniqueConstraint(
|
||||
condition=Q(idempotency_key__isnull=False),
|
||||
fields=("idempotency_key",),
|
||||
name="job_idempotency_key_unique",
|
||||
),
|
||||
),
|
||||
migrations.AddIndex(
|
||||
model_name="job",
|
||||
index=models.Index(
|
||||
condition=Q(status="queued"),
|
||||
fields=["available_at", "-priority", "created_at", "id"],
|
||||
name="job_claim_idx",
|
||||
),
|
||||
),
|
||||
migrations.AddIndex(
|
||||
model_name="job",
|
||||
index=models.Index(condition=Q(status="running"), fields=["locked_until"], name="job_timeout_idx"),
|
||||
),
|
||||
migrations.AddIndex(
|
||||
model_name="jobevent",
|
||||
index=models.Index(fields=["job", "id"], name="job_event_job_id_idx"),
|
||||
),
|
||||
migrations.AddIndex(
|
||||
model_name="jobevent",
|
||||
index=models.Index(fields=["id"], name="job_event_id_idx"),
|
||||
),
|
||||
]
|
||||
1
backend/jobs/migrations/__init__.py
Normal file
1
backend/jobs/migrations/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
|
||||
108
backend/jobs/models.py
Normal file
108
backend/jobs/models.py
Normal file
@@ -0,0 +1,108 @@
|
||||
import uuid
|
||||
|
||||
from django.db import models
|
||||
from django.db.models import Q
|
||||
|
||||
|
||||
class Job(models.Model):
|
||||
class Status(models.TextChoices):
|
||||
QUEUED = "queued", "Queued"
|
||||
RUNNING = "running", "Running"
|
||||
SUCCEEDED = "succeeded", "Succeeded"
|
||||
FAILED = "failed", "Failed"
|
||||
|
||||
id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False)
|
||||
type = models.CharField(max_length=120)
|
||||
payload = models.JSONField(default=dict, blank=True)
|
||||
status = models.CharField(max_length=20, choices=Status.choices, default=Status.QUEUED, db_index=True)
|
||||
priority = models.SmallIntegerField(default=50)
|
||||
available_at = models.DateTimeField(db_index=True)
|
||||
attempts = models.PositiveIntegerField(default=0)
|
||||
max_attempts = models.PositiveIntegerField(default=3)
|
||||
idempotency_key = models.CharField(max_length=255, null=True, blank=True)
|
||||
locked_by = models.CharField(max_length=255, null=True, blank=True)
|
||||
locked_until = models.DateTimeField(null=True, blank=True)
|
||||
last_error = models.TextField(blank=True, default="")
|
||||
result = models.JSONField(null=True, blank=True)
|
||||
created_at = models.DateTimeField(auto_now_add=True, db_index=True)
|
||||
updated_at = models.DateTimeField(auto_now=True)
|
||||
finished_at = models.DateTimeField(null=True, blank=True)
|
||||
|
||||
class Meta:
|
||||
ordering = ("-created_at",)
|
||||
constraints = [
|
||||
models.CheckConstraint(
|
||||
name="job_status_valid_shape",
|
||||
condition=(
|
||||
Q(
|
||||
status="queued",
|
||||
locked_by__isnull=True,
|
||||
locked_until__isnull=True,
|
||||
finished_at__isnull=True,
|
||||
)
|
||||
| Q(
|
||||
status="running",
|
||||
locked_by__isnull=False,
|
||||
locked_until__isnull=False,
|
||||
finished_at__isnull=True,
|
||||
)
|
||||
| Q(
|
||||
status__in=["succeeded", "failed"],
|
||||
locked_by__isnull=True,
|
||||
locked_until__isnull=True,
|
||||
finished_at__isnull=False,
|
||||
)
|
||||
),
|
||||
),
|
||||
models.CheckConstraint(
|
||||
name="job_attempts_valid",
|
||||
condition=Q(attempts__gte=0, max_attempts__gte=1, attempts__lte=models.F("max_attempts")),
|
||||
),
|
||||
models.UniqueConstraint(
|
||||
fields=["idempotency_key"],
|
||||
condition=Q(idempotency_key__isnull=False),
|
||||
name="job_idempotency_key_unique",
|
||||
),
|
||||
]
|
||||
indexes = [
|
||||
models.Index(
|
||||
fields=["available_at", "-priority", "created_at", "id"],
|
||||
name="job_claim_idx",
|
||||
condition=Q(status="queued"),
|
||||
),
|
||||
models.Index(fields=["locked_until"], name="job_timeout_idx", condition=Q(status="running")),
|
||||
]
|
||||
|
||||
def __str__(self):
|
||||
return f"{self.type}:{self.id}"
|
||||
|
||||
|
||||
class JobEvent(models.Model):
|
||||
class Type(models.TextChoices):
|
||||
CREATED = "created", "Created"
|
||||
CLAIMED = "claimed", "Claimed"
|
||||
PROGRESS = "progress", "Progress"
|
||||
LEASE_RENEWED = "lease_renewed", "Lease renewed"
|
||||
SUCCEEDED = "succeeded", "Succeeded"
|
||||
RETRY_SCHEDULED = "retry_scheduled", "Retry scheduled"
|
||||
FAILED = "failed", "Failed"
|
||||
TIMEOUT_REQUEUED = "timeout_requeued", "Timeout requeued"
|
||||
MANUAL_RETRY = "manual_retry", "Manual retry"
|
||||
|
||||
job = models.ForeignKey(Job, on_delete=models.CASCADE, related_name="events")
|
||||
type = models.CharField(max_length=40, choices=Type.choices, db_index=True)
|
||||
attempt = models.PositiveIntegerField(default=0)
|
||||
worker_id = models.CharField(max_length=255, null=True, blank=True)
|
||||
message = models.TextField(blank=True, default="")
|
||||
data = models.JSONField(default=dict, blank=True)
|
||||
created_at = models.DateTimeField(auto_now_add=True, db_index=True)
|
||||
|
||||
class Meta:
|
||||
ordering = ("id",)
|
||||
indexes = [
|
||||
models.Index(fields=["job", "id"], name="job_event_job_id_idx"),
|
||||
models.Index(fields=["id"], name="job_event_id_idx"),
|
||||
]
|
||||
|
||||
def __str__(self):
|
||||
return f"{self.type}:{self.job_id}:{self.id}"
|
||||
71
backend/jobs/serializers.py
Normal file
71
backend/jobs/serializers.py
Normal file
@@ -0,0 +1,71 @@
|
||||
from django.conf import settings
|
||||
from rest_framework import serializers
|
||||
|
||||
from jobs.models import Job, JobEvent
|
||||
|
||||
|
||||
class JobSerializer(serializers.ModelSerializer):
|
||||
class Meta:
|
||||
model = Job
|
||||
fields = (
|
||||
"id",
|
||||
"type",
|
||||
"payload",
|
||||
"status",
|
||||
"priority",
|
||||
"available_at",
|
||||
"attempts",
|
||||
"max_attempts",
|
||||
"idempotency_key",
|
||||
"locked_by",
|
||||
"locked_until",
|
||||
"last_error",
|
||||
"result",
|
||||
"created_at",
|
||||
"updated_at",
|
||||
"finished_at",
|
||||
)
|
||||
read_only_fields = (
|
||||
"id",
|
||||
"status",
|
||||
"attempts",
|
||||
"locked_by",
|
||||
"locked_until",
|
||||
"last_error",
|
||||
"result",
|
||||
"created_at",
|
||||
"updated_at",
|
||||
"finished_at",
|
||||
)
|
||||
|
||||
|
||||
class JobCreateSerializer(serializers.Serializer):
|
||||
type = serializers.CharField(max_length=120)
|
||||
payload = serializers.JSONField(required=False, default=dict)
|
||||
priority = serializers.IntegerField(required=False, min_value=-32768, max_value=32767, default=50)
|
||||
available_at = serializers.DateTimeField(required=False, allow_null=True, default=None)
|
||||
max_attempts = serializers.IntegerField(
|
||||
required=False,
|
||||
min_value=1,
|
||||
default=lambda: settings.JOB_WORKER_MAX_ATTEMPTS_DEFAULT,
|
||||
)
|
||||
idempotency_key = serializers.CharField(
|
||||
required=False,
|
||||
allow_blank=True,
|
||||
allow_null=True,
|
||||
max_length=255,
|
||||
default=None,
|
||||
)
|
||||
|
||||
def validate_idempotency_key(self, value):
|
||||
if value is None:
|
||||
return None
|
||||
value = value.strip()
|
||||
return value or None
|
||||
|
||||
|
||||
class JobEventSerializer(serializers.ModelSerializer):
|
||||
class Meta:
|
||||
model = JobEvent
|
||||
fields = ("id", "job", "type", "attempt", "worker_id", "message", "data", "created_at")
|
||||
read_only_fields = fields
|
||||
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,
|
||||
}
|
||||
185
backend/jobs/tests/test_services.py
Normal file
185
backend/jobs/tests/test_services.py
Normal file
@@ -0,0 +1,185 @@
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from datetime import timedelta
|
||||
|
||||
import pytest
|
||||
from django.conf import settings
|
||||
from django.db import close_old_connections, connection
|
||||
from django.utils import timezone
|
||||
|
||||
from jobs.models import Job, JobEvent
|
||||
from jobs.services import (
|
||||
claim_next_job,
|
||||
cleanup_expired_jobs,
|
||||
complete_job,
|
||||
create_job,
|
||||
fail_job,
|
||||
renew_lease,
|
||||
retry_failed_job,
|
||||
)
|
||||
from jobs.worker import JobWorkerRunner
|
||||
|
||||
|
||||
@pytest.mark.django_db
|
||||
def test_successful_claim_sets_lock_and_event():
|
||||
job = create_job(job_type="demo.success").job
|
||||
|
||||
claimed = claim_next_job(worker_id="worker-a", lease_seconds=30)
|
||||
|
||||
assert claimed.id == job.id
|
||||
assert claimed.status == Job.Status.RUNNING
|
||||
assert claimed.locked_by == "worker-a"
|
||||
assert claimed.locked_until is not None
|
||||
assert claimed.attempts == 1
|
||||
assert JobEvent.objects.filter(job=job, type=JobEvent.Type.CLAIMED, worker_id="worker-a").count() == 1
|
||||
|
||||
|
||||
@pytest.mark.django_db(transaction=True)
|
||||
def test_concurrent_claim_prevention_with_postgresql():
|
||||
if connection.vendor != "postgresql":
|
||||
pytest.skip("SKIP LOCKED concurrency is a PostgreSQL behavior.")
|
||||
create_job(job_type="demo.success")
|
||||
|
||||
def claim(index):
|
||||
close_old_connections()
|
||||
try:
|
||||
job = claim_next_job(worker_id=f"worker-{index}", lease_seconds=30)
|
||||
return str(job.id) if job else None
|
||||
finally:
|
||||
close_old_connections()
|
||||
|
||||
with ThreadPoolExecutor(max_workers=10) as executor:
|
||||
claimed_ids = list(executor.map(claim, range(10)))
|
||||
|
||||
assert len([job_id for job_id in claimed_ids if job_id]) == 1
|
||||
assert JobEvent.objects.filter(type=JobEvent.Type.CLAIMED).count() == 1
|
||||
|
||||
|
||||
@pytest.mark.django_db
|
||||
def test_priority_and_available_at_ordering_are_deterministic():
|
||||
future = timezone.now() + timedelta(hours=1)
|
||||
low = create_job(job_type="demo.success", priority=10).job
|
||||
high = create_job(job_type="demo.success", priority=90).job
|
||||
create_job(job_type="demo.success", priority=100, available_at=future)
|
||||
|
||||
claimed = claim_next_job(worker_id="worker-a", lease_seconds=30)
|
||||
|
||||
assert claimed.id == high.id
|
||||
next_claim = claim_next_job(worker_id="worker-b", lease_seconds=30)
|
||||
assert next_claim.id == low.id
|
||||
|
||||
|
||||
@pytest.mark.django_db
|
||||
def test_idempotent_create_returns_existing_job():
|
||||
first = create_job(job_type="demo.success", idempotency_key="invoice-123")
|
||||
second = create_job(job_type="demo.success", idempotency_key="invoice-123")
|
||||
|
||||
assert first.created is True
|
||||
assert second.created is False
|
||||
assert second.job.id == first.job.id
|
||||
assert Job.objects.count() == 1
|
||||
assert JobEvent.objects.filter(type=JobEvent.Type.CREATED).count() == 1
|
||||
|
||||
|
||||
@pytest.mark.django_db
|
||||
def test_failure_schedules_exponential_retry():
|
||||
job = create_job(job_type="demo.fail", max_attempts=3).job
|
||||
claimed = claim_next_job(worker_id="worker-a", lease_seconds=30)
|
||||
before = timezone.now()
|
||||
|
||||
updated = fail_job(claimed.id, worker_id="worker-a", attempt=1, error="boom")
|
||||
|
||||
assert updated.status == Job.Status.QUEUED
|
||||
assert updated.attempts == 1
|
||||
assert updated.available_at >= before + timedelta(seconds=settings.JOB_WORKER_BACKOFF_BASE_SECONDS)
|
||||
assert updated.locked_by is None
|
||||
assert JobEvent.objects.filter(job=job, type=JobEvent.Type.RETRY_SCHEDULED).exists()
|
||||
|
||||
|
||||
@pytest.mark.django_db
|
||||
def test_exhausted_attempts_become_failed():
|
||||
job = create_job(job_type="demo.fail", max_attempts=1).job
|
||||
claimed = claim_next_job(worker_id="worker-a", lease_seconds=30)
|
||||
|
||||
updated = fail_job(claimed.id, worker_id="worker-a", attempt=1, error="boom")
|
||||
|
||||
assert updated.status == Job.Status.FAILED
|
||||
assert updated.finished_at is not None
|
||||
assert updated.locked_by is None
|
||||
assert JobEvent.objects.filter(job=job, type=JobEvent.Type.FAILED).exists()
|
||||
|
||||
|
||||
@pytest.mark.django_db
|
||||
def test_expired_lease_cleanup_requeues_or_fails():
|
||||
requeue = create_job(job_type="demo.timeout", max_attempts=2).job
|
||||
fail = create_job(job_type="demo.timeout", max_attempts=1).job
|
||||
requeue_claim = claim_next_job(worker_id="worker-a", lease_seconds=30)
|
||||
fail_claim = claim_next_job(worker_id="worker-b", lease_seconds=30)
|
||||
|
||||
Job.objects.filter(id__in=[requeue_claim.id, fail_claim.id]).update(
|
||||
locked_until=timezone.now() - timedelta(seconds=1)
|
||||
)
|
||||
|
||||
assert cleanup_expired_jobs(batch_size=10) == 2
|
||||
|
||||
requeue.refresh_from_db()
|
||||
fail.refresh_from_db()
|
||||
assert requeue.status == Job.Status.QUEUED
|
||||
assert fail.status == Job.Status.FAILED
|
||||
assert JobEvent.objects.filter(job=requeue, type=JobEvent.Type.TIMEOUT_REQUEUED).exists()
|
||||
assert JobEvent.objects.filter(job=fail, type=JobEvent.Type.FAILED).exists()
|
||||
|
||||
|
||||
@pytest.mark.django_db
|
||||
def test_stale_worker_cannot_complete_old_attempt():
|
||||
job = create_job(job_type="demo.timeout", max_attempts=3).job
|
||||
first = claim_next_job(worker_id="worker-a", lease_seconds=30)
|
||||
Job.objects.filter(id=first.id).update(locked_until=timezone.now() - timedelta(seconds=1))
|
||||
cleanup_expired_jobs(batch_size=10)
|
||||
Job.objects.filter(id=job.id).update(available_at=timezone.now() - timedelta(seconds=1))
|
||||
second = claim_next_job(worker_id="worker-b", lease_seconds=30)
|
||||
|
||||
stale_result = complete_job(job.id, worker_id="worker-a", attempt=1, result={"ok": True})
|
||||
|
||||
assert stale_result is None
|
||||
job.refresh_from_db()
|
||||
assert job.status == Job.Status.RUNNING
|
||||
assert job.locked_by == "worker-b"
|
||||
assert job.attempts == second.attempts == 2
|
||||
|
||||
|
||||
@pytest.mark.django_db
|
||||
def test_lease_renewal_only_current_owner_and_attempt():
|
||||
job = create_job(job_type="demo.slow").job
|
||||
claimed = claim_next_job(worker_id="worker-a", lease_seconds=5)
|
||||
old_locked_until = claimed.locked_until
|
||||
|
||||
assert renew_lease(job.id, worker_id="worker-b", attempt=1, lease_seconds=60) is None
|
||||
assert renew_lease(job.id, worker_id="worker-a", attempt=2, lease_seconds=60) is None
|
||||
renewed = renew_lease(job.id, worker_id="worker-a", attempt=1, lease_seconds=60)
|
||||
|
||||
assert renewed is not None
|
||||
assert renewed.locked_until > old_locked_until
|
||||
assert JobEvent.objects.filter(job=job, type=JobEvent.Type.LEASE_RENEWED).count() == 1
|
||||
|
||||
|
||||
@pytest.mark.django_db
|
||||
def test_manual_retry_resets_failed_job():
|
||||
job = create_job(job_type="demo.fail", max_attempts=1).job
|
||||
claimed = claim_next_job(worker_id="worker-a", lease_seconds=30)
|
||||
fail_job(claimed.id, worker_id="worker-a", attempt=1, error="boom")
|
||||
|
||||
retried = retry_failed_job(job.id)
|
||||
|
||||
assert retried.status == Job.Status.QUEUED
|
||||
assert retried.attempts == 0
|
||||
assert retried.finished_at is None
|
||||
assert JobEvent.objects.filter(job=job, type=JobEvent.Type.MANUAL_RETRY).exists()
|
||||
|
||||
|
||||
def test_runner_request_stop_prevents_more_claiming():
|
||||
runner = JobWorkerRunner(thread_count=1)
|
||||
assert runner.stop_event.is_set() is False
|
||||
|
||||
runner.request_stop()
|
||||
|
||||
assert runner.stop_event.is_set() is True
|
||||
19
backend/jobs/urls.py
Normal file
19
backend/jobs/urls.py
Normal file
@@ -0,0 +1,19 @@
|
||||
from django.urls import path
|
||||
|
||||
from jobs.views import (
|
||||
GlobalJobEventsAPIView,
|
||||
JobDetailAPIView,
|
||||
JobEventsAPIView,
|
||||
JobListCreateAPIView,
|
||||
JobStatsAPIView,
|
||||
RetryJobAPIView,
|
||||
)
|
||||
|
||||
urlpatterns = [
|
||||
path("jobs/", JobListCreateAPIView.as_view(), name="jobs"),
|
||||
path("jobs/stats/", JobStatsAPIView.as_view(), name="job-stats"),
|
||||
path("jobs/<uuid:pk>/", JobDetailAPIView.as_view(), name="job-detail"),
|
||||
path("jobs/<uuid:pk>/events/", JobEventsAPIView.as_view(), name="job-events"),
|
||||
path("jobs/<uuid:pk>/retry/", RetryJobAPIView.as_view(), name="job-retry"),
|
||||
path("job-events/", GlobalJobEventsAPIView.as_view(), name="global-job-events"),
|
||||
]
|
||||
96
backend/jobs/views.py
Normal file
96
backend/jobs/views.py
Normal file
@@ -0,0 +1,96 @@
|
||||
from django.db import connection
|
||||
from rest_framework import generics, status
|
||||
from rest_framework.exceptions import NotFound, ValidationError
|
||||
from rest_framework.response import Response
|
||||
from rest_framework.views import APIView
|
||||
|
||||
from jobs.models import Job, JobEvent
|
||||
from jobs.serializers import JobCreateSerializer, JobEventSerializer, JobSerializer
|
||||
from jobs.services import create_job, job_stats, retry_failed_job
|
||||
|
||||
|
||||
class HealthAPIView(APIView):
|
||||
def get(self, request):
|
||||
try:
|
||||
with connection.cursor() as cursor:
|
||||
cursor.execute("SELECT 1")
|
||||
cursor.fetchone()
|
||||
except Exception as exc:
|
||||
return Response(
|
||||
{"ok": False, "database": "error", "detail": str(exc)},
|
||||
status=status.HTTP_503_SERVICE_UNAVAILABLE,
|
||||
)
|
||||
return Response({"ok": True, "database": "ok"})
|
||||
|
||||
|
||||
class JobListCreateAPIView(APIView):
|
||||
def get(self, request):
|
||||
queryset = Job.objects.order_by("-created_at")
|
||||
status_filter = request.query_params.get("status")
|
||||
if status_filter:
|
||||
queryset = queryset.filter(status=status_filter)
|
||||
return Response(JobSerializer(queryset, many=True).data)
|
||||
|
||||
def post(self, request):
|
||||
serializer = JobCreateSerializer(data=request.data)
|
||||
serializer.is_valid(raise_exception=True)
|
||||
result = create_job(
|
||||
job_type=serializer.validated_data["type"],
|
||||
payload=serializer.validated_data.get("payload") or {},
|
||||
priority=serializer.validated_data.get("priority", 50),
|
||||
available_at=serializer.validated_data.get("available_at"),
|
||||
max_attempts=serializer.validated_data.get("max_attempts"),
|
||||
idempotency_key=serializer.validated_data.get("idempotency_key"),
|
||||
)
|
||||
return Response(
|
||||
JobSerializer(result.job).data,
|
||||
status=status.HTTP_201_CREATED if result.created else status.HTTP_200_OK,
|
||||
)
|
||||
|
||||
|
||||
class JobDetailAPIView(generics.RetrieveAPIView):
|
||||
queryset = Job.objects.all()
|
||||
serializer_class = JobSerializer
|
||||
|
||||
|
||||
class JobEventsAPIView(APIView):
|
||||
def get(self, request, pk):
|
||||
if not Job.objects.filter(id=pk).exists():
|
||||
raise NotFound("Job not found.")
|
||||
events = JobEvent.objects.filter(job_id=pk).order_by("id")
|
||||
return Response(JobEventSerializer(events, many=True).data)
|
||||
|
||||
|
||||
class GlobalJobEventsAPIView(APIView):
|
||||
def get(self, request):
|
||||
queryset = JobEvent.objects.order_by("id")
|
||||
after_id = request.query_params.get("after_id")
|
||||
event_type = request.query_params.get("type")
|
||||
job_id = request.query_params.get("job_id")
|
||||
if after_id:
|
||||
queryset = queryset.filter(id__gt=after_id)
|
||||
if event_type:
|
||||
queryset = queryset.filter(type=event_type)
|
||||
if job_id:
|
||||
queryset = queryset.filter(job_id=job_id)
|
||||
try:
|
||||
limit = min(int(request.query_params.get("limit", "100")), 500)
|
||||
except ValueError as exc:
|
||||
raise ValidationError({"limit": "Must be an integer."}) from exc
|
||||
return Response(JobEventSerializer(queryset[:limit], many=True).data)
|
||||
|
||||
|
||||
class JobStatsAPIView(APIView):
|
||||
def get(self, request):
|
||||
return Response(job_stats())
|
||||
|
||||
|
||||
class RetryJobAPIView(APIView):
|
||||
def post(self, request, pk):
|
||||
try:
|
||||
job = retry_failed_job(pk)
|
||||
except Job.DoesNotExist as exc:
|
||||
raise NotFound("Job not found.") from exc
|
||||
except ValueError as exc:
|
||||
raise ValidationError({"detail": str(exc)}) from exc
|
||||
return Response(JobSerializer(job).data)
|
||||
90
backend/jobs/worker.py
Normal file
90
backend/jobs/worker.py
Normal file
@@ -0,0 +1,90 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import socket
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
|
||||
from django.conf import settings
|
||||
from django.db import close_old_connections
|
||||
|
||||
from jobs.handlers import execute_handler
|
||||
from jobs.services import claim_next_job, cleanup_expired_jobs, complete_job, fail_job
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class JobWorkerRunner:
|
||||
def __init__(self, *, thread_count: int | None = None):
|
||||
self.thread_count = thread_count or settings.JOB_WORKER_THREADS
|
||||
self.stop_event = threading.Event()
|
||||
self.threads: list[threading.Thread] = []
|
||||
|
||||
def request_stop(self, *_args):
|
||||
logger.info("Worker shutdown requested; no new jobs will be claimed.")
|
||||
self.stop_event.set()
|
||||
|
||||
def start(self):
|
||||
for index in range(self.thread_count):
|
||||
thread = threading.Thread(target=self._run_thread, args=(index,), name=f"job-worker-{index + 1}")
|
||||
thread.start()
|
||||
self.threads.append(thread)
|
||||
|
||||
def wait(self):
|
||||
while any(thread.is_alive() for thread in self.threads):
|
||||
for thread in self.threads:
|
||||
thread.join(timeout=0.25)
|
||||
|
||||
def run_forever(self):
|
||||
self.start()
|
||||
self.wait()
|
||||
|
||||
def _worker_id(self, index: int) -> str:
|
||||
return f"{socket.gethostname()}:{os.getpid()}:thread-{index + 1}:{uuid.uuid4().hex[:8]}"
|
||||
|
||||
def _run_thread(self, index: int):
|
||||
worker_id = self._worker_id(index)
|
||||
logger.info("Worker thread started: %s", worker_id)
|
||||
last_cleanup = 0.0
|
||||
poll_seconds = settings.JOB_WORKER_POLL_INTERVAL_MS / 1000
|
||||
while not self.stop_event.is_set():
|
||||
close_old_connections()
|
||||
now = time.monotonic()
|
||||
if now - last_cleanup >= settings.JOB_WORKER_CLEANUP_INTERVAL_SECONDS:
|
||||
try:
|
||||
cleanup_expired_jobs(batch_size=settings.JOB_WORKER_CLEANUP_BATCH_SIZE)
|
||||
except Exception:
|
||||
logger.exception("Expired job cleanup failed.")
|
||||
last_cleanup = now
|
||||
|
||||
try:
|
||||
job = claim_next_job(worker_id=worker_id, lease_seconds=settings.JOB_WORKER_LEASE_SECONDS)
|
||||
except Exception:
|
||||
logger.exception("Job claim failed.")
|
||||
self.stop_event.wait(poll_seconds)
|
||||
continue
|
||||
|
||||
if job is None:
|
||||
self.stop_event.wait(poll_seconds)
|
||||
continue
|
||||
|
||||
self._execute_claimed_job(job, worker_id=worker_id)
|
||||
|
||||
close_old_connections()
|
||||
logger.info("Worker thread stopped: %s", worker_id)
|
||||
|
||||
def _execute_claimed_job(self, job, *, worker_id: str):
|
||||
attempt = job.attempts
|
||||
try:
|
||||
result = execute_handler(job, worker_id=worker_id)
|
||||
except Exception as exc:
|
||||
updated_job = fail_job(job.id, worker_id=worker_id, attempt=attempt, error=str(exc))
|
||||
if updated_job is None:
|
||||
logger.info("Worker lost ownership before failing job %s attempt %s.", job.id, attempt)
|
||||
return
|
||||
|
||||
updated_job = complete_job(job.id, worker_id=worker_id, attempt=attempt, result=result)
|
||||
if updated_job is None:
|
||||
logger.info("Worker lost ownership before completing job %s attempt %s.", job.id, attempt)
|
||||
Reference in New Issue
Block a user