SERVER-42227: Limit the number of tasks burn_in_test will generate

2 files changed, 31 insertions, 1 deletions

diff --git a/buildscripts/burn_in_tests.py b/buildscripts/burn_in_tests.py
index c3dd28f1a95..7e9b94939d4 100644
--- a/buildscripts/burn_in_tests.py
+++ b/buildscripts/burn_in_tests.py
@@ -44,6 +44,7 @@ AVG_TEST_TIME_MULTIPLIER = 3
 CONFIG_FILE = "../src/.evergreen.yml"
 REPEAT_SUITES = 2
 EVERGREEN_FILE = "etc/evergreen.yml"
+MAX_TASKS_TO_CREATE = 1000
 MIN_AVG_TEST_OVERFLOW_SEC = 60
 MIN_AVG_TEST_TIME_SEC = 5 * 60
 # The executor_file and suite_files defaults are required to make the suite resolver work
@@ -640,7 +641,12 @@ def create_generate_tasks_file(evg_api, options, tests_by_task):
     evg_config = Configuration()
     evg_config = create_generate_tasks_config(evg_api, evg_config, options, tests_by_task,
                                               include_gen_task=True)
-    _write_json_file(evg_config.to_map(), options.generate_tasks_file)
+    json_config = evg_config.to_map()
+    tasks_to_create = len(json_config.get('tasks', []))
+    if tasks_to_create > MAX_TASKS_TO_CREATE:
+        LOGGER.warning("Attempting to create more tasks than max(%d), aborting", tasks_to_create)
+        sys.exit(1)
+    _write_json_file(json_config, options.generate_tasks_file)
 
 
 def run_tests(no_exec, tests_by_task, resmoke_cmd, report_file):
diff --git a/buildscripts/tests/test_burn_in_tests.py b/buildscripts/tests/test_burn_in_tests.py
index 90b83758675..5ed523561a4 100644
--- a/buildscripts/tests/test_burn_in_tests.py
+++ b/buildscripts/tests/test_burn_in_tests.py
@@ -782,6 +782,30 @@ class TestCreateGenerateTasksFile(unittest.TestCase):
             self.assertEqual(len(execution_tasks), 1)
             self.assertEqual(execution_tasks[0], burn_in.BURN_IN_TESTS_GEN_TASK)
 
+    @patch("buildscripts.burn_in_tests._write_json_file")
+    @patch("buildscripts.burn_in_tests.sys.exit")
+    @patch("buildscripts.burn_in_tests.create_generate_tasks_config")
+    def test_cap_on_task_generate(self, gen_tasks_config_mock, exit_mock, write_mock):
+        evg_api = MagicMock()
+        options = MagicMock()
+        tests_by_task = MagicMock()
+
+        task_list = [f"task_{i}" for i in range(1005)]
+
+        evg_config = MagicMock()
+        evg_config.to_map.return_value = {
+            "tasks": task_list,
+        }
+
+        gen_tasks_config_mock.return_value = evg_config
+
+        exit_mock.side_effect = ValueError('exiting')
+        with self.assertRaises(ValueError):
+            burn_in.create_generate_tasks_file(evg_api, options, tests_by_task)
+
+        exit_mock.assert_called_once()
+        write_mock.assert_not_called()
+
 
 class UpdateReportDataTests(unittest.TestCase):
     def test_update_report_data_nofile(self):