diff --git a/src/bin/scheduler/main.rs b/src/bin/scheduler/main.rs index ea04323c..17b6c524 100644 --- a/src/bin/scheduler/main.rs +++ b/src/bin/scheduler/main.rs @@ -75,7 +75,7 @@ fn run() -> Result<(), Terminate> { write_phase(&SchedulerPhase::RCCSetup, &global_config)?; info!("RCC-specific setup started"); - let (plans, rcc_setup_failures) = setup::rcc::setup(&global_config, plans); + let (plans, rcc_setup_failures) = setup::rcc::setup(&global_config, plans)?; write_setup_failures( general_setup_failures .into_iter() diff --git a/src/bin/scheduler/setup/api.rs b/src/bin/scheduler/setup/api.rs index 16b626b9..49757f77 100644 --- a/src/bin/scheduler/setup/api.rs +++ b/src/bin/scheduler/setup/api.rs @@ -1,5 +1,7 @@ use crate::internal_config::Plan; use robotmk::results::SetupFailure; +use robotmk::termination::Cancelled; +use tokio_util::sync::CancellationToken; #[derive(Debug)] pub struct Error { @@ -31,10 +33,16 @@ pub fn skip(plans: Vec) -> (Box, Vec) { (Box::new(SetupStepSuccess {}), plans) } -pub fn run_steps(steps: Vec) -> (Vec, Vec) { +pub fn run_steps( + steps: Vec, + cancellation_token: &CancellationToken, +) -> Result<(Vec, Vec), Cancelled> { let mut plans = Vec::new(); let mut errors = Vec::new(); for (step, affected_plans) in steps.into_iter() { + if cancellation_token.is_cancelled() { + return Err(Cancelled); + } match step.setup() { Ok(()) => { plans.extend(affected_plans); @@ -56,5 +64,5 @@ pub fn run_steps(steps: Vec) -> (Vec, Vec) { } } } - (plans, errors) + Ok((plans, errors)) } diff --git a/src/bin/scheduler/setup/directories.rs b/src/bin/scheduler/setup/directories.rs index 4f9c33a9..bfa23cdb 100644 --- a/src/bin/scheduler/setup/directories.rs +++ b/src/bin/scheduler/setup/directories.rs @@ -14,6 +14,7 @@ use robotmk::environment::Environment; use robotmk::fs::{create_dir_all, remove_dir_all, remove_file}; use robotmk::results::{plan_results_directory, SetupFailure}; use robotmk::session::Session; +use robotmk::termination::Cancelled; use robotmk::termination::Terminate; use std::collections::HashSet; @@ -30,7 +31,7 @@ pub fn setup( setup_managed_directory(&global_config.managed_directory)?; setup_results_directory(global_config, &plans)?; - Ok(run_setup_steps(global_config, plans)) + Ok(run_setup_steps(global_config, plans)?) } fn setup_working_directory(working_directory: &Utf8Path, plans: &[Plan]) -> AnyhowResult<()> { @@ -83,7 +84,10 @@ fn clean_up_results_directory( Ok(results_directory_lock.release()?) } -fn run_setup_steps(config: &GlobalConfig, mut plans: Vec) -> (Vec, Vec) { +fn run_setup_steps( + config: &GlobalConfig, + mut plans: Vec, +) -> Result<(Vec, Vec), Cancelled> { let gather_requirements = [ gather_managed_directories, #[cfg(windows)] @@ -107,13 +111,14 @@ fn run_setup_steps(config: &GlobalConfig, mut plans: Vec) -> (Vec, V plan_count, setup_steps.iter().map(|s| s.1.len()).sum::() ); - let (surviving_plans, current_errors) = run_steps(setup_steps); + let (surviving_plans, current_errors) = + run_steps(setup_steps, &config.cancellation_token)?; failures.extend(current_errors); surviving_plans }; } sort_plans_by_grouping(&mut plans); - (plans, failures) + Ok((plans, failures)) } struct StepCreate { diff --git a/src/bin/scheduler/setup/rcc.rs b/src/bin/scheduler/setup/rcc.rs index 3de41bef..d57648d2 100644 --- a/src/bin/scheduler/setup/rcc.rs +++ b/src/bin/scheduler/setup/rcc.rs @@ -9,6 +9,7 @@ use robotmk::results::SetupFailure; #[cfg(windows)] use robotmk::session::CurrentSession; use robotmk::session::{RunSpec, Session}; +use robotmk::termination::Cancelled; use robotmk::termination::Outcome; use anyhow::{anyhow, Context}; @@ -18,7 +19,10 @@ use robotmk::config::RCCProfileConfig; use std::vec; use tokio_util::sync::CancellationToken; -pub fn setup(global_config: &GlobalConfig, plans: Vec) -> (Vec, Vec) { +pub fn setup( + global_config: &GlobalConfig, + plans: Vec, +) -> Result<(Vec, Vec), Cancelled> { run_setup_steps(global_config, plans) } @@ -26,7 +30,10 @@ pub fn rcc_setup_working_directory(working_directory: &Utf8Path) -> Utf8PathBuf working_directory.join("rcc_setup") } -fn run_setup_steps(config: &GlobalConfig, mut plans: Vec) -> (Vec, Vec) { +fn run_setup_steps( + config: &GlobalConfig, + mut plans: Vec, +) -> Result<(Vec, Vec), Cancelled> { let gather_requirements = [ #[cfg(windows)] gather_rcc_binary_permissions, @@ -50,13 +57,14 @@ fn run_setup_steps(config: &GlobalConfig, mut plans: Vec) -> (Vec, V plan_count, setup_steps.iter().map(|s| s.1.len()).sum::() ); - let (surviving_plans, current_errors) = run_steps(setup_steps); + let (surviving_plans, current_errors) = + run_steps(setup_steps, &config.cancellation_token)?; failures.extend(current_errors); surviving_plans }; } sort_plans_by_grouping(&mut plans); - (plans, failures) + Ok((plans, failures)) } #[cfg(windows)]