DEVPROD-17564 Generate multiversion binary selection tasks (#99)

sleaux · web-flow · commit f7bbc7567fe8 · 2025-05-08T20:48:53.000Z
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,4 +1,7 @@
 # Changelog
+## 3.0.0 - 2025-05-07
+* Generate multiversion binary selection tasks
+
 ## 2.0.0 - 2025-04-23
 * Allow splitting large tasks on required variants based on total test runtime
 * Use the AWS-sdk for accessing the s3 bucket of test statistics
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -2,7 +2,7 @@
 name = "mongo-task-generator"
 description = "Dynamically split evergreen tasks into subtasks for testing the 10gen/mongo project."
 license = "Apache-2.0"
-version = "2.0.0"
+version = "3.0.0"
 repository = "https://github.com/mongodb/mongo-task-generator"
 authors = ["DevProd Correctness Team <devprod-correctness-team@mongodb.com>"]
 edition = "2018"
@@ -23,7 +23,7 @@ serde = { version = "1.0.206", features = ["derive"] }
 serde_json = "1.0.124"
 serde_yaml = "0.9.33"
 shellexpand = "3.1.0"
-shrub-rs = "0.5.5"
+shrub-rs = "0.6.0"
 tokio = { version = "1.39.2", features = ["full"] }
 tracing = "0.1.40"
 tracing-subscriber = { version = "0.3.18", features = ["json", "fmt", "std"] }
diff --git a/src/evergreen/evg_config_utils.rs b/src/evergreen/evg_config_utils.rs
@@ -6,6 +6,7 @@ use lazy_static::lazy_static;
 use regex::Regex;
 use shrub_rs::models::commands::EvgCommand::Function;
 use shrub_rs::models::params::ParamValue;
+use shrub_rs::models::task::TaskDependency;
 use shrub_rs::models::{commands::FunctionCall, task::EvgTask, variant::BuildVariant};
 
 use crate::evergreen_names::{
@@ -101,6 +102,12 @@ pub trait EvgConfigUtils: Sync + Send {
     /// List of task names the task depends on.
     fn get_task_dependencies(&self, task: &EvgTask) -> Vec<String>;
 
+    fn get_task_ref_dependencies(
+        &self,
+        task_name: &str,
+        build_variant: &BuildVariant,
+    ) -> Option<Vec<TaskDependency>>;
+
     /// Lookup the given variable in the vars section of the 'generate resmoke task' func.
     ///
     /// # Arguments
@@ -430,6 +437,20 @@ impl EvgConfigUtils for EvgConfigUtilsImpl {
         dependencies.unwrap_or_default()
     }
 
+    fn get_task_ref_dependencies(
+        &self,
+        task_name: &str,
+        build_variant: &BuildVariant,
+    ) -> Option<Vec<TaskDependency>> {
+        for task_ref in &build_variant.tasks {
+            if task_ref.name == task_name {
+                let dependencies = task_ref.depends_on.clone();
+                return dependencies;
+            }
+        }
+        None
+    }
+
     /// Lookup the given variable in the vars section of the 'generate resmoke task' func.
     ///
     /// # Arguments
diff --git a/src/evergreen_names.rs b/src/evergreen_names.rs
@@ -44,6 +44,8 @@ pub const BURN_IN_TESTS: &str = "burn_in_tests_gen";
 pub const BURN_IN_TAGS: &str = "burn_in_tags_gen";
 /// Name of burn_in_tasks task.
 pub const BURN_IN_TASKS: &str = "burn_in_tasks_gen";
+/// Name of multiversion binary selection task.
+pub const MULTIVERSION_BINARY_SELECTION: &str = "select_multiversion_binaries";
 
 // Vars
 /// Variable that indicates a task is a fuzzer.
diff --git a/src/lib.rs b/src/lib.rs
@@ -25,7 +25,8 @@ use evergreen::{
 };
 use evergreen_names::{
     BURN_IN_TAGS, BURN_IN_TAG_COMPILE_TASK_DEPENDENCY, BURN_IN_TAG_INCLUDE_BUILD_VARIANTS,
-    BURN_IN_TASKS, BURN_IN_TESTS, ENTERPRISE_MODULE, GENERATOR_TASKS, UNIQUE_GEN_SUFFIX_EXPANSION,
+    BURN_IN_TASKS, BURN_IN_TESTS, ENTERPRISE_MODULE, GENERATOR_TASKS,
+    MULTIVERSION_BINARY_SELECTION, UNIQUE_GEN_SUFFIX_EXPANSION,
 };
 use generate_sub_tasks_config::GenerateSubTasksConfig;
 use resmoke::{
@@ -35,7 +36,7 @@ use resmoke::{
 use services::config_extraction::{ConfigExtractionService, ConfigExtractionServiceImpl};
 use shrub_rs::models::{
     project::EvgProject,
-    task::{EvgTask, TaskRef},
+    task::{EvgTask, TaskDependency, TaskRef},
     variant::{BuildVariant, DisplayTask},
 };
 use task_types::{
@@ -708,6 +709,7 @@ impl GenerateTasksService for GenerateTasksServiceImpl {
                 .infer_build_variant_platform(build_variant);
             let mut gen_config = GeneratedConfig::new();
             let mut generating_tasks = vec![];
+            let mut includes_multiversion_tasks = false;
             for task in &build_variant.tasks {
                 if task.name == BURN_IN_TAGS {
                     if self.gen_burn_in {
@@ -748,12 +750,50 @@ impl GenerateTasksService for GenerateTasksServiceImpl {
                     gen_config
                         .display_tasks
                         .push(generated_task.build_display_task());
+
+                    // If a task is a multiversion task and it has dependencies overriden on the task
+                    // reference of the build variant, add the MULTIVERSION_BINARY_SELECTION task to
+                    // those overrides too.
+                    let mut task_ref_dependencies = self
+                        .evg_config_utils
+                        .get_task_ref_dependencies(&task.name, build_variant);
+
+                    if generated_task.is_multiversion() && task_ref_dependencies.is_some() {
+                        task_ref_dependencies = Some(
+                            [
+                                task_ref_dependencies.unwrap(),
+                                vec![TaskDependency {
+                                    name: MULTIVERSION_BINARY_SELECTION.to_string(),
+                                    variant: None,
+                                }],
+                            ]
+                            .concat(),
+                        );
+                    } else {
+                        task_ref_dependencies = None;
+                    }
+
+                    if generated_task.is_multiversion() {
+                        includes_multiversion_tasks = true;
+                    }
+
                     gen_config
                         .gen_task_specs
-                        .extend(generated_task.build_task_ref(large_distro));
+                        .extend(generated_task.build_task_ref(large_distro, task_ref_dependencies));
                 }
             }
 
+            // If any generated task is multiversion, ensure the
+            // MULTIVERSION_BINARY_SELECTION task is added to the build variant.
+            if includes_multiversion_tasks {
+                gen_config.gen_task_specs.push(TaskRef {
+                    name: MULTIVERSION_BINARY_SELECTION.to_string(),
+                    distros: None,
+                    activate: Some(false),
+                    depends_on: None,
+                });
+            }
+
             if !generating_tasks.is_empty() {
                 // Put all the "_gen" tasks into a display task to hide them from view.
                 gen_config.display_tasks.push(DisplayTask {
@@ -1016,6 +1056,7 @@ pub async fn build_s3_client() -> aws_sdk_s3::Client {
 #[cfg(test)]
 mod tests {
     use rstest::rstest;
+    use shrub_rs::models::task::TaskDependency;
 
     use crate::{
         evergreen::evg_config_utils::MultiversionGenerateTaskConfig,
@@ -1143,6 +1184,14 @@ mod tests {
             todo!()
         }
 
+        fn get_task_ref_dependencies(
+            &self,
+            _task_name: &str,
+            _build_variant: &BuildVariant,
+        ) -> Option<Vec<TaskDependency>> {
+            todo!()
+        }
+
         fn get_gen_task_var<'a>(&self, _task: &'a EvgTask, _var: &str) -> Option<&'a str> {
             todo!()
         }
diff --git a/src/services/config_extraction.rs b/src/services/config_extraction.rs
@@ -7,10 +7,10 @@ use crate::{
     evergreen::evg_config_utils::EvgConfigUtils,
     evergreen_names::{
         CONTINUE_ON_FAILURE, FUZZER_PARAMETERS, IDLE_TIMEOUT, LARGE_DISTRO_EXPANSION,
-        LAST_VERSIONS_EXPANSION, MULTIVERSION, NO_MULTIVERSION_GENERATE_TASKS, NPM_COMMAND,
-        NUM_FUZZER_FILES, NUM_FUZZER_TASKS, REPEAT_SUITES, RESMOKE_ARGS, RESMOKE_JOBS_MAX,
-        SHOULD_SHUFFLE_TESTS, UNIQUE_GEN_SUFFIX_EXPANSION, USE_LARGE_DISTRO, USE_XLARGE_DISTRO,
-        XLARGE_DISTRO_EXPANSION,
+        LAST_VERSIONS_EXPANSION, MULTIVERSION, MULTIVERSION_BINARY_SELECTION,
+        NO_MULTIVERSION_GENERATE_TASKS, NPM_COMMAND, NUM_FUZZER_FILES, NUM_FUZZER_TASKS,
+        REPEAT_SUITES, RESMOKE_ARGS, RESMOKE_JOBS_MAX, SHOULD_SHUFFLE_TESTS,
+        UNIQUE_GEN_SUFFIX_EXPANSION, USE_LARGE_DISTRO, USE_XLARGE_DISTRO, XLARGE_DISTRO_EXPANSION,
     },
     generate_sub_tasks_config::GenerateSubTasksConfig,
     task_types::{
@@ -127,7 +127,15 @@ impl ConfigExtractionServiceImpl {
     ///
     /// List of tasks that should be included as dependencies.
     fn determine_task_dependencies(&self, task_def: &EvgTask) -> Vec<String> {
-        let depends_on = self.evg_config_utils.get_task_dependencies(task_def);
+        let mut depends_on = self.evg_config_utils.get_task_dependencies(task_def);
+
+        if self
+            .evg_config_utils
+            .get_task_tags(task_def)
+            .contains(MULTIVERSION)
+        {
+            depends_on.push(MULTIVERSION_BINARY_SELECTION.to_string());
+        }
 
         depends_on
             .into_iter()
@@ -399,12 +407,14 @@ mod tests {
     // Tests for determine_task_dependencies.
     #[rstest]
     #[case(
-        vec![], vec![]
+        vec![], vec![], vec![]
     )]
-    #[case(vec!["dependency_0", "dependency_1"], vec!["dependency_0", "dependency_1"])]
-    #[case(vec!["dependency_0", "generating_task"], vec!["dependency_0"])]
+    #[case(vec!["dependency_0", "dependency_1"], vec![], vec!["dependency_0", "dependency_1"])]
+    #[case(vec!["dependency_0", "generating_task"], vec![], vec!["dependency_0"])]
+    #[case(vec!["dependency_0", "generating_task"], vec!["multiversion".to_string()], vec!["dependency_0", MULTIVERSION_BINARY_SELECTION])]
     fn test_determine_task_dependencies(
         #[case] depends_on: Vec<&str>,
+        #[case] tags: Vec<String>,
         #[case] expected_deps: Vec<&str>,
     ) {
         let config_extraction_service = build_mocked_config_extraction_service();
@@ -418,6 +428,7 @@ mod tests {
                     })
                     .collect(),
             ),
+            tags: Some(tags),
             ..Default::default()
         };
 
diff --git a/src/task_types/burn_in_tests.rs b/src/task_types/burn_in_tests.rs
@@ -430,7 +430,7 @@ impl BurnInService for BurnInServiceImpl {
 
         gen_config
             .gen_task_specs
-            .extend(generated_task.build_task_ref(large_distro));
+            .extend(generated_task.build_task_ref(large_distro, None));
         gen_config
             .display_tasks
             .push(generated_task.build_display_task());
@@ -730,6 +730,14 @@ mod tests {
             todo!()
         }
 
+        fn get_task_ref_dependencies(
+            &self,
+            _task_name: &str,
+            _build_variant: &BuildVariant,
+        ) -> Option<Vec<TaskDependency>> {
+            todo!()
+        }
+
         fn get_gen_task_var<'a>(&self, _task: &'a EvgTask, _var: &str) -> Option<&'a str> {
             todo!()
         }
diff --git a/src/task_types/fuzzer_tasks.rs b/src/task_types/fuzzer_tasks.rs
@@ -439,7 +439,7 @@ mod tests {
             ],
         };
 
-        let task_refs = fuzzer_task.build_task_ref(Some("distro".to_string()));
+        let task_refs = fuzzer_task.build_task_ref(Some("distro".to_string()), None);
 
         for task in task_refs {
             assert_eq!(task.distros.as_ref(), None);
diff --git a/src/task_types/generated_suite.rs b/src/task_types/generated_suite.rs
@@ -1,8 +1,10 @@
 use shrub_rs::models::{
-    task::{EvgTask, TaskRef},
+    task::{EvgTask, TaskDependency, TaskRef},
     variant::DisplayTask,
 };
 
+use crate::evergreen_names::MULTIVERSION_BINARY_SELECTION;
+
 /// Definition of a generated sub task.
 #[derive(Clone, Debug, Default)]
 pub struct GeneratedSubTask {
@@ -48,18 +50,37 @@ pub trait GeneratedSuite: Sync + Send {
         }
     }
 
+    fn is_multiversion(&self) -> bool {
+        self.sub_tasks()
+            .iter()
+            .any(|task| match &task.evg_task.depends_on {
+                Some(deps) => deps
+                    .iter()
+                    .any(|dep| dep.name == MULTIVERSION_BINARY_SELECTION),
+                _ => false,
+            })
+    }
+
     /// Build a shrub task reference for this generated task.
-    fn build_task_ref(&self, distro: Option<String>) -> Vec<TaskRef> {
+    fn build_task_ref(
+        &self,
+        distro: Option<String>,
+        depends_on: Option<Vec<TaskDependency>>,
+    ) -> Vec<TaskRef> {
         self.sub_tasks()
             .iter()
             .map(|sub_task| {
                 let mut large_distro = None;
                 if sub_task.use_large_distro || sub_task.use_xlarge_distro {
                     large_distro = distro.clone();
                 }
-                sub_task
+                let mut task_ref = sub_task
                     .evg_task
-                    .get_reference(large_distro.map(|d| vec![d]), Some(false))
+                    .get_reference(large_distro.map(|d| vec![d]), Some(false));
+
+                task_ref.depends_on = depends_on.clone();
+
+                task_ref
             })
             .collect()
     }
diff --git a/src/task_types/resmoke_tasks.rs b/src/task_types/resmoke_tasks.rs
@@ -1053,7 +1053,7 @@ mod tests {
                 .collect(),
         };
 
-        let task_refs = gen_suite.build_task_ref(Some(distro.clone()));
+        let task_refs = gen_suite.build_task_ref(Some(distro.clone()), None);
 
         for (i, task) in task_refs.iter().enumerate() {
             assert_eq!(task.name, format!("sub_suite_name_{}", i));
diff --git a/tests/data/evergreen.yml b/tests/data/evergreen.yml
@@ -7942,6 +7942,9 @@ buildvariants:
     test_flags: >-
       --runAllFeatureFlagTests
       --excludeWithAnyTags=incompatible_with_shard_merge
+  depends_on:
+    - name: archive_dist_test
+      variant: enterprise-rhel-80-64-bit-dynamic-required
   tasks: &enterprise-rhel-80-64-bit-dynamic-all-feature-flags-tasks
   - name: change_streams_per_shard_cursor_passthrough
   - name: cqf
