DEVPROD-16911 make number of sub-tasks configurable (#94)

Author: sleaux

Cargo.lock

@@ -2,10 +2,11 @@
 name = "mongo-task-generator"
 description = "Dynamically split evergreen tasks into subtasks for testing the 10gen/mongo project."
 license = "Apache-2.0"
-version = "1.0.0"
+version = "1.1.0"
 repository = "https://github.com/mongodb/mongo-task-generator"
 authors = ["DevProd Correctness Team <[email protected]>"]
 edition = "2018"
+rust-version = "1.75"
 
 [dependencies]
 anyhow = "1.0.86"

Cargo.toml

@@ -99,7 +99,7 @@ Looking at a [sample resmoke-based](https://github.com/mongodb/mongo/blob/852c5d
 ```
 
 Like fuzzer tasks, task generation is indicated by including the `"generate resmoke tasks"` function.
-Additionally, the 3 parameters here will impact how the task is generated.
+Additionally, the 4 parameters here will impact how the task is generated.
 
 * **suite**: By default, the name of the task (with the `_gen` suffix stripped off) will be used
   to determine which resmoke suite to base the generated sub-tasks on. This can be overwritten with
@@ -112,6 +112,7 @@ Additionally, the 3 parameters here will impact how the task is generated.
   variable is set to `"true"`, certain tasks will use an even larger distro that can be defined with
   the `xlarge_distro_name` expansion in the build variant. When the `xlarge_distro_name` expansion
   is not defined, it will fallback to the defined `large_distro_name` expansion in the build variant
+* **num_tasks**: The number of generated sub-tasks to split into. (Default 5).
 
 **Note**: If a task has the `use_large_distro` value defined, but is added to a build variant
 without a `large_distro_name`, it will trigger a failure. This can be supported by using the

docs/generating_tasks.md

@@ -58,7 +58,7 @@ mod utils;
 const BURN_IN_TESTS_PREFIX: &str = "burn_in_tests";
 const BURN_IN_TASKS_PREFIX: &str = "burn_in_tasks";
 const BURN_IN_BV_SUFFIX: &str = "generated-by-burn-in-tags";
-const MAX_SUB_TASKS_PER_TASK: usize = 5;
+const DEFAULT_SUB_TASKS_PER_TASK: usize = 5;
 
 type GenTaskCollection = HashMap<String, Box<dyn GeneratedSuite>>;
 
@@ -206,11 +206,8 @@ impl Dependencies {
                 32,
             )));
         let enterprise_dir = evg_config_service.get_module_dir(ENTERPRISE_MODULE);
-        let gen_resmoke_config = GenResmokeConfig::new(
-            MAX_SUB_TASKS_PER_TASK,
-            execution_config.use_task_split_fallback,
-            enterprise_dir,
-        );
+        let gen_resmoke_config =
+            GenResmokeConfig::new(execution_config.use_task_split_fallback, enterprise_dir);
         let gen_resmoke_task_service = Arc::new(GenResmokeTaskServiceImpl::new(
             task_history_service,
             discovery_service,

src/lib.rs

@@ -18,6 +18,7 @@ use crate::{
         multiversion::MultiversionService, resmoke_tasks::ResmokeGenParams,
     },
     utils::task_name::remove_gen_suffix,
+    DEFAULT_SUB_TASKS_PER_TASK,
 };
 
 /// Interface for performing extractions of evergreen project configuration.
@@ -248,6 +249,13 @@ impl ConfigExtractionService for ConfigExtractionServiceImpl {
                 .evg_config_utils
                 .lookup_build_variant_expansion(UNIQUE_GEN_SUFFIX_EXPANSION, variant);
         }
+        let num_tasks = match self
+            .evg_config_utils
+            .get_gen_task_var(task_def, "num_tasks")
+        {
+            Some(str) => str.parse().unwrap(),
+            _ => DEFAULT_SUB_TASKS_PER_TASK,
+        };
 
         Ok(ResmokeGenParams {
             task_name,
@@ -289,6 +297,7 @@ impl ConfigExtractionService for ConfigExtractionServiceImpl {
             pass_through_vars: self.evg_config_utils.get_gen_task_vars(task_def),
             platform,
             gen_task_suffix,
+            num_tasks,
         })
     }