DEVPROD-18259 Generate tasks for suites that run via bazel (#104)

Author: sleaux

CHANGELOG.md

@@ -1,4 +1,7 @@
 # Changelog
+## 3.3.0 - 2025-06-10
+* Generate tasks for suites that run via bazel
+
 ## 3.0.1 - 2025-05-22
 * Add multiversion binary selection dependency to burn-in tasks when needed.
 

Cargo.lock

@@ -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 = "3.2.0"
+version = "3.3.0"
 repository = "https://github.com/mongodb/mongo-task-generator"
 authors = ["DevProd Correctness Team <[email protected]>"]
 edition = "2018"

Cargo.toml

@@ -117,14 +117,27 @@ Options:
           Disable evergreen task-history queries and use task splitting fallback
       --resmoke-command <RESMOKE_COMMAND>
           Command to invoke resmoke [default: "python buildscripts/resmoke.py"]
+      --include-fully-disabled-feature-tests
+          If the generator should include tests that are tagged with fully disabled features
       --generate-sub-tasks-config <GENERATE_SUB_TASKS_CONFIG>
           File containing configuration for generating sub-tasks
       --burn-in
           Generate burn_in related tasks
       --burn-in-tests-command <BURN_IN_TESTS_COMMAND>
           Command to invoke burn_in_tests [default: "python buildscripts/burn_in_tests.py run"]
-      --s3-test-stats-bucket <S3_TEST_STATS_BUCKETT>
+      --s3-test-stats-bucket <S3_TEST_STATS_BUCKET>
           S3 bucket to get test stats from [default: mongo-test-stats]
+      --test-runtime-per-required-subtask <TEST_RUNTIME_PER_REQUIRED_SUBTASK>
+          [default: 3600]
+      --large-required-task-runtime-threshold <LARGE_REQUIRED_TASK_RUNTIME_THRESHOLD>
+          [default: 7200]
+      --default-subtasks-per-task <DEFAULT_SUBTASKS_PER_TASK>
+          [default: 5]
+      --max-subtasks-per-task <MAX_SUBTASKS_PER_TASK>
+          [default: 10]
+      --bazel-suite-configs <BAZEL_SUITE_CONFIGS>
+          YAML file mapping mapping bazel target names of suite configs to their file location location
+      
   -h, --help
           Print help
 ```

README.md

@@ -30,6 +30,8 @@ pub struct MultiversionGenerateTaskConfig {
     pub suite_name: String,
     /// Old version to run testing against.
     pub old_version: String,
+    /// The bazel test target, if it is a bazel-based resmoke task.
+    pub bazel_target: Option<String>,
 }
 
 pub trait EvgConfigUtils: Sync + Send {
@@ -361,8 +363,12 @@ impl EvgConfigUtils for EvgConfigUtilsImpl {
         let generated_task_name = remove_gen_suffix(&task.name);
 
         if let Some(vars) = optional_vars {
-            if let Some(ParamValue::String(suite_var)) = vars.get("suite") {
-                suite_var
+            if let Some(ParamValue::String(suite)) = vars.get("suite") {
+                if is_bazel_suite(suite) {
+                    get_bazel_suite_name(suite)
+                } else {
+                    suite
+                }
             } else {
                 generated_task_name
             }
@@ -406,11 +412,18 @@ impl EvgConfigUtils for EvgConfigUtilsImpl {
             get_func_vars_by_name(task, INITIALIZE_MULTIVERSION_TASKS)
         {
             let mut multiversion_generate_tasks = vec![];
-            for (suite_name, old_version) in multiversion_task_map {
+            for (suite, old_version) in multiversion_task_map {
+                let (suite_name, bazel_target) = if is_bazel_suite(suite) {
+                    (get_bazel_suite_name(suite).to_string(), Some(suite.clone()))
+                } else {
+                    (suite.clone(), None)
+                };
+
                 if let ParamValue::String(value) = old_version {
                     multiversion_generate_tasks.push(MultiversionGenerateTaskConfig {
-                        suite_name: suite_name.clone(),
+                        suite_name,
                         old_version: value.clone(),
+                        bazel_target,
                     });
                 }
             }
@@ -844,6 +857,33 @@ fn get_resmoke_vars(task: &EvgTask) -> Option<&HashMap<String, ParamValue>> {
     return get_func_vars_by_name(task, RUN_RESMOKE_TESTS);
 }
 
+/// Checks if a Resmoke suite is a bazel target.
+///
+/// # Arguments
+///
+/// * `suite` - A suite name from Evergreen YAML.
+///
+/// # Returns
+///
+/// True if the suite looks like a bazel target (e.g. starts with `//`).
+pub fn is_bazel_suite(suite: &str) -> bool {
+    suite.starts_with("//")
+}
+
+/// Get a suite name from a bazel target.
+///
+/// # Arguments
+///
+/// * `target` - A bazel target.
+///
+/// # Returns
+///
+/// A useful suite name, just the name of the target without the bazel package prefix.
+pub fn get_bazel_suite_name(target: &str) -> &str {
+    let (_, name) = target.rsplit_once(':').unwrap();
+    name
+}
+
 #[cfg(test)]
 mod tests {
     use std::collections::BTreeMap;
@@ -1413,10 +1453,12 @@ mod tests {
             MultiversionGenerateTaskConfig {
                 suite_name: "mv_suite1_last_continuous".to_string(),
                 old_version: "last-continuous".to_string(),
+                bazel_target: None,
             },
             MultiversionGenerateTaskConfig {
                 suite_name: "mv_suite1_last_lts".to_string(),
                 old_version: "last-lts".to_string(),
+                bazel_target: None,
             },
         ];
         assert!(multiversion_generate_tasks

src/evergreen/evg_config_utils.rs

@@ -17,6 +17,8 @@ pub const SETUP_JSTESTFUZZ: &str = "setup jstestfuzz";
 pub const RUN_FUZZER: &str = "run jstestfuzz";
 /// Function to run generated tasks.
 pub const RUN_GENERATED_TESTS: &str = "run generated tests";
+/// Function to run generated tasks via 'bazel test'.
+pub const RUN_GENERATED_TESTS_VIA_BAZEL: &str = "run generated tests via bazel";
 
 // Function for multi-version tests.
 /// Function to do setup for multi-version testing.
@@ -67,6 +69,8 @@ pub const MULTIVERSION_EXCLUDE_TAG: &str = "multiversion_exclude_tags_version";
 pub const REQUIRE_MULTIVERSION_SETUP: &str = "require_multiversion_setup";
 /// Arguments to pass to resmoke command.
 pub const RESMOKE_ARGS: &str = "resmoke_args";
+/// Arguments to pass to bazel command.
+pub const BAZEL_ARGS: &str = "bazel_args";
 /// Name of suite being executed.
 pub const SUITE_NAME: &str = "suite";
 /// Location where generation task configuration is stored in S3.

src/evergreen_names.rs

@@ -51,6 +51,8 @@ use tokio::{runtime::Handle, task::JoinHandle, time};
 use tracing::{event, Level};
 use utils::fs_service::FsServiceImpl;
 
+use crate::resmoke::resmoke_proxy::BazelConfigs;
+
 mod evergreen;
 mod evergreen_names;
 mod generate_sub_tasks_config;
@@ -150,6 +152,7 @@ pub struct ExecutionConfiguration<'a> {
     /// S3 bucket to get test stats from.
     pub s3_test_stats_bucket: &'a str,
     pub subtask_limits: SubtaskLimits,
+    pub bazel_suite_configs: Option<PathBuf>,
 }
 
 #[derive(Debug, Clone)]
@@ -193,10 +196,15 @@ impl Dependencies {
         s3_client: aws_sdk_s3::Client,
     ) -> Result<Self> {
         let fs_service = Arc::new(FsServiceImpl::new());
+        let bazel_suite_configs = match execution_config.bazel_suite_configs {
+            Some(path) => BazelConfigs::from_yaml_file(&path).unwrap_or_default(),
+            None => BazelConfigs::default(),
+        };
         let discovery_service = Arc::new(ResmokeProxy::new(
             execution_config.resmoke_command,
             execution_config.skip_covered_tests,
             execution_config.include_fully_disabled_feature_tests,
+            bazel_suite_configs,
         ));
         let multiversion_service = Arc::new(MultiversionServiceImpl::new(
             discovery_service.get_multiversion_config()?,
@@ -240,6 +248,11 @@ impl Dependencies {
             fs_service,
             gen_resmoke_config,
             execution_config.subtask_limits,
+            execution_config
+                .target_directory
+                .to_str()
+                .unwrap_or("")
+                .to_string(),
         ));
         let gen_task_service = Arc::new(GenerateTasksServiceImpl::new(
             evg_config_service,

src/lib.rs

@@ -152,6 +152,10 @@ struct Args {
     // Maximum number of subtasks that can be generated for tasks
     #[clap(long, default_value = DEFAULT_MAX_SUBTASKS_PER_TASK)]
     max_subtasks_per_task: usize,
+
+    /// YAML file mapping mapping bazel target names of suite configs to their file location location
+    #[clap(long, value_parser)]
+    bazel_suite_configs: Option<PathBuf>,
 }
 
 /// Configure logging for the command execution.
@@ -194,6 +198,7 @@ async fn main() {
             default_subtasks_per_task: args.default_subtasks_per_task,
             large_required_task_runtime_threshold: args.large_required_task_runtime_threshold,
         },
+        bazel_suite_configs: args.bazel_suite_configs.as_ref().map(|p| expand_path(p)),
     };
     let s3_client = build_s3_client().await;
     let deps = Dependencies::new(execution_config, s3_client).unwrap();

src/main.rs

@@ -4,7 +4,10 @@ use anyhow::Result;
 use serde::Deserialize;
 use tracing::{error, event, Level};
 
+use crate::evergreen::evg_config_utils::is_bazel_suite;
+
 use super::{external_cmd::run_command, resmoke_suite::ResmokeSuiteConfig};
+use std::collections::HashMap;
 
 /// Interface for discovering details about test suites.
 pub trait TestDiscovery: Send + Sync {
@@ -45,6 +48,7 @@ pub struct ResmokeProxy {
     skip_covered_tests: bool,
     /// True if test discovery should include tests that are tagged with fully disabled features.
     include_fully_disabled_feature_tests: bool,
+    bazel_suite_configs: BazelConfigs,
 }
 
 impl ResmokeProxy {
@@ -55,10 +59,12 @@ impl ResmokeProxy {
     /// * `resmoke_cmd` - Command to invoke resmoke.
     /// * `skip_covered_tests` - Whether the generator should skip tests run in more complex suites.
     /// * `include_fully_disabled_feature_tests` - If the generator should include tests that are tagged with fully disabled features.
+    /// * `bazel_suite_configs` - Optional bazel suite configurations.
     pub fn new(
         resmoke_cmd: &str,
         skip_covered_tests: bool,
         include_fully_disabled_feature_tests: bool,
+        bazel_suite_configs: BazelConfigs,
     ) -> Self {
         let cmd_parts: Vec<_> = resmoke_cmd.split(' ').collect();
         let cmd = cmd_parts[0];
@@ -68,10 +74,46 @@ impl ResmokeProxy {
             resmoke_script: script,
             skip_covered_tests,
             include_fully_disabled_feature_tests,
+            bazel_suite_configs,
         }
     }
 }
 
+#[derive(Debug, Clone, Default)]
+pub struct BazelConfigs {
+    /// Map of bazel resmoke config targets to their generated suite config YAMLs.
+    configs: HashMap<String, String>,
+}
+
+impl BazelConfigs {
+    pub fn from_yaml_file(path: &Path) -> Result<Self> {
+        let contents = std::fs::read_to_string(path)?;
+        let configs: Result<HashMap<String, String>, serde_yaml::Error> =
+            serde_yaml::from_str(&contents);
+        if configs.is_err() {
+            error!(
+                file = path.display().to_string(),
+                contents = &contents,
+                "Failed to parse bazel configs from yaml file",
+            );
+        }
+        Ok(Self { configs: configs? })
+    }
+
+    /// Get the generated suite config for a bazel resmoke target.
+    ///
+    /// # Arguments
+    ///
+    /// * `target` - Bazel resmoke test target, like "//buildscripts/resmoke:core".
+    ///
+    /// # Returns
+    ///
+    /// The path the the generated suite config YAML, like "bazel-out/buildscripts/resmoke/core_config.yml".
+    pub fn get(&self, target: &str) -> &str {
+        self.configs.get(&format!("{}_config", target)).unwrap()
+    }
+}
+
 /// Details about tests comprising a test suite.
 #[derive(Debug, Deserialize)]
 #[allow(dead_code)]
@@ -94,9 +136,15 @@ impl TestDiscovery for ResmokeProxy {
     ///
     /// A list of tests belonging to given suite.
     fn discover_tests(&self, suite_name: &str) -> Result<Vec<String>> {
+        let suite_config = if is_bazel_suite(suite_name) {
+            self.bazel_suite_configs.get(suite_name)
+        } else {
+            suite_name
+        };
+
         let mut cmd = vec![&*self.resmoke_cmd];
         cmd.append(&mut self.resmoke_script.iter().map(|s| s.as_str()).collect());
-        cmd.append(&mut vec!["test-discovery", "--suite", suite_name]);
+        cmd.append(&mut vec!["test-discovery", "--suite", suite_config]);
 
         // When running in a patch build, we use the --skipTestsCoveredByMoreComplexSuites
         // flag to tell Resmoke to exclude any tests in the given suite that will
@@ -114,7 +162,7 @@ impl TestDiscovery for ResmokeProxy {
 
         event!(
             Level::INFO,
-            suite_name,
+            suite_config,
             duration_ms = start.elapsed().as_millis() as u64,
             "Resmoke test discovery finished"
         );
@@ -146,9 +194,15 @@ impl TestDiscovery for ResmokeProxy {
     ///
     /// Resmoke configuration for the given suite.
     fn get_suite_config(&self, suite_name: &str) -> Result<ResmokeSuiteConfig> {
+        let suite_config = if is_bazel_suite(suite_name) {
+            self.bazel_suite_configs.get(suite_name)
+        } else {
+            suite_name
+        };
+
         let mut cmd = vec![&*self.resmoke_cmd];
         cmd.append(&mut self.resmoke_script.iter().map(|s| s.as_str()).collect());
-        cmd.append(&mut vec!["suiteconfig", "--suite", suite_name]);
+        cmd.append(&mut vec!["suiteconfig", "--suite", suite_config]);
         let cmd_output = run_command(&cmd).unwrap();
 
         Ok(ResmokeSuiteConfig::from_str(&cmd_output)?)