diff --git a/src/clusterfuzz/_internal/base/tasks/__init__.py b/src/clusterfuzz/_internal/base/tasks/__init__.py index 2688d26b79..7e74c3d5c9 100644 --- a/src/clusterfuzz/_internal/base/tasks/__init__.py +++ b/src/clusterfuzz/_internal/base/tasks/__init__.py @@ -25,6 +25,7 @@ from google.cloud import monitoring_v3 +from clusterfuzz._internal import swarming from clusterfuzz._internal.base import external_tasks from clusterfuzz._internal.base import feature_flags from clusterfuzz._internal.base import memoize @@ -932,7 +933,11 @@ def add_task(command, queue=None, wait_time=None, extra_info=None): - """Add a new task to the job queue.""" + """Adds a new task to the job queue. + + When `queue` is None, tasks for Swarming-applicable jobs are routed to the + job's platform queue instead of the default Linux queue. + """ if wait_time is None: wait_time = random.randint(1, TASK_CREATION_WAIT_INTERVAL) @@ -957,6 +962,9 @@ def add_task(command, external_tasks.add_external_task(command, argument, job) return + if queue is None and swarming.is_swarming_applicable(job_type, job=job): + queue = queue_for_platform(job.platform) + # Add the task. eta = utils.utcnow() + datetime.timedelta(seconds=wait_time) extra_info = extra_info or {} diff --git a/src/clusterfuzz/_internal/swarming/__init__.py b/src/clusterfuzz/_internal/swarming/__init__.py index 6358732451..0f775e2fa2 100644 --- a/src/clusterfuzz/_internal/swarming/__init__.py +++ b/src/clusterfuzz/_internal/swarming/__init__.py @@ -85,6 +85,13 @@ def is_swarming_job(job_name: str, job: data_types.Job | None = None) -> bool: return _get_instance_spec(swarming_config, job) is not None +def is_swarming_applicable(job_name: str, + job: data_types.Job | None = None) -> bool: + """Returns True if swarming remote execution is enabled and the job is a + swarming job.""" + return is_swarming_enabled() and is_swarming_job(job_name, job=job) + + def _get_instance_spec(swarming_config: local_config.SwarmingConfig, job: data_types.Job) -> dict | None: return swarming_config.get('mapping').get(job.platform, None) diff --git a/src/clusterfuzz/_internal/tests/core/base/tasks/tasks_test.py b/src/clusterfuzz/_internal/tests/core/base/tasks/tasks_test.py index ee5c0dd653..f18a46685d 100644 --- a/src/clusterfuzz/_internal/tests/core/base/tasks/tasks_test.py +++ b/src/clusterfuzz/_internal/tests/core/base/tasks/tasks_test.py @@ -696,3 +696,55 @@ def test_subqueue_and_base_platform_not_in_string_value(self): mock_flag.return_value = mock_flag_obj self.assertIsNone(tasks.get_task()) self.mock.get_fuzz_task.assert_not_called() + + +@test_utils.with_cloud_emulators('datastore') +class SwarmingNonFuzzTaskRoutingTest(unittest.TestCase): + """Tests for routing non-fuzz tasks of Swarming jobs to platform queues.""" + + def setUp(self): + helpers.patch_environ(self) + helpers.patch(self, [ + 'clusterfuzz._internal.base.tasks.bulk_add_tasks', + ]) + data_types.Job( + name='android_emulator_swarming', + platform='ANDROID_EMULATOR', + environment_string='IS_SWARMING_JOB = True\n').put() + data_types.Job( + name='android_emulator_non_swarming', + platform='ANDROID_EMULATOR', + environment_string='').put() + self.flag = data_types.FeatureFlag( + id='swarming_remote_execution', enabled=True) + self.flag.put() + + def _queue_used(self): + return self.mock.bulk_add_tasks.call_args[1]['queue'] + + def test_swarming_job_routes_to_platform_queue(self): + """Tests that a Swarming job task with no explicit queue is routed to its + platform queue when SWARMING_REMOTE_EXECUTION is enabled.""" + tasks.add_task('minimize', '123', 'android_emulator_swarming') + self.assertEqual(self._queue_used(), 'jobs-android-emulator') + + def test_swarming_disabled_keeps_default_queue(self): + """Tests that Swarming jobs keep the default queue when + SWARMING_REMOTE_EXECUTION is disabled.""" + self.flag.enabled = False + self.flag.put() + tasks.add_task('minimize', '123', 'android_emulator_swarming') + self.assertIsNone(self._queue_used()) + + def test_non_swarming_job_keeps_default_queue(self): + """Tests that non-Swarming jobs keep the default queue even when + SWARMING_REMOTE_EXECUTION is enabled.""" + tasks.add_task('minimize', '123', 'android_emulator_non_swarming') + self.assertIsNone(self._queue_used()) + + def test_explicit_queue_is_not_overridden(self): + """Tests that an explicitly provided queue is preserved for Swarming + jobs.""" + tasks.add_task( + 'minimize', '123', 'android_emulator_swarming', queue='custom-queue') + self.assertEqual(self._queue_used(), 'custom-queue')