Setup steps API: Check for cancellation before running a step

elabit · Nov 20, 2024 · 1187c18 · 1187c18
1 parent 1d5a1c5
commit 1187c18
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 11 deletions.
diff --git 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
@@ -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<Plan>) -> (Box<dyn SetupStep>, Vec<Plan>) {
     (Box::new(SetupStepSuccess {}), plans)
 }
 
-pub fn run_steps(steps: Vec<StepWithPlans>) -> (Vec<Plan>, Vec<SetupFailure>) {
+pub fn run_steps(
+    steps: Vec<StepWithPlans>,
+    cancellation_token: &CancellationToken,
+) -> Result<(Vec<Plan>, Vec<SetupFailure>), 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<StepWithPlans>) -> (Vec<Plan>, Vec<SetupFailure>) {
             }
         }
     }
-    (plans, errors)
+    Ok((plans, errors))
 }
diff --git 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<Plan>) -> (Vec<Plan>, Vec<SetupFailure>) {
+fn run_setup_steps(
+    config: &GlobalConfig,
+    mut plans: Vec<Plan>,
+) -> Result<(Vec<Plan>, Vec<SetupFailure>), Cancelled> {
     let gather_requirements = [
         gather_managed_directories,
         #[cfg(windows)]
@@ -107,13 +111,14 @@ fn run_setup_steps(config: &GlobalConfig, mut plans: Vec<Plan>) -> (Vec<Plan>, V
                 plan_count,
                 setup_steps.iter().map(|s| s.1.len()).sum::<usize>()
             );
-            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
@@ -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,15 +19,21 @@ use robotmk::config::RCCProfileConfig;
 use std::vec;
 use tokio_util::sync::CancellationToken;
 
-pub fn setup(global_config: &GlobalConfig, plans: Vec<Plan>) -> (Vec<Plan>, Vec<SetupFailure>) {
+pub fn setup(
+    global_config: &GlobalConfig,
+    plans: Vec<Plan>,
+) -> Result<(Vec<Plan>, Vec<SetupFailure>), Cancelled> {
     run_setup_steps(global_config, plans)
 }
 
 pub fn rcc_setup_working_directory(working_directory: &Utf8Path) -> Utf8PathBuf {
     working_directory.join("rcc_setup")
 }
 
-fn run_setup_steps(config: &GlobalConfig, mut plans: Vec<Plan>) -> (Vec<Plan>, Vec<SetupFailure>) {
+fn run_setup_steps(
+    config: &GlobalConfig,
+    mut plans: Vec<Plan>,
+) -> Result<(Vec<Plan>, Vec<SetupFailure>), Cancelled> {
     let gather_requirements = [
         #[cfg(windows)]
         gather_rcc_binary_permissions,
@@ -50,13 +57,14 @@ fn run_setup_steps(config: &GlobalConfig, mut plans: Vec<Plan>) -> (Vec<Plan>, V
                 plan_count,
                 setup_steps.iter().map(|s| s.1.len()).sum::<usize>()
             );
-            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)]