Use the future callback to convert jobs to cancelled

Author: davegaeddert

bolt-jobs/bolt/jobs/migrations/0016_remove_job_worker_uuid_remove_jobresult_worker_uuid.py

@@ -0,0 +1,20 @@
+# Generated by Bolt 5.0.dev20240121220125 on 2024-01-22 16:53
+
+from bolt.db import migrations
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("boltqueue", "0015_job_worker_uuid_jobresult_worker_uuid_and_more"),
+    ]
+
+    operations = [
+        migrations.RemoveField(
+            model_name="job",
+            name="worker_uuid",
+        ),
+        migrations.RemoveField(
+            model_name="jobresult",
+            name="worker_uuid",
+        ),
+    ]

bolt-jobs/bolt/jobs/models.py

@@ -61,7 +61,7 @@ class Meta:
     def __str__(self):
         return f"{self.job_class} [{self.uuid}]"
 
-    def convert_to_job(self, *, worker_uuid=None):
+    def convert_to_job(self):
         """
         JobRequests are the pending jobs that are waiting to be executed.
         We immediately convert them to JobResults when they are picked up.
@@ -76,7 +76,6 @@ def convert_to_job(self, *, worker_uuid=None):
                 retries=self.retries,
                 retry_attempt=self.retry_attempt,
                 unique_key=self.unique_key,
-                worker_uuid=worker_uuid,
             )
 
             # Delete the pending JobRequest now
@@ -118,9 +117,6 @@ class Job(models.Model):
     created_at = models.DateTimeField(auto_now_add=True, db_index=True)
     started_at = models.DateTimeField(blank=True, null=True, db_index=True)
 
-    # To associate with a worker
-    worker_uuid = models.UUIDField(blank=True, null=True, db_index=True)
-
     # From the JobRequest
     job_request_uuid = models.UUIDField(db_index=True)
     job_class = models.CharField(max_length=255, db_index=True)
@@ -168,8 +164,6 @@ def convert_to_result(self, *, status, error=""):
                 ended_at=timezone.now(),
                 error=error,
                 status=status,
-                # From the worker
-                worker_uuid=self.worker_uuid,
                 # From the Job
                 job_uuid=self.uuid,
                 started_at=self.started_at,
@@ -248,9 +242,6 @@ class JobResult(models.Model):
     uuid = models.UUIDField(default=uuid.uuid4, editable=False, unique=True)
     created_at = models.DateTimeField(auto_now_add=True, db_index=True)
 
-    # To associate with a worker
-    worker_uuid = models.UUIDField(blank=True, null=True, db_index=True)
-
     # From the Job
     job_uuid = models.UUIDField(db_index=True)
     started_at = models.DateTimeField(blank=True, null=True, db_index=True)

bolt-jobs/bolt/jobs/workers.py

@@ -2,8 +2,8 @@
 import multiprocessing
 import os
 import time
-import uuid
-from concurrent.futures import ProcessPoolExecutor
+from concurrent.futures import Future, ProcessPoolExecutor
+from functools import partial
 
 from bolt.db import transaction
 from bolt.logs import app_logger
@@ -28,8 +28,6 @@ def __init__(self, max_processes=None, max_jobs_per_process=None, stats_every=No
         self.max_processes = self.executor._max_workers
         self.max_jobs_per_process = max_jobs_per_process
 
-        self.uuid = uuid.uuid4()
-
         self._is_shutting_down = False
 
     def run(self):
@@ -63,37 +61,25 @@ def run(self):
                     job_request.source,
                 )
 
-                job = job_request.convert_to_job(worker_uuid=self.uuid)
+                job = job_request.convert_to_job()
 
             job_uuid = str(job.uuid)  # Make a str copy
 
             # Release these now
             del job_request
             del job
 
-            self.executor.submit(process_job, job_uuid)
+            future = self.executor.submit(process_job, job_uuid)
+            future.add_done_callback(partial(future_finished_callback, job_uuid))
 
     def shutdown(self):
         if self._is_shutting_down:
             # Already shutting down somewhere else
             return
 
-        self._is_shutting_down = True
-
         logger.info("Job worker shutdown started")
-
-        # Make an attept to immediatelly move any unstarted jobs from this worker
-        # to the JobResults as cancelled. If they have a retry, they'll be picked
-        # up by the next worker process.
-        for job in Job.objects.filter(worker_uuid=self.uuid, started_at__isnull=True):
-            job.convert_to_result(status=JobResultStatuses.CANCELLED)
-
-        # Now shutdown the process pool.
-        # There's still some chance of a race condition here where we
-        # just deleted a Job and it was still picked up, but we'll log
-        # missing jobs as warnings instead of exceptions.
+        self._is_shutting_down = True
         self.executor.shutdown(wait=True, cancel_futures=True)
-
         logger.info("Job worker shutdown complete")
 
     def maybe_log_stats(self):
@@ -149,6 +135,15 @@ def check_job_results(self):
         JobResult.objects.retry_failed_jobs()
 
 
+def future_finished_callback(job_uuid: str, future: Future):
+    if future.cancelled():
+        logger.warning("Job cancelled job_uuid=%s", job_uuid)
+        job = Job.objects.get(uuid=job_uuid)
+        job.convert_to_result(status=JobResultStatuses.CANCELLED)
+    else:
+        logger.debug("Job finished job_uuid=%s", job_uuid)
+
+
 def process_job(job_uuid):
     try:
         worker_pid = os.getpid()