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,36 @@ def test_exception(self):
     self.assertFalse(clean_exit)
     self.assertEqual('payload', payload)
 
+  def test_max_executions(self):
+    """Test that the loop breaks after MAX_EXECUTIONS iterations."""
+    from clusterfuzz._internal.system import environment
+    environment._initial_environment = None
+    os.environ['MAX_EXECUTIONS'] = '3'
+    # Use side_effect to avoid raising an exception so it loops cleanly.
+    self.mock.process_command.side_effect = None
+
+    _, 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.error')
+  def test_max_executions_invalid(self, mock_error):
+    """Test that an invalid MAX_EXECUTIONS logs an error and is removed."""
+    from clusterfuzz._internal.system import environment
+    environment._initial_environment = None
+    os.environ['MAX_EXECUTIONS'] = 'invalid'
+
+    # We need the loop to exit otherwise it will run indefinitely because we removed the limit.
+    # We can make process_command succeed once, then raise an exception so it exits.
+    self.mock.process_command.side_effect = [None, Exception('exit')]
+
+    _, _, _ = run_bot.task_loop()
+
+    mock_error.assert_any_call('Invalid value for MAX_EXECUTIONS: invalid')
+    self.assertEqual(2, 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,26 @@ def schedule_utask_mains():
         task.pubsub_task.cancel_lease_ack()
 
 
+def _get_max_executions():
+  """Returns the MAX_EXECUTIONS limit as an int, or None if invalid/unset."""
+  val = environment.get_value('MAX_EXECUTIONS')
+  if not val:
+    return None
+  try:
+    return int(val)
+  except ValueError:
+    logs.error(f'Invalid value for MAX_EXECUTIONS: {val}')
+    environment.remove_key('MAX_EXECUTIONS')
+    return None
+
+
 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
   while True:
     stacktrace = ''
     exception_occurred = False
@@ -191,6 +205,13 @@ def task_loop():
       time.sleep(utils.random_number(1, failure_wait_interval))
       break
 
+    execution_count += 1
+    max_executions = _get_max_executions()
+    if max_executions and execution_count >= max_executions:
+      logs.info(f'Reached MAX_EXECUTIONS limit ({max_executions}). Exiting.')
+      clean_exit = True
+      break
+
   task_payload = task.payload() if task else None
   return stacktrace, clean_exit, task_payload