1from __future__ import annotations
2
3import datetime
4import time
5import traceback
6from typing import TYPE_CHECKING, Any, Self
7from uuid import UUID
8
9from opentelemetry.semconv._incubating.attributes.messaging_attributes import (
10 MESSAGING_CONSUMER_GROUP_NAME,
11 MESSAGING_MESSAGE_ID,
12 MESSAGING_OPERATION_NAME,
13)
14from opentelemetry.semconv.attributes.error_attributes import ERROR_TYPE
15from opentelemetry.trace import Link, SpanContext, SpanKind, TraceFlags
16from plain.logs import get_framework_logger
17from plain.postgres import transaction, types
18from plain.postgres.expressions import F
19from plain.runtime import settings
20from plain.utils import timezone
21
22from plain import postgres
23
24from .exceptions import DeferJob
25from .otel import (
26 operation_duration_histogram,
27 process_metric_attributes,
28 queue_wait_duration_histogram,
29 record_consumed,
30 record_span_error,
31 tracer,
32)
33from .registry import jobs_registry
34
35if TYPE_CHECKING:
36 from .jobs import Job
37
38__all__ = [
39 "JobProcess",
40 "JobRequest",
41 "JobResult",
42 "JobResultStatuses",
43 "WorkerHeartbeat",
44]
45
46logger = get_framework_logger()
47
48
49class JobRequestQuerySet(postgres.QuerySet["JobRequest"]):
50 def ready_to_run(self) -> Self:
51 """JobRequests with no scheduling constraint or whose `start_at` is past."""
52 return self.filter(
53 postgres.Q(start_at__isnull=True) | postgres.Q(start_at__lte=timezone.now())
54 )
55
56 def scheduled(self) -> Self:
57 """JobRequests scheduled to start in the future."""
58 return self.filter(start_at__gt=timezone.now())
59
60
61@postgres.register_model
62class JobRequest(postgres.Model):
63 """
64 Keep all pending job requests in a single table.
65 """
66
67 created_at = types.DateTimeField(create_now=True)
68 uuid = types.UUIDField(generate=True)
69
70 job_class = types.TextField(max_length=255)
71 parameters: dict[str, Any] | None = types.JSONField(required=False, allow_null=True)
72 priority = types.SmallIntegerField(default=0)
73 source = types.TextField(required=False, default="")
74 queue = types.TextField(default="default", max_length=255)
75
76 retries = types.SmallIntegerField(default=0)
77 retry_attempt = types.SmallIntegerField(default=0)
78
79 concurrency_key = types.TextField(max_length=255, required=False, default="")
80
81 start_at = types.DateTimeField(required=False, allow_null=True)
82
83 # OpenTelemetry trace context
84 trace_id = types.TextField(max_length=34, required=False, allow_null=True)
85 span_id = types.TextField(max_length=18, required=False, allow_null=True)
86
87 # expires_at = postgres.DateTimeField(required=False, allow_null=True)
88
89 query: JobRequestQuerySet = JobRequestQuerySet()
90
91 model_options = postgres.Options(
92 ordering=["-priority", "-created_at"],
93 indexes=[
94 postgres.Index(
95 name="plainjobs_jobrequest_priority_idx", fields=["priority"]
96 ),
97 postgres.Index(
98 name="plainjobs_jobrequest_created_at_idx", fields=["created_at"]
99 ),
100 postgres.Index(name="plainjobs_jobrequest_queue_idx", fields=["queue"]),
101 postgres.Index(
102 name="plainjobs_jobrequest_start_at_idx", fields=["start_at"]
103 ),
104 postgres.Index(
105 name="plainjobs_jobrequest_concurrency_key_idx",
106 fields=["concurrency_key"],
107 ),
108 # Used for job grouping queries
109 postgres.Index(
110 name="job_request_concurrency_key",
111 fields=["job_class", "concurrency_key"],
112 ),
113 ],
114 constraints=[
115 postgres.UniqueConstraint(
116 fields=["uuid"], name="plainjobs_jobrequest_unique_uuid"
117 ),
118 ],
119 )
120
121 def __str__(self) -> str:
122 return f"{self.job_class} [{self.uuid}]"
123
124 def convert_to_job_process(self, *, worker_id: UUID) -> JobProcess:
125 """
126 JobRequests are the pending jobs that are waiting to be executed.
127 We immediately convert them to JobProcess when they are picked up.
128
129 worker_id stamps ownership: rescue_stale_workers uses it to find
130 which jobs belonged to a worker whose heartbeat went stale. Required —
131 every JobProcess has an owning worker, and the NOT NULL column
132 constraint is what stops pre-heartbeat workers from inserting
133 unrescuable rows during a rolling upgrade.
134 """
135 with transaction.atomic():
136 result = JobProcess.query.create(
137 job_request_uuid=self.uuid,
138 requested_at=self.created_at,
139 job_class=self.job_class,
140 parameters=self.parameters,
141 priority=self.priority,
142 source=self.source,
143 queue=self.queue,
144 retries=self.retries,
145 retry_attempt=self.retry_attempt,
146 concurrency_key=self.concurrency_key,
147 trace_id=self.trace_id,
148 span_id=self.span_id,
149 worker_id=worker_id,
150 )
151
152 # Delete the pending JobRequest now
153 self.delete()
154
155 return result
156
157
158class JobQuerySet(postgres.QuerySet["JobProcess"]):
159 def running(self) -> Self:
160 return self.filter(started_at__isnull=False)
161
162 def waiting(self) -> Self:
163 return self.filter(started_at__isnull=True)
164
165
166@postgres.register_model
167class JobProcess(postgres.Model):
168 """
169 All active jobs are stored in this table.
170 """
171
172 uuid = types.UUIDField(generate=True)
173 created_at = types.DateTimeField(create_now=True)
174 started_at = types.DateTimeField(required=False, allow_null=True)
175
176 # From the JobRequest
177 job_request_uuid = types.UUIDField()
178 requested_at = types.DateTimeField(required=False, allow_null=True)
179 job_class = types.TextField(max_length=255)
180 parameters: dict[str, Any] | None = types.JSONField(required=False, allow_null=True)
181 priority = types.SmallIntegerField(default=0)
182 source = types.TextField(required=False, default="")
183 queue = types.TextField(default="default", max_length=255)
184 retries = types.SmallIntegerField(default=0)
185 retry_attempt = types.SmallIntegerField(default=0)
186 concurrency_key = types.TextField(max_length=255, required=False, default="")
187
188 # OpenTelemetry trace context
189 trace_id = types.TextField(max_length=34, required=False, allow_null=True)
190 span_id = types.TextField(max_length=18, required=False, allow_null=True)
191
192 worker_id = types.UUIDField()
193
194 query: JobQuerySet = JobQuerySet()
195
196 model_options = postgres.Options(
197 ordering=["-created_at"],
198 indexes=[
199 postgres.Index(
200 name="plainjobs_jobprocess_created_at_idx", fields=["created_at"]
201 ),
202 postgres.Index(name="plainjobs_jobprocess_queue_idx", fields=["queue"]),
203 postgres.Index(
204 name="plainjobs_jobprocess_concurrency_key_idx",
205 fields=["concurrency_key"],
206 ),
207 postgres.Index(
208 name="plainjobs_jobprocess_started_at_idx", fields=["started_at"]
209 ),
210 postgres.Index(
211 name="plainjobs_jobprocess_job_request_uuid_idx",
212 fields=["job_request_uuid"],
213 ),
214 postgres.Index(
215 name="plainjobs_jobprocess_worker_id_idx", fields=["worker_id"]
216 ),
217 # Used for job grouping queries
218 postgres.Index(
219 name="job_concurrency_key",
220 fields=["job_class", "concurrency_key"],
221 ),
222 ],
223 constraints=[
224 postgres.UniqueConstraint(
225 fields=["uuid"], name="plainjobs_job_unique_uuid"
226 ),
227 ],
228 )
229
230 def revert_to_job_request(self) -> JobRequest:
231 """Undo convert_to_job_process — put the job back in the request queue."""
232 with transaction.atomic():
233 job_request = JobRequest.query.create(
234 uuid=self.job_request_uuid,
235 job_class=self.job_class,
236 parameters=self.parameters,
237 priority=self.priority,
238 source=self.source,
239 queue=self.queue,
240 retries=self.retries,
241 retry_attempt=self.retry_attempt,
242 concurrency_key=self.concurrency_key,
243 trace_id=self.trace_id,
244 span_id=self.span_id,
245 )
246 self.delete()
247 return job_request
248
249 def run(self) -> JobResult:
250 links = []
251 if self.trace_id and self.span_id:
252 try:
253 links.append(
254 Link(
255 SpanContext(
256 trace_id=int(self.trace_id, 16),
257 span_id=int(self.span_id, 16),
258 is_remote=True,
259 trace_flags=TraceFlags(TraceFlags.SAMPLED),
260 )
261 )
262 )
263 except (ValueError, TypeError):
264 logger.warning(
265 "Invalid trace context for job",
266 extra={"job_uuid": self.uuid},
267 )
268
269 metric_attributes: dict[str, Any] = process_metric_attributes(
270 self.queue, self.job_class
271 )
272 start_time = time.perf_counter()
273 try:
274 with tracer.start_as_current_span(
275 f"process {self.queue}",
276 kind=SpanKind.CONSUMER,
277 attributes={
278 **metric_attributes,
279 MESSAGING_OPERATION_NAME: "process",
280 MESSAGING_MESSAGE_ID: str(self.uuid),
281 MESSAGING_CONSUMER_GROUP_NAME: self.queue,
282 },
283 links=links,
284 ) as span:
285 # This is how we know it has been picked up.
286 # Keep `started_at` as a local: reading `self.started_at` back
287 # through the descriptor types as `datetime | None` (the field
288 # is `allow_null=True`), which doesn't subtract cleanly below.
289 started_at = timezone.now()
290 self.started_at = started_at
291 self.update(fields=["started_at"])
292
293 if self.requested_at:
294 queue_wait = (started_at - self.requested_at).total_seconds()
295 queue_wait_duration_histogram.record(queue_wait, metric_attributes)
296
297 try:
298 job = jobs_registry.load_job(self.job_class, self.parameters)
299 job.job_process = self
300
301 try:
302 job.run()
303 except DeferJob as e:
304 # Job deferred - not an error, log at INFO level
305 logger.info(
306 "Job deferred",
307 extra={
308 "delay": e.delay,
309 "increment_retries": e.increment_retries,
310 "job_class": self.job_class,
311 "job_process_uuid": self.uuid,
312 },
313 )
314 result = self.defer(job=job, defer_exception=e)
315 if result.retry_job_request_uuid is None:
316 # Re-enqueue was blocked by should_enqueue() —
317 # either the default uniqueness rule (a peer
318 # exists) or a user override (rate limit, custom
319 # rule). Same treatment as the initial-enqueue
320 # path's `job.enqueue.skipped`: not an error,
321 # just visibility on the consumer span.
322 span.set_attribute("plain.jobs.defer.skipped", True)
323 return result
324
325 return self.convert_to_result(status=JobResultStatuses.SUCCESSFUL)
326
327 except Exception as e:
328 # Note: if a rescuer already wrote JobResult(LOST) for this
329 # row (heartbeat went stale during a long job, then the job
330 # actually finished), the convert_to_result below trips the
331 # unique constraint on job_process_uuid and produces a
332 # second log line. Rare; correct outcome; not worth
333 # pre-checking on every successful job.
334 logger.exception("Job failed")
335 error_type = record_span_error(span, e)
336 metric_attributes[ERROR_TYPE] = error_type
337 return self.convert_to_result(
338 status=JobResultStatuses.ERRORED,
339 error="".join(traceback.format_tb(e.__traceback__)),
340 error_type=error_type,
341 )
342 finally:
343 duration = time.perf_counter() - start_time
344 operation_duration_histogram.record(duration, metric_attributes)
345
346 def defer(self, *, job: Job, defer_exception: DeferJob) -> JobResult:
347 """Defer this job by re-enqueueing it for later execution.
348
349 Atomically deletes the JobProcess, re-enqueues the job, and creates
350 a JobResult. The concurrency slot is released before re-enqueue so
351 the new request's own `should_enqueue()` check can pass.
352
353 If `should_enqueue()` blocks the re-enqueue, the framework honors
354 that signal — same convention as `run_in_worker()` and `retry_job()`,
355 which both return `None` silently in the same situation. The
356 JobResult is still `DEFERRED` but `retry_job_request_uuid` is
357 `None`, the error message records that the re-enqueue was skipped,
358 and the caller stamps `plain.jobs.defer.skipped=True` on the
359 consumer span so this case is queryable in APM without surfacing
360 as an exception.
361 """
362 # Calculate new retry_attempt based on increment_retries
363 retry_attempt = (
364 self.retry_attempt + 1
365 if defer_exception.increment_retries
366 else self.retry_attempt
367 )
368
369 with transaction.atomic():
370 # 1. Save JobProcess state and delete (releases concurrency slot)
371 job_process_uuid = self.uuid
372 job_request_uuid = self.job_request_uuid
373 requested_at = self.requested_at
374 started_at = self.started_at
375 self.delete()
376
377 # 2. Re-enqueue job (concurrency check can now pass)
378 new_job_request = job.run_in_worker(
379 queue=self.queue,
380 delay=defer_exception.delay,
381 priority=self.priority,
382 retries=self.retries,
383 retry_attempt=retry_attempt,
384 concurrency_key=self.concurrency_key,
385 )
386
387 if new_job_request is None:
388 error = (
389 f"Deferred for {defer_exception.delay} seconds "
390 f"(re-enqueue skipped: should_enqueue() returned False "
391 f"for concurrency_key '{self.concurrency_key}')"
392 )
393 retry_job_request_uuid = None
394 else:
395 error = f"Deferred for {defer_exception.delay} seconds"
396 retry_job_request_uuid = new_job_request.uuid
397
398 # 3. Create JobResult (linking to new request if one was created)
399 result = JobResult.query.create(
400 ended_at=timezone.now(),
401 error=error,
402 status=JobResultStatuses.DEFERRED,
403 retry_job_request_uuid=retry_job_request_uuid,
404 # From the JobProcess
405 job_process_uuid=job_process_uuid,
406 started_at=started_at,
407 # From the JobRequest
408 job_request_uuid=job_request_uuid,
409 requested_at=requested_at,
410 job_class=self.job_class,
411 parameters=self.parameters,
412 priority=self.priority,
413 source=self.source,
414 queue=self.queue,
415 retries=self.retries,
416 retry_attempt=self.retry_attempt,
417 concurrency_key=self.concurrency_key,
418 trace_id=self.trace_id,
419 span_id=self.span_id,
420 )
421
422 # Counter ticks for the DEFERRED outcome too — defer() bypasses
423 # convert_to_result, so without this the deferred path would not
424 # show up in the consumed counter.
425 record_consumed(result)
426 return result
427
428 def convert_to_result(
429 self,
430 *,
431 status: str,
432 error: str = "",
433 error_type: str | None = None,
434 fire_hook: bool = True,
435 ) -> JobResult:
436 """
437 Convert this JobProcess to a JobResult.
438
439 error_type, when supplied, is the OTel-style exception name (matching
440 the spec's `error.type` attribute). It rides along to the consumed
441 counter so dashboards can group ERRORED jobs by exception class. Only
442 the live exception-driven paths supply it — rescue (LOST) and direct
443 cancellations have no exception object to derive it from.
444
445 fire_hook controls whether on_aborted dispatches synchronously. The
446 rescue path passes fire_hook=False so it can dispatch hooks AFTER its
447 outer transaction commits — otherwise a hook DB error would mark the
448 connection for rollback even though dispatch_aborted_hook catches the
449 exception, poisoning the rescue commit.
450 """
451 with transaction.atomic():
452 result = JobResult.query.create(
453 ended_at=timezone.now(),
454 error=error,
455 status=status,
456 # From the JobProcess
457 job_process_uuid=self.uuid,
458 started_at=self.started_at,
459 # From the JobRequest
460 job_request_uuid=self.job_request_uuid,
461 requested_at=self.requested_at,
462 job_class=self.job_class,
463 parameters=self.parameters,
464 priority=self.priority,
465 source=self.source,
466 queue=self.queue,
467 retries=self.retries,
468 retry_attempt=self.retry_attempt,
469 concurrency_key=self.concurrency_key,
470 trace_id=self.trace_id,
471 span_id=self.span_id,
472 )
473
474 # Delete the JobProcess now
475 self.delete()
476
477 # Counter ticks for every terminal status — the live SUCCESSFUL/ERRORED
478 # paths plus the LOST/CANCELLED paths that don't flow through
479 # JobProcess.run()'s finally. The outcome attribute lets dashboards
480 # split throughput by final status; error_type is forwarded for ERRORED
481 # jobs caught by the live path.
482 record_consumed(result, error_type=error_type)
483
484 # Fire Job.on_aborted outside the atomic block so a raise in user code
485 # can't roll back the framework's bookkeeping. Only for terminal
486 # statuses run() couldn't observe.
487 if fire_hook and status in (
488 JobResultStatuses.LOST,
489 JobResultStatuses.CANCELLED,
490 ):
491 result.dispatch_aborted_hook()
492
493 return result
494
495 def as_json(self) -> dict[str, str | int | dict | None]:
496 """A JSON-compatible representation to make it easier to reference in Sentry or logging"""
497 return {
498 "uuid": str(self.uuid),
499 "created_at": self.created_at.isoformat(),
500 "started_at": self.started_at.isoformat() if self.started_at else None,
501 "job_request_uuid": str(self.job_request_uuid),
502 "job_class": self.job_class,
503 "parameters": self.parameters,
504 "priority": self.priority,
505 "source": self.source,
506 "queue": self.queue,
507 "retries": self.retries,
508 "retry_attempt": self.retry_attempt,
509 "concurrency_key": self.concurrency_key,
510 "trace_id": self.trace_id,
511 "span_id": self.span_id,
512 }
513
514
515class JobResultQuerySet(postgres.QuerySet["JobResult"]):
516 def successful(self) -> Self:
517 return self.filter(status=JobResultStatuses.SUCCESSFUL)
518
519 def cancelled(self) -> Self:
520 return self.filter(status=JobResultStatuses.CANCELLED)
521
522 def lost(self) -> Self:
523 return self.filter(status=JobResultStatuses.LOST)
524
525 def errored(self) -> Self:
526 return self.filter(status=JobResultStatuses.ERRORED)
527
528 def retried(self) -> Self:
529 return self.filter(
530 postgres.Q(retry_job_request_uuid__isnull=False)
531 | postgres.Q(retry_attempt__gt=0)
532 )
533
534 def failed(self) -> Self:
535 return self.filter(
536 status__in=[
537 JobResultStatuses.ERRORED,
538 JobResultStatuses.LOST,
539 JobResultStatuses.CANCELLED,
540 ]
541 )
542
543 def retryable(self) -> Self:
544 return self.failed().filter(
545 retry_job_request_uuid__isnull=True,
546 retries__gt=0,
547 retry_attempt__lt=F("retries"),
548 )
549
550 def retry_failed_jobs(self) -> None:
551 for result in self.retryable():
552 try:
553 result.retry_job()
554 except Exception:
555 # If something went wrong (like a job class being deleted)
556 # then we immediately increment the retry_attempt on the existing obj
557 # so it won't retry forever.
558 logger.exception(
559 "Failed to retry job, incrementing retry_attempt",
560 extra={"result": str(result)},
561 )
562 result.retry_attempt += 1
563 result.update(fields=["retry_attempt"])
564
565
566class JobResultStatuses(postgres.TextChoices):
567 SUCCESSFUL = "SUCCESSFUL", "Successful"
568 ERRORED = "ERRORED", "Errored" # Threw an error
569 CANCELLED = "CANCELLED", "Cancelled" # Interrupted by shutdown/deploy
570 DEFERRED = "DEFERRED", "Deferred" # Intentionally rescheduled (will run again)
571 LOST = (
572 "LOST",
573 "Lost",
574 ) # Either process lost, lost in transit, or otherwise never finished
575
576
577@postgres.register_model
578class JobResult(postgres.Model):
579 """
580 All in-process and completed jobs are stored in this table.
581 """
582
583 uuid = types.UUIDField(generate=True)
584 created_at = types.DateTimeField(create_now=True)
585
586 # From the Job
587 job_process_uuid = types.UUIDField()
588 started_at = types.DateTimeField(required=False, allow_null=True)
589 ended_at = types.DateTimeField(required=False, allow_null=True)
590 error = types.TextField(required=False, default="")
591 status = types.TextField(
592 max_length=20,
593 choices=JobResultStatuses.choices,
594 )
595
596 # From the JobRequest
597 job_request_uuid = types.UUIDField()
598 requested_at = types.DateTimeField(required=False, allow_null=True)
599 job_class = types.TextField(max_length=255)
600 parameters: dict[str, Any] | None = types.JSONField(required=False, allow_null=True)
601 priority = types.SmallIntegerField(default=0)
602 source = types.TextField(required=False, default="")
603 queue = types.TextField(default="default", max_length=255)
604 retries = types.SmallIntegerField(default=0)
605 retry_attempt = types.SmallIntegerField(default=0)
606 concurrency_key = types.TextField(max_length=255, required=False, default="")
607
608 # Retries
609 retry_job_request_uuid = types.UUIDField(required=False, allow_null=True)
610
611 # OpenTelemetry trace context
612 trace_id = types.TextField(max_length=34, required=False, allow_null=True)
613 span_id = types.TextField(max_length=18, required=False, allow_null=True)
614
615 query: JobResultQuerySet = JobResultQuerySet()
616
617 model_options = postgres.Options(
618 ordering=["-created_at"],
619 indexes=[
620 postgres.Index(
621 name="plainjobs_jobresult_created_at_idx", fields=["created_at"]
622 ),
623 postgres.Index(name="plainjobs_jobresult_status_idx", fields=["status"]),
624 ],
625 constraints=[
626 postgres.UniqueConstraint(
627 fields=["uuid"], name="plainjobs_jobresult_unique_uuid"
628 ),
629 # One JobProcess produces exactly one JobResult. Guards the
630 # rescue-vs-late-finish race: if our heartbeat goes stale during a
631 # DB outage, a peer rescuer creates JobResult(LOST) for our
632 # JobProcess. When our subprocess eventually finishes and calls
633 # convert_to_result on the now-deleted JobProcess, the second
634 # insert hits this constraint and is swallowed by process_job's
635 # outer except — instead of silently producing two divergent
636 # results for the same run.
637 postgres.UniqueConstraint(
638 fields=["job_process_uuid"],
639 name="plainjobs_jobresult_unique_job_process_uuid",
640 ),
641 ],
642 )
643
644 def dispatch_aborted_hook(self) -> None:
645 """
646 Load the Job class and call its on_aborted hook with this result.
647
648 Errors loading the class or running the hook are logged but suppressed
649 so JobProcess → JobResult bookkeeping is never blocked by user code or
650 stale registrations.
651 """
652 try:
653 job = jobs_registry.load_job(self.job_class, self.parameters)
654 except Exception:
655 logger.exception(
656 "Failed to load job for on_aborted hook",
657 extra={"job_class": self.job_class},
658 )
659 return
660
661 try:
662 job.on_aborted(self)
663 except Exception:
664 logger.exception(
665 "Job.on_aborted raised",
666 extra={"job_class": self.job_class},
667 )
668
669 def retry_job(self, delay: int | None = None) -> JobRequest | None:
670 retry_attempt = self.retry_attempt + 1
671 job = jobs_registry.load_job(self.job_class, self.parameters)
672
673 if delay is None:
674 retry_delay = job.calculate_retry_delay(retry_attempt)
675 else:
676 retry_delay = delay
677
678 with transaction.atomic():
679 result = job.run_in_worker(
680 # Pass most of what we know through so it stays consistent
681 queue=self.queue,
682 delay=retry_delay,
683 priority=self.priority,
684 retries=self.retries,
685 retry_attempt=retry_attempt,
686 concurrency_key=self.concurrency_key,
687 )
688 if result:
689 self.retry_job_request_uuid = result.uuid
690 self.update(fields=["retry_job_request_uuid"])
691 return result
692
693 return None
694
695
696@postgres.register_model
697class WorkerHeartbeat(postgres.Model):
698 """
699 A live registration row written by each worker process while it's running.
700
701 Workers create a row at startup, bump `last_heartbeat_at` periodically, and
702 delete it on clean shutdown. `rescue_stale_workers` finds rows whose
703 heartbeat is older than `JOBS_HEARTBEAT_TIMEOUT` and rescues their
704 in-flight jobs.
705 """
706
707 worker_id = types.UUIDField()
708 hostname = types.TextField(max_length=255)
709 pid = types.IntegerField()
710 queues: list[str] = types.JSONField()
711 started_at = types.DateTimeField(create_now=True)
712 last_heartbeat_at = types.DateTimeField()
713
714 model_options = postgres.Options(
715 ordering=["-last_heartbeat_at"],
716 constraints=[
717 # The unique constraint provides the worker_id lookup index.
718 postgres.UniqueConstraint(
719 fields=["worker_id"],
720 name="plainjobs_workerheartbeat_unique_worker_id",
721 ),
722 ],
723 )
724
725 def __str__(self) -> str:
726 return f"WorkerHeartbeat({self.worker_id} on {self.hostname}:{self.pid})"
727
728
729def heartbeat_cutoff() -> datetime.datetime:
730 """The timestamp before which a WorkerHeartbeat is considered stale.
731
732 Single source of truth — rescue, admin display, and OTel gauges all
733 consult this so they agree on which workers are alive.
734 """
735 return timezone.now() - datetime.timedelta(seconds=settings.JOBS_HEARTBEAT_TIMEOUT)
736
737
738def rescue_stale_workers() -> list[JobResult]:
739 """
740 Convert in-flight JobProcess rows from dead workers to JobResult(LOST).
741
742 A worker is dead when its WorkerHeartbeat is older than
743 JOBS_HEARTBEAT_TIMEOUT. Detection is heartbeat-based, not time-based, so
744 a long-running legitimate job is safe as long as its worker keeps
745 heartbeating.
746
747 Returns the JobResults whose on_aborted hook still needs to fire. The
748 caller dispatches them, interleaving heartbeat ticks if needed — a slow
749 or large batch of hooks would otherwise starve the calling worker's
750 heartbeat and trigger false-positive rescue from a peer.
751
752 This is a free function (not a queryset method) because rescue is
753 inherently global: filtering would let one rescuer claim a dead heartbeat
754 without converting all of that worker's jobs, stranding the rest forever.
755 """
756 cutoff = heartbeat_cutoff()
757 dead_workers = WorkerHeartbeat.query.filter(last_heartbeat_at__lt=cutoff)
758
759 pending_hooks: list[JobResult] = []
760 for worker in dead_workers:
761 # Per-worker rescue is atomic: the heartbeat delete (claim) and every
762 # JobProcess→JobResult conversion either all commit, or all roll
763 # back. Without this, a mid-loop failure would leave the heartbeat
764 # deleted but some JobProcesses still stamped with the dead worker_id
765 # — stranded forever with no heartbeat to match.
766 #
767 # on_aborted hooks are deferred: dispatching them inside the atomic
768 # block would let a hook's DB error mark the connection for rollback
769 # (even though dispatch_aborted_hook swallows the exception),
770 # aborting the rescue commit.
771 worker_hooks: list[JobResult] = []
772 try:
773 with transaction.atomic():
774 # Atomic claim. If another rescuer also saw this dead
775 # heartbeat, only one of us deletes a row. The loser sees 0
776 # affected and skips.
777 claimed = WorkerHeartbeat.query.filter(
778 worker_id=worker.worker_id,
779 last_heartbeat_at__lt=cutoff,
780 ).delete()
781 if not claimed:
782 continue
783
784 # list() materializes the queryset before the loop body
785 # starts deleting rows, so iteration can't skip entries.
786 for job in list(JobProcess.query.filter(worker_id=worker.worker_id)):
787 result = job.convert_to_result(
788 status=JobResultStatuses.LOST, fire_hook=False
789 )
790 worker_hooks.append(result)
791 except Exception:
792 # One dead worker's failure shouldn't abort rescue of others. The
793 # next rescue tick will retry this worker (heartbeat was rolled
794 # back, so it's still discoverable).
795 logger.exception(
796 "Failed to rescue jobs for dead worker",
797 extra={"worker_id": str(worker.worker_id)},
798 )
799 continue
800
801 # Rescue committed. Hand hooks back for the caller to dispatch.
802 pending_hooks.extend(worker_hooks)
803
804 return pending_hooks