SERVER-47274: Refactor task generation in evergreen

1 files changed, 26 insertions, 91 deletions

diff --git a/buildscripts/patch_builds/task_generation.py b/buildscripts/patch_builds/task_generation.py
index 576ac95b15a..2ca8a7e184f 100644
--- a/buildscripts/patch_builds/task_generation.py
+++ b/buildscripts/patch_builds/task_generation.py
@@ -1,24 +1,19 @@
 """Utilities to help generate evergreen tasks."""
-from typing import Optional, List
+from __future__ import annotations
 
-from shrub.command import CommandDefinition
-from shrub.config import Configuration
-from shrub.operations import CmdTimeoutUpdate
-from shrub.task import TaskDependency
-from shrub.variant import TaskSpec, DisplayTaskDefinition
+from typing import Any, Dict, Optional, List
 
+from shrub.v2 import FunctionCall, ShrubProject
+from shrub.v2.command import timeout_update, ShrubCommand
+from structlog import get_logger
 
-def _cmd_by_name(cmd_name):
-    """
-    Create a command definition of a function with the given name.
-
-    :param cmd_name: Name of function.
-    :return: Command Definition for function.
-    """
-    return CommandDefinition().function(cmd_name)
+LOGGER = get_logger(__name__)
+MAX_SHRUB_TASKS_FOR_SINGLE_TASK = 1000
 
 
-def resmoke_commands(run_tests_fn_name, run_tests_vars, timeout_info, use_multiversion=None):
+def resmoke_commands(run_tests_fn_name: str, run_tests_vars: Dict[str, Any],
+                     timeout_info: TimeoutInfo,
+                     use_multiversion: Optional[str] = None) -> List[ShrubCommand]:
     """
     Create a list of commands to run a resmoke task.
 
@@ -30,9 +25,9 @@ def resmoke_commands(run_tests_fn_name, run_tests_vars, timeout_info, use_multiv
     """
     commands = [
         timeout_info.cmd,
-        _cmd_by_name("do setup"),
-        _cmd_by_name("do multiversion setup") if use_multiversion else None,
-        _cmd_by_name(run_tests_fn_name).vars(run_tests_vars),
+        FunctionCall("do setup"),
+        FunctionCall("do multiversion setup") if use_multiversion else None,
+        FunctionCall(run_tests_fn_name, run_tests_vars),
     ]
 
     return [cmd for cmd in commands if cmd]
@@ -75,13 +70,7 @@ class TimeoutInfo(object):
     def cmd(self):
         """Create a command that sets timeouts as specified."""
         if not self.use_defaults:
-            timeout_cmd = CmdTimeoutUpdate()
-            if self.timeout:
-                timeout_cmd.timeout(self.timeout)
-
-            if self.exec_timeout:
-                timeout_cmd.exec_timeout(self.exec_timeout)
-            return timeout_cmd.validate().resolve()
+            return timeout_update(exec_timeout_secs=self.exec_timeout, timeout_secs=self.timeout)
 
         return None
 
@@ -92,70 +81,16 @@ class TimeoutInfo(object):
         return f"<exec_timeout={self.exec_timeout}, timeout={self.timeout}>"
 
 
-class TaskList(object):
-    """A list of evergreen tasks to be generated together."""
-
-    def __init__(self, evg_config: Configuration):
-        """
-        Create a list of evergreen tasks to create.
-
-        :param evg_config: Evergreen configuration to add tasks to.
-        """
-        self.evg_config = evg_config
-        self.task_specs = []
-        self.task_names = []
-
-    def add_task(self, name: str, commands: [CommandDefinition],
-                 depends_on: Optional[List[str]] = None, distro: Optional[str] = None):
-        """
-        Add a new task to the task list.
-
-        :param name: Name of task to add.
-        :param commands: List of commands comprising task.
-        :param depends_on: Any dependencies for the task.
-        :param distro: Distro task should be run on.
-        """
-        task = self.evg_config.task(name)
-        task.commands(commands)
-
-        if depends_on:
-            for dep in depends_on:
-                task.dependency(TaskDependency(dep))
-
-        task_spec = TaskSpec(name)
-        if distro:
-            task_spec.distro(distro)
-        self.task_specs.append(task_spec)
-        self.task_names.append(name)
-
-    def display_task(self, display_name: str, existing_tasks: Optional[List[str]] = None) \
-            -> DisplayTaskDefinition:
-        """
-        Create a display task for the list of tasks.
-
-        Note: This function should be called after all calls to `add_task` have been done.
-
-        :param display_name: Name of display tasks.
-        :param existing_tasks: Any existing tasks that should be part of the display task.
-        :return: Display task object.
-        """
-        execution_tasks = self.task_names
-        if existing_tasks:
-            execution_tasks.extend(existing_tasks)
-
-        display_task = DisplayTaskDefinition(display_name).execution_tasks(execution_tasks)
-        return display_task
-
-    def add_to_variant(self, variant_name: str, display_name: Optional[str] = None,
-                       existing_tasks: Optional[List[str]] = None):
-        """
-        Add this task list to a build variant.
+def validate_task_generation_limit(shrub_project: ShrubProject) -> bool:
+    """
+    Determine if this shrub configuration generates less than the limit.
 
-        :param variant_name: Variant to add to.
-        :param display_name: Display name to add tasks under.
-        :param existing_tasks: Any existing tasks that should be added to the display group.
-        """
-        variant = self.evg_config.variant(variant_name)
-        variant.tasks(self.task_specs)
-        if display_name:
-            variant.display_task(self.display_task(display_name, existing_tasks))
+    :param shrub_project: Shrub configuration to validate.
+    :return: True if the configuration is under the limit.
+    """
+    tasks_to_create = len(shrub_project.all_tasks())
+    if tasks_to_create > MAX_SHRUB_TASKS_FOR_SINGLE_TASK:
+        LOGGER.warning("Attempting to create more tasks than max, aborting", tasks=tasks_to_create,
+                       max=MAX_SHRUB_TASKS_FOR_SINGLE_TASK)
+        return False
+    return True