feat: requeue pipeline runs missing from queue

Signed-off-by: Keshav Priyadarshi <git@keshav.space>

Author: keshav-space

vulnerabilities/schedules.py

@@ -124,11 +124,34 @@ def mark_stale_runs():
 
     from vulnerabilities.models import PipelineRun
 
-    for job in PipelineRun.objects.filter(run_start_date__isnull=False).filter(
+    stale_jobs = PipelineRun.objects.filter(run_start_date__isnull=False).filter(
         Q(run_end_date__isnull=True) | Q(run_exitcode__isnull=True)
-    ):
+    )
+    stale_jobs_count = stale_jobs.count()
+
+    for job in stale_jobs.iterator(chunk_size=1000):
         job.set_run_staled()
 
+    log.info(f"Marked {stale_jobs_count} unfinished jobs as stale.")
+
+
+def requeue_missing_jobs():
+    """Requeue pipeline runs that have not started and are missing from the queue."""
+
+    from vulnerabilities.models import PipelineRun
+    from vulnerabilities.tasks import enqueue_run
+    from vulnerabilities.tasks import is_job_in_queue
+
+    missing_jobs_count = 0
+    for job in PipelineRun.objects.filter(run_start_date__isnull=True).iterator(chunk_size=1000):
+        if not is_job_in_queue(job_id=job.run_id):
+            enqueue_run(run=job)
+            missing_jobs_count += 1
+
+    log.info(f"Requeued {missing_jobs_count} missing jobs.")
+
 
 def refresh_runs():
+    """Refresh the stale and missing jobs."""
     mark_stale_runs()
+    requeue_missing_jobs()

vulnerabilities/tasks.py

@@ -115,9 +115,22 @@ def set_run_failure(job, connection, type, value, traceback):
     run.set_run_ended(exitcode=1, output=f"value={value} trace={traceback}")
 
 
+def enqueue_run(run):
+    pipeline_schedule = run.pipeline
+    queue = queues.get(pipeline_schedule.get_run_priority_display())
+
+    job = queue.enqueue(
+        execute_pipeline,
+        pipeline_schedule.pipeline_id,
+        run.run_id,
+        job_id=str(run.run_id),
+        on_failure=set_run_failure,
+        job_timeout=f"{pipeline_schedule.execution_timeout}h",
+    )
+
+
 def enqueue_pipeline(pipeline_id):
     pipeline_schedule = models.PipelineSchedule.objects.get(pipeline_id=pipeline_id)
-    queue = queues.get(pipeline_schedule.get_run_priority_display())
 
     if pipeline_schedule.status in [
         models.PipelineRun.Status.RUNNING,
@@ -134,14 +147,8 @@ def enqueue_pipeline(pipeline_id):
     run = models.PipelineRun.objects.create(
         pipeline=pipeline_schedule,
     )
-    job = queue.enqueue(
-        execute_pipeline,
-        pipeline_id,
-        run.run_id,
-        job_id=str(run.run_id),
-        on_failure=set_run_failure,
-        job_timeout=f"{pipeline_schedule.execution_timeout}h",
-    )
+
+    enqueue_run(run)
 
 
 def dequeue_job(job_id):
@@ -152,6 +159,12 @@ def dequeue_job(job_id):
             queue.remove(job_id)
 
 
+def is_job_in_queue(job_id):
+    """Return whether a job exists in any queue."""
+
+    return any(job_id in queue.jobs for queue in queues.values())
+
+
 def compute_queue_load_factor():
     """
     Compute worker load per queue.