Add MAX_EXECUTIONS environment variable limit

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."""

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