diff --git a/bindings/trycp_runner/src/cli.rs b/bindings/trycp_runner/src/cli.rs index 6cd2b4d0..f6cc656a 100644 --- a/bindings/trycp_runner/src/cli.rs +++ b/bindings/trycp_runner/src/cli.rs @@ -25,6 +25,12 @@ pub struct WindTunnelTryCPScenarioCli { #[clap(long, default_value = "1")] pub instances_per_target: u8, + /// The minimum number of agents required for the scenario to run + /// + /// If the number of running agents drops below this value the scenario will fail. + #[clap(long)] + pub min_required_agents: Option, + /// Assign a behaviour to a number of agents. Specify the behaviour and number of agents to assign /// it to in the format `behaviour:count`. For example `--behaviour=login:5`. /// @@ -89,6 +95,7 @@ impl TryInto for WindTunnelTryCPScenarioCli { // Pack values together and extract by agent id in helpers. connection_string: targets.nodes.join(","), agents: Some(required_agents), + min_required_agents: self.min_required_agents, behaviour: self.behaviour, duration: self.duration, soak: self.soak, diff --git a/framework/runner/src/cli.rs b/framework/runner/src/cli.rs index b0bcc1a5..881d274f 100644 --- a/framework/runner/src/cli.rs +++ b/framework/runner/src/cli.rs @@ -11,6 +11,12 @@ pub struct WindTunnelScenarioCli { #[clap(long)] pub agents: Option, + /// The minimum number of agents required for the scenario to run + /// + /// If the number of running agents drops below this value the scenario will fail. + #[clap(long)] + pub min_required_agents: Option, + /// Assign a behaviour to a number of agents. Specify the behaviour and number of agents to assign /// it to in the format `behaviour:count`. For example `--behaviour=login:5`. /// diff --git a/framework/runner/src/definition.rs b/framework/runner/src/definition.rs index 71052622..7766d5fa 100644 --- a/framework/runner/src/definition.rs +++ b/framework/runner/src/definition.rs @@ -24,6 +24,7 @@ pub struct ScenarioDefinitionBuilder, default_duration_s: Option, + default_min_required_agents: Option, setup_fn: Option>, setup_agent_fn: Option>, agent_behaviour: HashMap>, @@ -41,6 +42,7 @@ pub struct ScenarioDefinition pub(crate) name: String, pub(crate) assigned_behaviours: Vec, pub(crate) duration_s: Option, + pub(crate) min_required_agents: usize, pub(crate) connection_string: String, pub(crate) no_progress: bool, pub(crate) reporter: ReporterOpt, @@ -92,6 +94,7 @@ impl ScenarioDefinitionBuilde cli, default_agent_count: None, default_duration_s: None, + default_min_required_agents: None, setup_fn: None, setup_agent_fn: None, agent_behaviour: HashMap::new(), @@ -116,6 +119,15 @@ impl ScenarioDefinitionBuilde self } + /// Sets the minimum number of agents required for this scenario. + /// + /// The scenario will fail if the number of running agents drops below this amount. + /// This can be overridden when the scenario is run using the `--min-required-agents` flag. + pub fn with_default_min_required_agents(mut self, default_min_required_agents: usize) -> Self { + self.default_min_required_agents = Some(default_min_required_agents); + self + } + /// Sets the global setup hook for this scenario. It will be run once, before any agents are started. pub fn use_setup(mut self, setup_fn: GlobalHookMut) -> Self { self.setup_fn = Some(setup_fn); @@ -178,6 +190,13 @@ impl ScenarioDefinitionBuilde // Priority given to the CLI, then the default value provided by the scenario, then default to 1 let resolved_agent_count = self.cli.agents.or(self.default_agent_count).unwrap_or(1); + // Priority given to the CLI, then the default value provided by the scenario, then default to 1 + let min_required_agents = self + .cli + .min_required_agents + .or(self.default_min_required_agents) + .unwrap_or(1); + // Check that the user hasn't requested behaviours that aren't registered in the scenario. let registered_behaviours = self .agent_behaviour @@ -204,6 +223,7 @@ impl ScenarioDefinitionBuilde name: self.name, assigned_behaviours: build_assigned_behaviours(&self.cli, resolved_agent_count)?, duration_s: resolved_duration, + min_required_agents, connection_string: self.cli.connection_string, no_progress: self.cli.no_progress, reporter: self.cli.reporter, @@ -255,6 +275,7 @@ mod tests { &crate::cli::WindTunnelScenarioCli { connection_string: "".to_string(), agents: None, + min_required_agents: None, behaviour: vec![], duration: None, soak: false, @@ -276,6 +297,7 @@ mod tests { &crate::cli::WindTunnelScenarioCli { connection_string: "".to_string(), agents: None, + min_required_agents: None, behaviour: vec![], // Not specified duration: None, soak: false, @@ -297,6 +319,7 @@ mod tests { &crate::cli::WindTunnelScenarioCli { connection_string: "".to_string(), agents: None, + min_required_agents: None, behaviour: vec![("login".to_string(), 3)], // 3 of 5 duration: None, soak: false, @@ -320,6 +343,7 @@ mod tests { &crate::cli::WindTunnelScenarioCli { connection_string: "".to_string(), agents: None, + min_required_agents: None, behaviour: vec![("login".to_string(), 30)], // 30 of 5 duration: None, soak: false, diff --git a/framework/runner/src/run.rs b/framework/runner/src/run.rs index 0f180da1..31e78239 100644 --- a/framework/runner/src/run.rs +++ b/framework/runner/src/run.rs @@ -3,7 +3,7 @@ use std::sync::atomic::AtomicUsize; use std::sync::Arc; use std::time::Duration; -use anyhow::Context; +use anyhow::{anyhow, Context}; use wind_tunnel_core::prelude::{AgentBailError, ShutdownHandle, ShutdownSignalError}; use wind_tunnel_instruments::ReportConfig; @@ -216,5 +216,12 @@ pub fn run( println!("#RunId: [{}]", run_id); - Ok(agents_run_to_completion.load(std::sync::atomic::Ordering::Acquire)) + let agents_run_to_completion = + agents_run_to_completion.load(std::sync::atomic::Ordering::Acquire); + + if agents_run_to_completion < definition.min_required_agents { + Err(anyhow!("Not enough agents ran scenario to completion: expected {}, actual {agents_run_to_completion}", definition.min_required_agents)) + } else { + Ok(agents_run_to_completion) + } } diff --git a/framework/runner/tests/hook_error_handling.rs b/framework/runner/tests/hook_error_handling.rs index 864f38b2..ab5bf145 100644 --- a/framework/runner/tests/hook_error_handling.rs +++ b/framework/runner/tests/hook_error_handling.rs @@ -21,6 +21,7 @@ fn sample_cli_cfg() -> WindTunnelScenarioCli { WindTunnelScenarioCli { connection_string: "test_connection_string".to_string(), agents: None, + min_required_agents: None, behaviour: vec![], duration: None, soak: false, @@ -59,6 +60,7 @@ fn capture_error_in_agent_setup() { sample_cli_cfg(), ) .with_default_duration_s(5) + .with_default_min_required_agents(0) .use_agent_setup(agent_setup); let result = run(scenario); @@ -161,3 +163,57 @@ fn capture_error_in_teardown() { assert!(result.is_ok()); } + +#[test] +fn error_if_default_minimum_required_agents_not_met() { + fn agent_behaviour(_: &mut AgentContext) -> HookResult { + Err(AgentBailError::default().into()) + } + + let scenario = ScenarioDefinitionBuilder::::new( + "error_if_minimum_required_agents_not_met", + sample_cli_cfg(), + ) + .with_default_duration_s(5) + .use_agent_behaviour(agent_behaviour); + + let result = run(scenario); + + assert!(result.is_err()); + assert_eq!( + result.unwrap_err().to_string(), + "Not enough agents ran scenario to completion: expected 1, actual 0" + ); +} + +#[test] +fn error_if_set_minimum_required_agents_not_met() { + fn agent_behaviour( + ctx: &mut AgentContext, + ) -> HookResult { + if ctx.agent_index() == 0 { + Ok(()) + } else { + Err(AgentBailError::default().into()) + } + } + + let scenario = ScenarioDefinitionBuilder::::new( + "error_if_minimum_required_agents_not_met", + WindTunnelScenarioCli { + agents: Some(2), + ..sample_cli_cfg() + }, + ) + .with_default_duration_s(5) + .with_default_min_required_agents(2) + .use_agent_behaviour(agent_behaviour); + + let result = run(scenario); + + assert!(result.is_err()); + assert_eq!( + result.unwrap_err().to_string(), + "Not enough agents ran scenario to completion: expected 2, actual 1" + ); +}