From ff2652ed540f28fd3924b34cf192672c3d9c4880 Mon Sep 17 00:00:00 2001
From: Oliver Chang <oliverchang@users.noreply.github.com>
Date: Wed, 14 Aug 2024 00:30:36 +1000
Subject: [PATCH] Fix update_mappings_for_fuzzer when there are stale job
 references. (#4160)

---
 .../_internal/fuzzing/fuzzer_selection.py     | 21 ++++++++++++++++++-
 .../core/fuzzing/fuzzer_selection_test.py     | 16 ++++++++++++++
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/src/clusterfuzz/_internal/fuzzing/fuzzer_selection.py b/src/clusterfuzz/_internal/fuzzing/fuzzer_selection.py
index d747a3b4fe..34d928ced2 100644
--- a/src/clusterfuzz/_internal/fuzzing/fuzzer_selection.py
+++ b/src/clusterfuzz/_internal/fuzzing/fuzzer_selection.py
@@ -31,7 +31,8 @@
 def update_mappings_for_fuzzer(fuzzer, mappings=None):
   """Clear existing mappings for a fuzzer, and replace them."""
   if mappings is None:
-    mappings = fuzzer.jobs
+    # Make a copy in case we need to modify it.
+    mappings = fuzzer.jobs.copy()
 
   query = data_types.FuzzerJob.query()
   query = query.filter(data_types.FuzzerJob.fuzzer == fuzzer.name)
@@ -47,7 +48,22 @@ def update_mappings_for_fuzzer(fuzzer, mappings=None):
     jobs = {job.name: job for job in jobs}
   else:
     jobs = {}
+
+  fuzzer_modified = False
+
   for job_name in mappings:
+    if job_name not in jobs:
+      # Job references a deleted job, clean it up.
+      try:
+        fuzzer.jobs.remove(job_name)
+        fuzzer_modified = True
+      except ValueError:
+        # If `mappings` was provided via an argument, it's possible it won't
+        # exist in `fuzzer.jobs`.
+        pass
+
+      continue
+
     mapping = old_mappings.pop(job_name, None)
     if not mapping:
       mapping = data_types.FuzzerJob()
@@ -59,6 +75,9 @@ def update_mappings_for_fuzzer(fuzzer, mappings=None):
   ndb_utils.put_multi(new_mappings)
   ndb_utils.delete_multi([m.key for m in list(old_mappings.values())])
 
+  if fuzzer_modified:
+    fuzzer.put()
+
 
 def update_mappings_for_job(job, mappings):
   """Clear existing mappings for a job, and replace them."""
diff --git a/src/clusterfuzz/_internal/tests/core/fuzzing/fuzzer_selection_test.py b/src/clusterfuzz/_internal/tests/core/fuzzing/fuzzer_selection_test.py
index 9c2e450c91..eaa8bf792e 100644
--- a/src/clusterfuzz/_internal/tests/core/fuzzing/fuzzer_selection_test.py
+++ b/src/clusterfuzz/_internal/tests/core/fuzzing/fuzzer_selection_test.py
@@ -16,6 +16,7 @@
 import os
 import unittest
 
+from google.cloud import ndb
 import parameterized
 
 from clusterfuzz._internal.datastore import data_types
@@ -131,6 +132,21 @@ def test_mapping_removed(self):
     mappings = _get_job_list_for_fuzzer(fuzzer)
     self.assertEqual([], mappings)
 
+  def test_nonexistent_job(self):
+    """Ensure stale, non-existent job references are handled."""
+    fuzzer = data_types.Fuzzer()
+    fuzzer.name = 'adding_jobs'
+    fuzzer.jobs = ['does_not_exist', 'job_1']
+    fuzzer.put()
+
+    fuzzer_selection.update_mappings_for_fuzzer(fuzzer)
+
+    fuzzer = ndb.Key(data_types.Fuzzer, fuzzer.key.id()).get()
+    self.assertCountEqual(['job_1'], fuzzer.jobs)
+
+    mappings = _get_job_list_for_fuzzer(fuzzer)
+    self.assertCountEqual(['job_1'], mappings)
+
 
 @test_utils.with_cloud_emulators('datastore')
 class UpdateMappingsForJobTest(unittest.TestCase):