prepare unscheduled_tasks

Signed-off-by: Javan Lacerda <javanlacerda@google.com>

Author: javanlacerda

src/clusterfuzz/_internal/remote_task/remote_task_gate.py

@@ -21,8 +21,10 @@
 
 import collections
 import random
-from typing import List
+import typing
 
+from clusterfuzz._internal.datastore import data_types
+from clusterfuzz._internal.datastore import ndb_utils
 from clusterfuzz._internal.metrics import logs
 from clusterfuzz._internal.remote_task import remote_task_adapters
 from clusterfuzz._internal.remote_task import remote_task_types
@@ -105,13 +107,39 @@ def get_job_frequency(self):
     logs.info('Job frequencies', frequencies=frequencies)
     return frequencies
 
+  def prepare_unscheduled_tasks(
+      self, unscheduled_remote_tasks: typing.List[remote_task_types.RemoteTask]
+  ) -> typing.List[remote_task_types.RemoteTask]:
+    """Prepares the unscheduled remote tasks to be sent back to the 
+      preprocess queue.
+
+      The messages in the preprocess queue expects tasks with 
+      the properties followng task_name, fuzzer, job, eta. 
+      The utasks has the signed url as argument instead of the fuzzer,
+      then for sending it back to the preprocess, it recovers its fuzzer
+      and override the argument.
+    """
+
+    job_names = {task.job_type for task in unscheduled_remote_tasks}
+    query = data_types.FuzzerJob.query(
+        data_types.FuzzerJob.job.IN(list(job_names)))
+    fuzzer_jobs = ndb_utils.get_all_from_query(query)
+    fuzzer_jobs_name_mapped = {job.name: job for job in fuzzer_jobs}
+    for task in unscheduled_remote_tasks:
+      if task.job_type not in fuzzer_jobs_name_mapped:
+        logs.error(f'{task.job_type} not found.')
+        continue
+      task.argument = fuzzer_jobs_name_mapped[task.job_type].argument
+
+    return unscheduled_remote_tasks
+
   def create_utask_main_job(self, module, job_type, input_download_url):
     adapter_id = self._get_adapter()
     service = self._service_map[adapter_id]
     return service.create_utask_main_job(module, job_type, input_download_url)
 
-  def create_utask_main_jobs(self,
-                             remote_tasks: List[remote_task_types.RemoteTask]):
+  def create_utask_main_jobs(
+      self, remote_tasks: typing.List[remote_task_types.RemoteTask]):
     """Creates a batch of remote tasks, distributing them across backends.
 
     This method handles two cases:
@@ -146,13 +174,16 @@ def create_utask_main_jobs(self,
         adapter_id = list(frequencies.keys())[i % len(frequencies)]
         tasks_by_adapter[adapter_id].append(task)
 
-    results = []
+    unscheduled_tasks = []
     for adapter_id, tasks in tasks_by_adapter.items():
       if tasks:
         try:
           logs.info(f'Sending {len(tasks)} tasks to {adapter_id}.')
           service = self._service_map[adapter_id]
-          results.extend(service.create_utask_main_jobs(tasks))
+          unscheduled_tasks.extend(service.create_utask_main_jobs(tasks))
         except Exception:  # pylint: disable=broad-except
           logs.error(f'Failed to send {len(tasks)} tasks to {adapter_id}.')
-    return results
+
+    prepared_unscheduled_tasks = self.prepare_unscheduled_tasks(
+        unscheduled_tasks)
+    return prepared_unscheduled_tasks

src/clusterfuzz/_internal/tests/core/remote_task/remote_task_gate_test.py

@@ -55,6 +55,12 @@ def setUp(self):
     self.patcher.start()
     self.addCleanup(self.patcher.stop)
 
+    self.mock_prepare_unscheduled_tasks = mock.patch.object(
+        remote_task_gate.RemoteTaskGate,
+        'prepare_unscheduled_tasks',
+        side_effect=lambda x: x).start()
+    self.addCleanup(mock.patch.stopall)
+
   def test_init(self):
     """Tests that the RemoteTaskGate initializes correctly and creates
     service map."""
@@ -197,3 +203,33 @@ def test_create_utask_main_jobs_full_gcp_batch(self, mock_get_job_frequency):
     self.mock_gcp_batch_service.create_utask_main_jobs.assert_called_once_with(
         tasks)
     self.mock_k8s_service.create_utask_main_jobs.assert_not_called()
+
+  @mock.patch.object(remote_task_gate.RemoteTaskGate, 'get_job_frequency')
+  def test_create_utask_main_jobs_calls_prepare_unscheduled_tasks(
+      self, mock_get_job_frequency):
+    """Tests that create_utask_main_jobs calls prepare_unscheduled_tasks with
+    unscheduled tasks."""
+    tasks = [
+        remote_task_types.RemoteTask('c', 'j', 'u1'),
+    ]
+    unscheduled_tasks = [
+        remote_task_types.RemoteTask('c', 'j', 'u1'),
+    ]
+    prepared_tasks = [
+        remote_task_types.RemoteTask('c', 'j', 'u1'),
+    ]
+    prepared_tasks[0].argument = 'arg'
+
+    mock_get_job_frequency.return_value = {'kubernetes': 1.0, 'gcp_batch': 0.0}
+    self.mock_k8s_service.create_utask_main_jobs.return_value = unscheduled_tasks
+
+    self.mock_prepare_unscheduled_tasks.side_effect = None
+    self.mock_prepare_unscheduled_tasks.return_value = prepared_tasks
+
+    gate = remote_task_gate.RemoteTaskGate()
+    result = gate.create_utask_main_jobs(tasks)
+
+    self.mock_k8s_service.create_utask_main_jobs.assert_called_once_with(tasks)
+    self.mock_prepare_unscheduled_tasks.assert_called_once_with(
+        unscheduled_tasks)
+    self.assertEqual(result, prepared_tasks)