Merge pull request #2449 from aboutcode-org/2437-refresh-run

Refresh stale and missing pipeline jobs

Author: keshav-space

vulnerabilities/management/commands/run_scheduler.py

@@ -12,6 +12,7 @@
 
 from vulnerabilities import models
 from vulnerabilities.schedules import clear_zombie_pipeline_schedules
+from vulnerabilities.schedules import refresh_runs
 from vulnerabilities.schedules import scheduled_job_exists
 from vulnerabilities.schedules import update_pipeline_schedule
 
@@ -38,5 +39,6 @@ class Command(rqscheduler.Command):
     def handle(self, *args, **kwargs):
         clear_zombie_pipeline_schedules()
         update_pipeline_schedule()
+        refresh_runs()
         init_pipeline_scheduled()
         super(Command, self).handle(*args, **kwargs)

vulnerabilities/pipelines/v2_importers/redhat_importer.py

@@ -30,8 +30,8 @@
 from vulnerabilities.pipelines import VulnerableCodeBaseImporterPipelineV2
 from vulnerabilities.pipes import extractcode_utils
 from vulnerabilities.severity_systems import REDHAT_AGGREGATE
+from vulnerabilities.utils import get_item
 from vulnerabilities.utils import load_json
-from vulntotal import vulntotal_utils
 
 
 class RedHatImporterPipeline(VulnerableCodeBaseImporterPipelineV2):
@@ -135,7 +135,7 @@ def parse_advisory(self, record):
                 )
             )
 
-        impacts = get_item(advisory, "product_tree", "branches", 0, "branches", default=[])
+        impacts = get_item(advisory, "product_tree", "branches", 0, "branches") or []
         for impact in impacts:
             if impact["category"] == "product_family":
                 continue
@@ -145,7 +145,6 @@ def parse_advisory(self, record):
                     "product",
                     "product_identification_helper",
                     "purl",
-                    default=None,
                 ):
                     if not purl.startswith("pkg:rpm/"):
                         continue
@@ -188,11 +187,3 @@ def clean_download(self):
 
     def on_failure(self):
         self.clean_download()
-
-
-def get_item(entity, *attributes, default=None):
-    try:
-        result = vulntotal_utils.get_item(entity, *attributes)
-    except (KeyError, IndexError, TypeError) as e:
-        result = default
-    return result

vulnerabilities/pipelines/v2_improvers/relate_severities.py

@@ -103,7 +103,6 @@ def relate_severities(self):
             matches = advisory_id_map.get(advisory.advisory_id, set())
             for target_id in matches:
                 if target_id != advisory.id:
-                    self.log(f"Relating advisory {advisory.avid} to {target_id}")
                     relations.append(
                         through(
                             from_advisoryv2_id=target_id,

vulnerabilities/pipes/group_advisories.py

@@ -23,8 +23,6 @@
 
 @transaction.atomic
 def delete_and_save_advisory_set(groups, package, relation=None):
-    print(f"Grouping advisories for package: {package.purl}")
-
     AdvisorySet.objects.filter(
         package=package,
         relation_type=relation,
@@ -101,8 +99,6 @@ def delete_and_save_advisory_set(groups, package, relation=None):
             batch_size=10000,
         )
 
-    print(f"Successfully saved advisory sets for package: {package.purl}")
-
 
 def group_advisory_for_package(package, logger=None):
     """
@@ -133,7 +129,6 @@ def group_advisory_for_package(package, logger=None):
         group_single_package_with_provided_advisories(
             package, affecting_advisories, fixed_by_advisories
         )
-        logger(f"Successfully rebuilt advisory sets for package {package.purl}")
     except Exception as e:
         if logger:
             logger(f"Failed rebuilding advisory sets for package {package.purl}: {e!r}")

vulnerabilities/schedules.py

@@ -116,3 +116,46 @@ def update_pipeline_schedule():
                 pipeline.run_priority = run_priority
                 pipeline.run_interval = run_interval
                 pipeline.save()
+
+
+def mark_stale_runs():
+    """Mark unfinished pipeline runs as stale."""
+    from django.db.models import Q
+
+    from vulnerabilities.models import PipelineRun
+
+    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):
+            pipeline_latest_run = job.pipeline.latest_run
+            if pipeline_latest_run and pipeline_latest_run.run_id == job.run_id:
+                enqueue_run(run=job)
+                missing_jobs_count += 1
+            else:
+                job.set_run_staled()
+
+    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.

vulnerabilities/tests/pipelines/v2_importers/test_redhat_importer_v2.py

@@ -17,6 +17,7 @@
 from vulnerabilities.models import PackageV2
 from vulnerabilities.pipelines.v2_importers.redhat_importer import RedHatImporterPipeline
 from vulnerabilities.tests import util_tests
+from vulnerabilities.tests.pipelines import TestLogger
 from vulnerabilities.utils import normalize_list
 
 TEST_DATA = Path(__file__).parent.parent.parent / "test_data" / "redhat" / "csaf_2_0"
@@ -25,10 +26,15 @@
 class TestArchLinuxImporterPipeline(TestCase):
     @patch("vulnerabilities.pipelines.v2_importers.redhat_importer.RedHatImporterPipeline.fetch")
     def test_redhat_advisories_v2(self, mock_fetch):
+        logger = TestLogger()
+
         mock_fetch.__name__ = "fetch"
         pipeline = RedHatImporterPipeline()
         pipeline.location = TEST_DATA
-        pipeline.execute()
+        pipeline.log = logger.write
+        exit_code, _ = pipeline.execute()
+
+        self.assertEqual(exit_code, 0)
         self.assertEqual(6, AdvisoryV2.objects.count())
         self.assertEqual(93, PackageV2.objects.count())
         expected_file = TEST_DATA.parent / "redhat_advisoryv2-expected.json"

vulnerabilities/tests/test_models.py

@@ -688,6 +688,23 @@ def test_pipelinerun_success(self):
         self.run1.set_run_ended(exitcode=0)
         self.assertEqual(self.run1.status, self.run1.Status.SUCCESS)
 
+    def test_pipelinerun_schedules_mark_stale_runs_unfinished_job(self):
+        from vulnerabilities.schedules import mark_stale_runs
+
+        self.run1.set_run_started()
+        self.assertEqual(self.run1.run_exitcode, None)
+        mark_stale_runs()
+        self.run1.refresh_from_db()
+        self.assertEqual(self.run1.run_exitcode, 88)
+
+    def test_pipelinerun_schedules_mark_stale_runs_job_not_yet_started(self):
+        from vulnerabilities.schedules import mark_stale_runs
+
+        self.assertEqual(self.run1.run_exitcode, None)
+        mark_stale_runs()
+        self.run1.refresh_from_db()
+        self.assertEqual(self.run1.run_exitcode, None)
+
 
 class TestPipelineScheduleModel(DjangoTestCase):
     def setUp(self):