SERVER-43186: Support max tests per generated sub-suite

(cherry picked from commit cdf0a059437f1ed417a22b3c574e2f73afc04233)
(cherry picked from commit b326fd656716e95418e563ff12368a3015994b5e)

2 files changed, 77 insertions, 6 deletions

diff --git a/buildscripts/evergreen_generate_resmoke_tasks.py b/buildscripts/evergreen_generate_resmoke_tasks.py
index d9a15215927..3c2f61da7bd 100755
--- a/buildscripts/evergreen_generate_resmoke_tasks.py
+++ b/buildscripts/evergreen_generate_resmoke_tasks.py
@@ -63,6 +63,7 @@ REQUIRED_CONFIG_KEYS = {
 }
 
 DEFAULT_CONFIG_VALUES = {
+    "max_tests_per_suite": 100,
     "resmoke_args": "",
     "resmoke_repeat_suites": 1,
     "run_multiple_jobs": "true",
@@ -74,6 +75,7 @@ DEFAULT_CONFIG_VALUES = {
 CONFIG_FORMAT_FN = {
     "fallback_num_sub_suites": int,
     "max_sub_suites": int,
+    "max_tests_per_suite": int,
     "target_resmoke_time": int,
 }
 
@@ -212,20 +214,56 @@ def divide_remaining_tests_among_suites(remaining_tests_runtimes, suites):
             suite_idx = 0
 
 
-def divide_tests_into_suites(tests_runtimes, max_time_seconds, max_suites=None):
+def _new_suite_needed(current_suite, test_runtime, max_suite_runtime, max_tests_per_suite):
+    """
+    Check if a new suite should be created for the given suite.
+
+    :param current_suite: Suite currently being added to.
+    :param test_runtime: Runtime of test being added.
+    :param max_suite_runtime: Max runtime of a single suite.
+    :param max_tests_per_suite: Max number of tests in a suite.
+    :return: True if a new test suite should be created.
+    """
+    if current_suite.get_runtime() + test_runtime > max_suite_runtime:
+        # Will adding this test put us over the target runtime?
+        return True
+
+    if max_tests_per_suite and current_suite.get_test_count() + 1 > max_tests_per_suite:
+        # Will adding this test put us over the max number of tests?
+        return True
+
+    return False
+
+
+def divide_tests_into_suites(tests_runtimes, max_time_seconds, max_suites=None,
+                             max_tests_per_suite=None):
     """
     Divide the given tests into suites.
 
-    Each suite should be able to execute in less than the max time specified.
+    Each suite should be able to execute in less than the max time specified. If a single
+    test has a runtime greater than `max_time_seconds`, it will be run in a suite on its own.
+
+    If max_suites is reached before assigning all tests to a suite, the remaining tests will be
+    divided up among the created suites.
+
+    Note: If `max_suites` is hit, suites may have more tests than `max_tests_per_suite` and may have
+    runtimes longer than `max_time_seconds`.
+
+    :param tests_runtimes: List of tuples containing test names and test runtimes.
+    :param max_time_seconds: Maximum runtime to add to a single bucket.
+    :param max_suites: Maximum number of suites to create.
+    :param max_tests_per_suite: Maximum number of tests to add to a single suite.
+    :return: List of Suite objects representing grouping of tests.
     """
     suites = []
     current_suite = Suite()
     last_test_processed = len(tests_runtimes)
-    LOGGER.debug("Determines suites for runtime", max_runtime_seconds=max_time_seconds)
+    LOGGER.debug("Determines suites for runtime", max_runtime_seconds=max_time_seconds,
+                 max_suites=max_suites, max_tests_per_suite=max_tests_per_suite)
     for idx, (test_file, runtime) in enumerate(tests_runtimes):
         LOGGER.debug("Adding test", test=test_file, test_runtime=runtime)
-        if current_suite.get_runtime() + runtime > max_time_seconds:
-            LOGGER.debug("Runtime greater than boundary", suite_runtime=current_suite.get_runtime(),
+        if _new_suite_needed(current_suite, runtime, max_time_seconds, max_tests_per_suite):
+            LOGGER.debug("Finished suite", suite_runtime=current_suite.get_runtime(),
                          test_runtime=runtime, max_time=max_time_seconds)
             if current_suite.get_test_count() > 0:
                 suites.append(current_suite)
@@ -631,7 +669,8 @@ class GenerateSubSuites(object):
             return self.calculate_fallback_suites()
         self.test_list = [info[0] for info in tests_runtimes]
         return divide_tests_into_suites(tests_runtimes, execution_time_secs,
-                                        self.config_options.max_sub_suites)
+                                        self.config_options.max_sub_suites,
+                                        self.config_options.max_tests_per_suite)
 
     def filter_existing_tests(self, tests_runtimes):
         """Filter out tests that do not exist in the filesystem."""
diff --git a/buildscripts/tests/test_evergreen_generate_resmoke_tasks.py b/buildscripts/tests/test_evergreen_generate_resmoke_tasks.py
index bd8f4b0a4a9..e108bfbf727 100644
--- a/buildscripts/tests/test_evergreen_generate_resmoke_tasks.py
+++ b/buildscripts/tests/test_evergreen_generate_resmoke_tasks.py
@@ -244,6 +244,37 @@ class DivideTestsIntoSuitesByMaxtimeTest(unittest.TestCase):
             total_tests += suite.get_test_count()
         self.assertEqual(total_tests, len(tests_runtimes))
 
+    def test_max_tests_per_suites_is_one(self):
+        max_time = 5
+        num_tests = 10
+        tests_runtimes = [("tests_{i}".format(i=i), i) for i in range(num_tests)]
+
+        suites = under_test.divide_tests_into_suites(tests_runtimes, max_time,
+                                                     max_tests_per_suite=1)
+
+        self.assertEqual(len(suites), num_tests)
+
+    def test_max_tests_per_suites_is_less_than_number_of_tests(self):
+        max_time = 100
+        num_tests = 10
+        tests_runtimes = [("tests_{i}".format(i=i), 1) for i in range(num_tests)]
+
+        suites = under_test.divide_tests_into_suites(tests_runtimes, max_time,
+                                                     max_tests_per_suite=2)
+
+        self.assertEqual(len(suites), num_tests // 2)
+
+    def test_max_suites_overrides_max_tests_per_suite(self):
+        max_time = 100
+        num_tests = 10
+        max_suites = 2
+        tests_runtimes = [("tests_{i}".format(i=i), 1) for i in range(num_tests)]
+
+        suites = under_test.divide_tests_into_suites(tests_runtimes, max_time,
+                                                     max_suites=max_suites, max_tests_per_suite=2)
+
+        self.assertEqual(len(suites), max_suites)
+
 
 class SuiteTest(unittest.TestCase):
     def test_adding_tests_increases_count_and_runtime(self):
@@ -620,6 +651,7 @@ class GenerateSubSuitesTest(unittest.TestCase):
         options = MagicMock()
         options.target_resmoke_time = 10
         options.fallback_num_sub_suites = 2
+        options.max_tests_per_suite = None
         return options
 
     @staticmethod