diff --git a/src/clusterfuzz/_internal/tests/core/bot/startup/run_bot_test.py b/src/clusterfuzz/_internal/tests/core/bot/startup/run_bot_test.py index 585ec248c32..2418a79036c 100644 --- a/src/clusterfuzz/_internal/tests/core/bot/startup/run_bot_test.py +++ b/src/clusterfuzz/_internal/tests/core/bot/startup/run_bot_test.py @@ -89,7 +89,9 @@ def setUp(self): 'clusterfuzz._internal.bot.tasks.commands.process_command', 'clusterfuzz._internal.bot.tasks.update_task.run', 'clusterfuzz._internal.bot.tasks.update_task.track_revision', + 'python.bot.startup.run_bot.update_task_enabled', ]) + self.mock.update_task_enabled.return_value = True self.task = mock.MagicMock() self.task.payload.return_value = 'payload' @@ -107,6 +109,33 @@ def test_exception(self): self.assertFalse(clean_exit) self.assertEqual('payload', payload) + def test_max_executions(self): + """Test that the loop breaks after MAX_TASK_EXECUTIONS iterations.""" + from clusterfuzz._internal.system import environment + environment._initial_environment = None + os.environ['MAX_TASK_EXECUTIONS'] = '3' + + _, clean_exit, payload = run_bot.task_loop() + + self.assertEqual(3, self.mock.process_command.call_count) + self.assertTrue(clean_exit) + self.assertEqual('payload', payload) + + @mock.patch('clusterfuzz._internal.metrics.logs.log_fatal_and_exit') + def test_max_executions_invalid(self, mock_log_fatal_and_exit): + """Test that an invalid MAX_TASK_EXECUTIONS logs a fatal error and exits.""" + from clusterfuzz._internal.system import environment + environment._initial_environment = None + os.environ['MAX_TASK_EXECUTIONS'] = 'invalid' + mock_log_fatal_and_exit.side_effect = SystemExit + + with self.assertRaises(SystemExit): + run_bot.task_loop() + + mock_log_fatal_and_exit.assert_any_call( + 'Invalid value for MAX_TASK_EXECUTIONS: invalid') + self.assertEqual(0, self.mock.process_command.call_count) + class LeaseAllTasksTest(unittest.TestCase): """Tests for lease_all_tasks.""" diff --git a/src/python/bot/startup/run_bot.py b/src/python/bot/startup/run_bot.py index 6825a28e149..671ddadd735 100644 --- a/src/python/bot/startup/run_bot.py +++ b/src/python/bot/startup/run_bot.py @@ -114,12 +114,27 @@ def schedule_utask_mains(): task.pubsub_task.cancel_lease_ack() +def _get_max_task_executions(): # pylint: disable=inconsistent-return-statements + """Returns the MAX_TASK_EXECUTIONS limit as an int, or None if + invalid/unset.""" + val = environment.get_value('MAX_TASK_EXECUTIONS') + if not val: + return None + try: + return int(val) + except ValueError: + logs.log_fatal_and_exit(f'Invalid value for MAX_TASK_EXECUTIONS: {val}') + + def task_loop(): """Executes tasks indefinitely.""" # Defer heavy task imports to prevent issues with multiprocessing.Process from clusterfuzz._internal.bot.tasks import commands clean_exit = False + execution_count = 0 + max_task_executions = _get_max_task_executions() + while True: stacktrace = '' exception_occurred = False @@ -191,6 +206,14 @@ def task_loop(): time.sleep(utils.random_number(1, failure_wait_interval)) break + execution_count += 1 + if max_task_executions and execution_count >= max_task_executions: + logs.info( + 'Reached MAX_TASK_EXECUTIONS limit. Exiting.', + max_task_executions=max_task_executions) + clean_exit = True + break + task_payload = task.payload() if task else None return stacktrace, clean_exit, task_payload