91 lines
3.2 KiB
Python
91 lines
3.2 KiB
Python
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)
|