1from __future__ import annotations
2
3import datetime
4import gc
5import multiprocessing
6import os
7import socket
8import threading
9import time
10import traceback
11import uuid
12from concurrent.futures import Future, ProcessPoolExecutor, wait
13from concurrent.futures.process import BrokenProcessPool
14from functools import partial
15from typing import TYPE_CHECKING, Any
16
17from opentelemetry import trace
18from plain.logs import get_framework_logger
19from plain.postgres import transaction
20from plain.postgres.db import return_database_connection
21from plain.postgres.otel import suppress_db_tracing
22from plain.runtime import settings
23from plain.utils import timezone
24from plain.utils.module_loading import import_string
25from plain.utils.os import get_cpu_count
26
27from .otel import WorkerMetrics, error_consumer_span, record_span_error, tracer
28from .registry import jobs_registry
29
30if TYPE_CHECKING:
31 from .models import JobProcess, JobResult
32
33# Models are NOT imported at the top of this file!
34# See comment on _worker_process_initializer() for explanation.
35
36logger = get_framework_logger()
37
38
39def _worker_process_initializer() -> None:
40 """Initialize Plain framework in worker process before processing jobs.
41
42 Why this is needed:
43 - We use multiprocessing with 'spawn' context (not 'fork')
44 - Spawn creates fresh Python processes, not forked copies
45 - When a spawned process starts, it re-imports this module BEFORE the initializer runs
46 - If we imported models at the top of this file, model registration would
47 happen before plain.runtime.setup(), causing PackageRegistryNotReady errors
48
49 Solution:
50 - This initializer runs plain.runtime.setup() FIRST in each worker process
51 - All model imports happen lazily inside functions (after setup completes)
52 - This ensures packages registry is ready before any models are accessed
53
54 Execution order in spawned worker:
55 1. Re-import workers.py (but models NOT imported yet - lazy!)
56 2. Run this initializer → plain.runtime.setup()
57 3. Execute process_job() → NOW it's safe to import models
58 """
59 from plain.runtime import setup
60
61 # Each spawned worker process needs to set up Plain
62 # (spawn context creates fresh processes, not forks)
63 setup()
64
65
66class Worker:
67 def __init__(
68 self,
69 queues: list[str],
70 jobs_schedule: list[Any] | None = None,
71 max_processes: int | None = None,
72 max_jobs_per_process: int | None = None,
73 max_pending_per_process: int = 10,
74 stats_every: int | None = None,
75 ) -> None:
76 if jobs_schedule is None:
77 jobs_schedule = []
78
79 if max_processes is None:
80 max_processes = get_cpu_count()
81
82 self.executor = ProcessPoolExecutor(
83 max_workers=max_processes,
84 max_tasks_per_child=max_jobs_per_process,
85 mp_context=multiprocessing.get_context("spawn"),
86 initializer=_worker_process_initializer,
87 )
88
89 self.queues = queues
90
91 # Filter the jobs schedule to those that are in the same queue as this worker
92 self.jobs_schedule = [
93 x for x in jobs_schedule if x[0].default_queue() in queues
94 ]
95
96 # How often to log the stats (in seconds)
97 self.stats_every = stats_every
98
99 self.max_processes = self.executor._max_workers # ty: ignore[unresolved-attribute]
100 self.max_jobs_per_process = max_jobs_per_process
101 self.max_pending_per_process = max_pending_per_process
102
103 self._is_shutting_down = False
104
105 # Maintenance baselines — each task runs when its interval has
106 # elapsed since these, so construction counts as the starting point.
107 now = time.time()
108 self._stats_logged_at = now
109 self._job_results_checked_at = now
110 self._jobs_schedule_checked_at = now
111
112 self.worker_id = uuid.uuid4()
113 self._hostname = socket.gethostname()
114 self._pid = os.getpid()
115 self._heartbeat_at = 0.0
116 # We refuse to claim JobRequests until a WorkerHeartbeat row exists
117 # for our worker_id. Otherwise an unregistered worker could stamp
118 # JobProcess rows with a worker_id that has no heartbeat row to ever
119 # match — rescue would never find them.
120 self._heartbeat_registered = False
121
122 # Track our own in-flight futures so the shutdown drain doesn't depend
123 # on ProcessPoolExecutor's private _pending_work_items dict (which
124 # isn't guaranteed to drain cleanly when cancel_futures=True is used).
125 # The mapping (Future → JobProcess.uuid) also lets _rescue_own_orphans
126 # reconcile DB rows against currently-tracked futures. Done callbacks
127 # fire from the executor's management thread, so all access goes
128 # through _inflight_lock to keep iteration safe.
129 self._inflight_futures: dict[Future, str] = {}
130 self._inflight_lock = threading.Lock()
131
132 self.metrics = WorkerMetrics(self)
133
134 def run(self) -> None:
135 logger.info(
136 "Starting Plain worker",
137 extra={
138 "registered_jobs": list(jobs_registry.jobs.keys()),
139 "queues": list(self.queues),
140 "jobs_schedule": [str(x) for x in self.jobs_schedule],
141 "stats_every": self.stats_every,
142 "max_processes": self.max_processes,
143 "max_jobs_per_process": self.max_jobs_per_process,
144 "max_pending_per_process": self.max_pending_per_process,
145 "pid": self._pid,
146 "worker_id": str(self.worker_id),
147 },
148 )
149
150 # Heartbeat writes here run outside any entry span — suppress their
151 # DB tracing so they don't export as single-span root traces.
152 with suppress_db_tracing():
153 self.register_heartbeat()
154 self._run_loop()
155 self._drain_with_heartbeat()
156 # Only reached on clean exit. On error/interrupt, control unwinds past
157 # this and the heartbeat row is left to go stale — rescue then picks
158 # up our in-flight jobs as LOST. Deleting the row here on error would
159 # lie about being alive and strand any JobProcess rows still stamped
160 # with this worker_id.
161 with suppress_db_tracing():
162 self.deregister_heartbeat()
163
164 def _discard_inflight(self, future: Future) -> None:
165 with self._inflight_lock:
166 self._inflight_futures.pop(future, None)
167
168 def _drain_with_heartbeat(self) -> None:
169 """Wait for in-flight jobs to finish while keeping our heartbeat alive.
170
171 Without continuing to heartbeat during drain, a long-running job
172 (e.g. multi-minute LLM turns) would let our row go stale and trigger
173 false-positive LOST conversions from another worker's rescue tick.
174
175 Drain is unbounded — a stuck job will block here until the platform
176 sends SIGKILL (Heroku's grace period, k8s terminationGracePeriod, etc.).
177 At that point the heartbeat goes stale and rescue picks up the orphans.
178 """
179 self.executor.shutdown(wait=False, cancel_futures=True)
180 while True:
181 # Snapshot under the lock — done callbacks mutate the set from
182 # the executor's management thread.
183 with self._inflight_lock:
184 snapshot = list(self._inflight_futures.keys())
185 if not snapshot:
186 break
187 # Sleep up to 1s, waking early if any future completes.
188 wait(snapshot, timeout=1)
189 try:
190 # No entry span is active during drain — suppress the
191 # heartbeat's DB tracing so each drain tick doesn't export
192 # a single-span root trace.
193 with suppress_db_tracing():
194 self.maybe_heartbeat()
195 except Exception:
196 logger.exception("Worker heartbeat failed during drain")
197 logger.info("Job worker shutdown complete")
198
199 def _run_loop(self) -> None:
200 consecutive_claim_failures = 0
201
202 while not self._is_shutting_down:
203 # Return last tick's pooled connection so the next checkout
204 # re-validates it — a server-side close (PG restart, failover)
205 # is then caught at checkout instead of erroring a tick — and so
206 # the loop doesn't hold a pool slot while idle between ticks.
207 return_database_connection()
208
209 # Only open the span when maintenance will actually run — a
210 # fully-idle tick exports no telemetry at all, instead of a
211 # single-span root trace per second per worker. A new maybe_*
212 # call here needs its due-predicate added to _maintenance_due().
213 if self._maintenance_due():
214 with tracer.start_as_current_span(
215 "worker loop", kind=trace.SpanKind.CONSUMER
216 ) as span:
217 try:
218 self.maybe_heartbeat()
219 self.maybe_log_stats()
220 self.maybe_check_job_results()
221 self.maybe_schedule_jobs()
222 except Exception as e:
223 # The catch is inside the span, so the SDK's auto-record
224 # on context exit won't fire — stamp the canonical
225 # failure signal explicitly. Log and continue: these
226 # tasks are ancillary to the main job processing.
227 record_span_error(span, e)
228 logger.exception("Worker maintenance task failed")
229
230 # Re-check shutdown after maintenance — a signal may have arrived
231 # between the loop condition and now. Don't pick up new work.
232 if self._is_shutting_down:
233 break
234
235 if not self._heartbeat_registered:
236 time.sleep(1)
237 continue
238
239 if len(self._inflight_futures) >= (
240 self.max_processes * self.max_pending_per_process
241 ):
242 # We don't want to convert too many JobRequests to Jobs,
243 # because anything not started yet will be cancelled on deploy etc.
244 # It's easier to leave them in the JobRequest db queue as long as possible.
245 time.sleep(0.5)
246 continue
247
248 try:
249 job = self._claim_job()
250 consecutive_claim_failures = 0
251 except Exception as e:
252 # A transient DB failure while claiming shouldn't kill the
253 # worker. With the claim's CLIENT spans suppressed there is
254 # no entry span to carry the failure, so emit one and log
255 # inside it so the record shares its trace.
256 with error_consumer_span(name="claim job", exc=e):
257 logger.exception("Failed to claim job")
258 consecutive_claim_failures += 1
259 if consecutive_claim_failures >= 30:
260 # This isn't a blip (e.g. schema drift or lost table
261 # permissions while the heartbeat table stays healthy).
262 # Crash so the supervisor restarts us visibly instead of
263 # looping forever while reporting a fresh heartbeat.
264 raise
265 time.sleep(1)
266 continue
267
268 if job is None:
269 # Potentially no jobs to process (who knows for how long)
270 # but sleep for a second to give the CPU and DB a break
271 time.sleep(1)
272 continue
273
274 # Signal may have fired during the DB queries above. Don't submit
275 # new work past shutdown — revert the JobProcess back to a
276 # JobRequest so the next worker generation picks it up.
277 if self._is_shutting_down:
278 # No entry span here either — same suppression as the claim.
279 with suppress_db_tracing():
280 job.revert_to_job_request()
281 break
282
283 job_process_uuid = str(job.uuid) # Make a str copy
284
285 try:
286 future = self.executor.submit(process_job, job_process_uuid)
287 with self._inflight_lock:
288 self._inflight_futures[future] = job_process_uuid
289 # If the future is already done, add_done_callback runs the
290 # callback inline on THIS thread — and its finally returns
291 # this thread's connection. Safe here because no atomic block
292 # or cursor is open at this point; keep it that way.
293 future.add_done_callback(
294 partial(future_finished_callback, job_process_uuid)
295 )
296 future.add_done_callback(self._discard_inflight)
297 except (BrokenProcessPool, RuntimeError):
298 # BrokenProcessPool: child OOM, segfault, or other crash.
299 # RuntimeError: executor already shut down (shutdown race).
300 # Either way, the job was already converted from JobRequest
301 # to JobProcess, so re-enqueue it before exiting.
302 logger.warning(
303 "Process pool broken, re-enqueuing job",
304 extra={"job_process_uuid": job_process_uuid},
305 )
306 # No entry span here either — same suppression as the claim.
307 with suppress_db_tracing():
308 job.revert_to_job_request()
309 break
310
311 def _claim_job(self) -> JobProcess | None:
312 """Atomically claim the next ready JobRequest as a JobProcess, or
313 return None when the queues are empty.
314
315 The poll is framework housekeeping that fires every ~1s per worker,
316 outside any entry span — untraced, its CLIENT spans would each export
317 as a single-span root trace. Suppress the whole claim transaction;
318 the meaningful telemetry for a claimed job is the `process {queue}`
319 CONSUMER span emitted later by JobProcess.run().
320 """
321 # Lazy import - see _worker_process_initializer() comment for why
322 from .models import JobRequest
323
324 with suppress_db_tracing(), transaction.atomic():
325 job_request = (
326 JobRequest.query.ready_to_run()
327 .filter(queue__in=self.queues)
328 .select_for_update(skip_locked=True)
329 .order_by("-priority", "-start_at", "-created_at")
330 .first()
331 )
332 if not job_request:
333 return None
334
335 logger.debug(
336 "Preparing to execute job",
337 extra={
338 "job_class": job_request.job_class,
339 "job_request_uuid": job_request.uuid,
340 "job_priority": job_request.priority,
341 "job_source": job_request.source,
342 "job_queue": job_request.queue,
343 },
344 )
345 return job_request.convert_to_job_process(worker_id=self.worker_id)
346
347 def shutdown(self) -> None:
348 if self._is_shutting_down:
349 # Already shutting down somewhere else
350 return
351
352 logger.info("Job worker shutdown requested")
353 # Just flip the flag — drain happens in _drain_with_heartbeat() so
354 # heartbeats keep firing during it. Blocking the signal handler with
355 # executor.shutdown(wait=True) here would let our row go stale.
356 self._is_shutting_down = True
357
358 def _maintenance_due(self) -> bool:
359 """At least one maybe_* task will do real work this tick.
360
361 Checked before opening the `worker loop` span so a fully-idle tick
362 emits no telemetry at all. Any task that is due gets wrapped in the
363 span — it's the OTel error-attribution entry span for maintenance.
364
365 Keep these predicates in lockstep with the maybe_* calls in
366 _run_loop — a task missing here only runs when another task happens
367 to be due.
368 """
369 now = time.time()
370 return (
371 self._heartbeat_due(now)
372 or self._stats_due(now)
373 or self._job_results_check_due(now)
374 or self._schedule_due(now)
375 )
376
377 def _stats_due(self, now: float) -> bool:
378 if not self.stats_every:
379 return False
380 return now - self._stats_logged_at > self.stats_every
381
382 def maybe_log_stats(self) -> None:
383 now = time.time()
384 if not self._stats_due(now):
385 return
386 self._stats_logged_at = now
387 self.log_stats()
388
389 def _job_results_check_due(self, now: float) -> bool:
390 # Only need to check once a minute
391 return now - self._job_results_checked_at > 60
392
393 def maybe_check_job_results(self) -> None:
394 now = time.time()
395 if not self._job_results_check_due(now):
396 return
397 self._job_results_checked_at = now
398 self.rescue_job_results()
399
400 def _create_heartbeat_row(self) -> None:
401 # Lazy import - see _worker_process_initializer() comment for why
402 from .models import WorkerHeartbeat
403
404 WorkerHeartbeat.query.create(
405 worker_id=self.worker_id,
406 hostname=self._hostname,
407 pid=self._pid,
408 queues=list(self.queues),
409 last_heartbeat_at=timezone.now(),
410 )
411
412 def _refresh_heartbeat(self) -> None:
413 # Lazy import - see _worker_process_initializer() comment for why
414 from .models import WorkerHeartbeat
415
416 updated = WorkerHeartbeat.query.filter(worker_id=self.worker_id).update(
417 last_heartbeat_at=timezone.now()
418 )
419 if not updated:
420 # Row was deleted — registration failed earlier, or another
421 # rescuer claimed us as dead. Recreate so we're discoverable.
422 logger.warning(
423 "Worker heartbeat row missing, re-registering",
424 extra={"worker_id": str(self.worker_id)},
425 )
426 self._create_heartbeat_row()
427
428 def register_heartbeat(self) -> None:
429 try:
430 self._create_heartbeat_row()
431 self._heartbeat_at = time.time()
432 self._heartbeat_registered = True
433 except Exception as e:
434 # Registration failure is non-fatal — maybe_heartbeat will retry.
435 # Until it succeeds, _heartbeat_registered stays False and the run
436 # loop won't claim work. That's serious enough to deserve error
437 # attribution, and there's no entry span here to carry it.
438 with error_consumer_span(name="worker heartbeat", exc=e):
439 logger.exception("Worker registration failed")
440 logger.warning(
441 "Worker heartbeat registration failed; worker will not claim "
442 "jobs until a heartbeat row is created",
443 extra={"worker_id": str(self.worker_id)},
444 )
445
446 def _heartbeat_due(self, now: float) -> bool:
447 return (
448 not self._heartbeat_registered
449 or now - self._heartbeat_at >= settings.JOBS_HEARTBEAT_INTERVAL
450 )
451
452 def maybe_heartbeat(self) -> None:
453 now = time.time()
454 if not self._heartbeat_due(now):
455 return
456
457 try:
458 self._refresh_heartbeat()
459 self._heartbeat_at = now
460 self._heartbeat_registered = True
461 except Exception as e:
462 # We don't know if the row exists (the update may have returned 0
463 # and the recreate may have raised). Mark unregistered so the run
464 # loop stops claiming work until the next tick succeeds. If the
465 # DB is unreachable for long enough, our row goes stale and
466 # rescue marks our jobs LOST — that's the intended behavior.
467 # The catch swallows the exception, so stamp the failure on its
468 # own CONSUMER span (logging inside it so the record shares its
469 # trace) — otherwise a heartbeat outage exports only
470 # healthy-looking spans (or, during drain, nothing at all).
471 self._heartbeat_registered = False
472 with error_consumer_span(name="worker heartbeat", exc=e):
473 logger.exception("Worker heartbeat failed")
474
475 def deregister_heartbeat(self) -> None:
476 # Lazy import - see _worker_process_initializer() comment for why
477 from .models import JobProcess, WorkerHeartbeat
478
479 try:
480 # If any JobProcess rows still reference this worker_id, a
481 # bookkeeping error during drain (e.g. future_finished_callback's
482 # own convert_to_result raised) left them stranded. Don't delete
483 # the heartbeat — let it go stale so rescue can pick them up.
484 if JobProcess.query.filter(worker_id=self.worker_id).exists():
485 logger.warning(
486 "Worker has remaining JobProcess rows at shutdown; "
487 "leaving heartbeat for rescue to claim",
488 extra={"worker_id": str(self.worker_id)},
489 )
490 return
491
492 WorkerHeartbeat.query.filter(worker_id=self.worker_id).delete()
493 except Exception:
494 # Best effort. A leftover row will be reclaimed by rescue when its
495 # heartbeat goes stale.
496 logger.exception("Failed to remove worker heartbeat")
497
498 def _schedule_due(self, now: float) -> bool:
499 if not self.jobs_schedule:
500 return False
501 # Only need to check once every 60 seconds
502 return now - self._jobs_schedule_checked_at > 60
503
504 def maybe_schedule_jobs(self) -> None:
505 if not self._schedule_due(time.time()):
506 return
507
508 for job, schedule in self.jobs_schedule:
509 next_start_at = schedule.next()
510
511 # Leverage the concurrency_key to group scheduled jobs
512 # with the same start time
513 schedule_concurrency_key = f"{job.default_concurrency_key()}:scheduled:{int(next_start_at.timestamp())}"
514
515 # Job's should_enqueue hook can control scheduling behavior
516 result = job.run_in_worker(
517 delay=next_start_at,
518 concurrency_key=schedule_concurrency_key,
519 )
520 # Result is None if should_enqueue returned False
521 if result:
522 logger.info(
523 "Scheduling job",
524 extra={
525 "job_class": result.job_class,
526 "job_queue": result.queue,
527 "job_start_at": result.start_at,
528 "job_schedule": schedule,
529 "concurrency_key": result.concurrency_key,
530 },
531 )
532
533 # Stamp only after the whole pass succeeds — a mid-pass failure
534 # (caught by the worker-loop catch) then retries next tick instead
535 # of waiting out the window and skipping the missed occurrence.
536 # Re-runs are deduped by the scheduled concurrency_key.
537 self._jobs_schedule_checked_at = time.time()
538
539 def log_stats(self) -> None:
540 # Lazy import - see _worker_process_initializer() comment for why
541 from .models import JobProcess, JobRequest
542
543 try:
544 num_proccesses = len(self.executor._processes)
545 except (AttributeError, TypeError):
546 # Depending on shutdown timing and internal behavior, this might not work
547 num_proccesses = 0
548
549 jobs_requested = JobRequest.query.filter(queue__in=self.queues).count()
550 jobs_processing = JobProcess.query.filter(queue__in=self.queues).count()
551
552 logger.info(
553 "Job worker stats",
554 extra={
555 "worker_processes": num_proccesses,
556 "worker_queues": ",".join(self.queues),
557 "jobs_requested": jobs_requested,
558 "jobs_processing": jobs_processing,
559 "worker_max_processes": self.max_processes,
560 "worker_max_jobs_per_process": self.max_jobs_per_process,
561 },
562 )
563
564 def rescue_job_results(self) -> None:
565 """Find any lost or failed jobs on this worker's queues and handle them.
566
567 Hooks are dispatched with maybe_heartbeat() interleaved so a slow or
568 large batch of on_aborted hooks can't starve our heartbeat and trigger
569 a false-positive LOST from a peer's rescue tick.
570 """
571 # Lazy import - see _worker_process_initializer() comment for why
572 from .models import JobResult, rescue_stale_workers
573
574 # rescue_stale_workers is global, not queue-scoped — a dead worker's
575 # heartbeat going stale is a global signal, and partial conversion
576 # would strand jobs.
577 global_hooks = rescue_stale_workers()
578 own_hooks = self._rescue_own_orphans()
579 self._dispatch_aborted_hooks(global_hooks + own_hooks)
580 JobResult.query.filter(queue__in=self.queues).retry_failed_jobs()
581
582 def _dispatch_aborted_hooks(self, results: list[JobResult]) -> None:
583 for result in results:
584 self.maybe_heartbeat()
585 result.dispatch_aborted_hook()
586
587 def _rescue_own_orphans(self) -> list[JobResult]:
588 """Convert any of our own JobProcess rows that aren't tracked by a
589 live future to LOST.
590
591 These shouldn't normally exist. The path that creates them: a future
592 completes, future_finished_callback runs, but convert_to_result raises
593 (transient DB error, peer-rescuer constraint conflict, etc.). The
594 exception escapes the callback into concurrent.futures (which logs and
595 moves on), _discard_inflight still fires from the second callback, and
596 the row is left in the DB with no path back — our heartbeat is fresh
597 so rescue_stale_workers won't see it.
598
599 Age threshold avoids racing with newly-claimed rows that haven't been
600 added to _inflight_futures yet (microsecond window between
601 convert_to_job_process and the dict insert in _run_loop).
602
603 Returns JobResults whose on_aborted hook the caller should dispatch.
604 """
605 from .models import JobProcess, JobResultStatuses
606
607 with self._inflight_lock:
608 inflight_uuids = list(self._inflight_futures.values())
609
610 cutoff = timezone.now() - datetime.timedelta(
611 seconds=settings.JOBS_HEARTBEAT_TIMEOUT
612 )
613 stranded = JobProcess.query.filter(
614 worker_id=self.worker_id,
615 created_at__lt=cutoff,
616 ).exclude(uuid__in=inflight_uuids)
617
618 pending_hooks: list[JobResult] = []
619 for orphan in list(stranded):
620 try:
621 result = orphan.convert_to_result(
622 status=JobResultStatuses.LOST,
623 error="JobProcess stranded — done-callback failed during conversion",
624 fire_hook=False,
625 )
626 except Exception:
627 logger.exception(
628 "Failed to rescue own orphan JobProcess",
629 extra={"job_process_uuid": str(orphan.uuid)},
630 )
631 continue
632 pending_hooks.append(result)
633 return pending_hooks
634
635
636def future_finished_callback(job_process_uuid: str, future: Future) -> None:
637 # Lazy import - see _worker_process_initializer() comment for why
638 from .models import JobProcess, JobResultStatuses
639
640 # This callback runs on the executor's done-callback thread with no entry
641 # span active, so suppress DB tracing — otherwise the orphan-check query
642 # on every completed job (and the conversions on cancel/failure) each
643 # export as a single-span root trace, scaling with job volume. The
644 # suppression is for framework bookkeeping only: Job.on_aborted is user
645 # code, so its dispatch is deferred to after the suppressed block
646 # (fire_hook=False here, dispatch_aborted_hook below).
647 aborted_result: JobResult | None = None
648 try:
649 with suppress_db_tracing():
650 if future.cancelled():
651 logger.warning(
652 "Job cancelled", extra={"job_process_uuid": job_process_uuid}
653 )
654 try:
655 job = JobProcess.query.get(uuid=job_process_uuid)
656 aborted_result = job.convert_to_result(
657 status=JobResultStatuses.CANCELLED, fire_hook=False
658 )
659 except JobProcess.DoesNotExist:
660 # Job may have already been cleaned up
661 pass
662 elif exception := future.exception():
663 # Process pool may have been killed (OOM/segfault), or process_job
664 # itself raised past its outer except (e.g. import failure).
665 logger.warning(
666 "Job failed",
667 extra={"job_process_uuid": job_process_uuid},
668 exc_info=exception,
669 )
670 try:
671 job = JobProcess.query.get(uuid=job_process_uuid)
672 # If started_at is set, run() was actively executing when the
673 # process died — user code may have set up state it expected to
674 # tear down. Use LOST so on_aborted fires. If started_at is
675 # unset, run() never got to execute (import failure, etc.), so
676 # ERRORED with no hook is correct.
677 if job.started_at is not None:
678 status = JobResultStatuses.LOST
679 else:
680 status = JobResultStatuses.ERRORED
681 result = job.convert_to_result(
682 status=status,
683 error="".join(traceback.format_exception(exception)),
684 fire_hook=False,
685 )
686 if status == JobResultStatuses.LOST:
687 aborted_result = result
688 except JobProcess.DoesNotExist:
689 # Job may have already been cleaned up
690 pass
691 else:
692 logger.debug(
693 "Job finished", extra={"job_process_uuid": job_process_uuid}
694 )
695 # Orphan check: process_job's outer except-Exception swallows any
696 # failure that escapes job.run() (middleware crash, OTel error, DB
697 # blip during convert_to_result, etc.). The future completes cleanly
698 # but the JobProcess row was never converted, and since our parent
699 # is still heartbeating, rescue_stale_workers won't see it as orphaned.
700 job = JobProcess.query.filter(uuid=job_process_uuid).first()
701 if job is None:
702 return
703 logger.warning(
704 "Job future completed but JobProcess survived; converting to ERRORED",
705 extra={"job_process_uuid": job_process_uuid},
706 )
707 try:
708 job.convert_to_result(
709 status=JobResultStatuses.ERRORED,
710 error="Job future completed without recording a result",
711 )
712 except Exception:
713 # A peer rescuer may have already created a JobResult(LOST) for
714 # this row, in which case the unique constraint on
715 # JobResult.job_process_uuid trips. Either way, the row is now
716 # accounted for — log and move on rather than letting the
717 # exception escape into the executor's done-callback machinery.
718 logger.exception(
719 "Failed to convert orphan JobProcess to ERRORED",
720 extra={"job_process_uuid": job_process_uuid},
721 )
722
723 if aborted_result is not None:
724 aborted_result.dispatch_aborted_hook()
725 finally:
726 # The done-callback thread gets its own thread-local pooled
727 # connection. Return it after each callback for the same reasons
728 # the run loop returns its connection every tick: checkout
729 # re-validation, and not holding a pool slot while idle.
730 return_database_connection()
731
732
733def process_job(job_process_uuid: str) -> None:
734 # Lazy import - see _worker_process_initializer() comment for why
735 from .models import JobProcess
736
737 try:
738 worker_pid = os.getpid()
739
740 job_process = JobProcess.query.get(uuid=job_process_uuid)
741
742 logger.info(
743 "Executing job",
744 extra={
745 "worker_pid": worker_pid,
746 "job_class": job_process.job_class,
747 "job_request_uuid": job_process.job_request_uuid,
748 "job_priority": job_process.priority,
749 "job_source": job_process.source,
750 "job_queue": job_process.queue,
751 },
752 )
753
754 def middleware_chain(job: JobProcess) -> JobResult:
755 return job.run()
756
757 for middleware_path in reversed(settings.JOBS_MIDDLEWARE):
758 middleware_class = import_string(middleware_path)
759 middleware_instance = middleware_class(middleware_chain)
760 middleware_chain = middleware_instance.process_job
761
762 job_result = middleware_chain(job_process)
763
764 assert job_result.ended_at is not None
765 assert job_result.started_at is not None
766 duration = job_result.ended_at - job_result.started_at
767 duration = duration.total_seconds()
768
769 if job_result.requested_at and job_result.started_at:
770 queue_time = (
771 job_result.started_at - job_result.requested_at
772 ).total_seconds()
773 else:
774 queue_time = None
775
776 logger.info(
777 "Completed job",
778 extra={
779 "worker_pid": worker_pid,
780 "job_class": job_result.job_class,
781 "job_process_uuid": job_result.job_process_uuid,
782 "job_request_uuid": job_result.job_request_uuid,
783 "job_result_uuid": job_result.uuid,
784 "job_priority": job_result.priority,
785 "job_source": job_result.source,
786 "job_queue": job_result.queue,
787 "job_duration": duration,
788 "job_queue_time": queue_time,
789 },
790 )
791 except Exception as e:
792 # Raising exceptions inside the worker process doesn't seem to be
793 # caught/shown anywhere as configured, so log it here. (A job
794 # catches its own user-code errors — this is for library errors:
795 # a failed JobProcess lookup, middleware that won't import or
796 # construct, or an error escaping run().) None of those has a
797 # *live* entry span, so stamp the failure on a one-off CONSUMER
798 # span and log inside it so the record carries its trace ids.
799 # For the rare library error escaping run(), run()'s own span
800 # already recorded the failure — the deliberate cost of this
801 # catch-all is that such an error reports on both spans, in
802 # exchange for the log never exporting span-less.
803 with error_consumer_span(name="process job", exc=e):
804 logger.exception("Job process errored")
805 finally:
806 return_database_connection()
807 gc.collect()