From b2d71d8affd4482857381b683458e9928164c502 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B6rg=20Herbel?= <joerg.herbel@checkmk.com>
Date: Wed, 20 Nov 2024 19:36:19 +0100
Subject: [PATCH] Use `Terminate` instead of `anyhow::Error` in more places

`Terminate` is more expressive and than `anyhow::Error`, since it explicitly
takes cancellation into account.
---
 src/bin/scheduler/logging.rs           |  7 +++++--
 src/bin/scheduler/scheduling/plans.rs  | 12 +++++++-----
 src/bin/scheduler/setup/directories.rs |  7 ++++---
 src/termination.rs                     | 19 +++++++++++++++++++
 4 files changed, 35 insertions(+), 10 deletions(-)

diff --git a/src/bin/scheduler/logging.rs b/src/bin/scheduler/logging.rs
index 145c4afc..815b906e 100644
--- a/src/bin/scheduler/logging.rs
+++ b/src/bin/scheduler/logging.rs
@@ -1,10 +1,10 @@
-use anyhow::Error;
 use camino::Utf8PathBuf;
 use flexi_logger::{
     Age, Cleanup, Criterion, DeferredNow, FileSpec, FlexiLoggerError, LogSpecification, Logger,
     LoggerHandle, Naming, Record,
 };
 use log::error;
+use std::fmt::Debug;
 
 pub const TIMESTAMP_FORMAT: &str = "%Y-%m-%dT%H.%M.%S%.f%z";
 
@@ -43,7 +43,10 @@ pub fn init(
     .start()
 }
 
-pub fn log_and_return_error(error: Error) -> Error {
+pub fn log_and_return_error<T>(error: T) -> T
+where
+    T: Debug,
+{
     error!("{error:?}");
     error
 }
diff --git a/src/bin/scheduler/scheduling/plans.rs b/src/bin/scheduler/scheduling/plans.rs
index 5d747b07..b8f57198 100644
--- a/src/bin/scheduler/scheduling/plans.rs
+++ b/src/bin/scheduler/scheduling/plans.rs
@@ -3,13 +3,14 @@ use crate::logging::TIMESTAMP_FORMAT;
 use robotmk::plans::run_attempts_with_rebot;
 use robotmk::results::{AttemptsConfig, PlanExecutionReport};
 
-use anyhow::{Context, Result as AnyhowResult};
+use anyhow::Context;
 use chrono::Utc;
 use log::info;
 use robotmk::section::WritePiggybackSection;
+use robotmk::termination::{ContextUnrecoverable, Terminate};
 use std::fs::create_dir_all;
 
-pub fn run_plan(plan: &Plan) -> AnyhowResult<()> {
+pub fn run_plan(plan: &Plan) -> Result<(), Terminate> {
     info!(
         "Running plan {} ({})",
         &plan.id,
@@ -21,7 +22,7 @@ pub fn run_plan(plan: &Plan) -> AnyhowResult<()> {
             plan.host.clone(),
             &plan.results_directory_locker,
         )
-        .context("Reporting plan results failed")?;
+        .context_unrecoverable("Reporting plan results failed")?;
     info!("Plan {} finished", &plan.id);
 
     Ok(())
@@ -48,7 +49,7 @@ fn format_source_for_logging(source: &Source) -> String {
     }
 }
 
-fn produce_plan_results(plan: &Plan) -> AnyhowResult<PlanExecutionReport> {
+fn produce_plan_results(plan: &Plan) -> Result<PlanExecutionReport, Terminate> {
     let timestamp = Utc::now();
     let output_directory = plan
         .working_directory
@@ -68,7 +69,8 @@ fn produce_plan_results(plan: &Plan) -> AnyhowResult<PlanExecutionReport> {
         &plan.cancellation_token,
         &output_directory,
     )
-    .context("Received termination signal while running plan")?;
+    .map_err(|cancelled| cancelled.into())
+    .context_unrecoverable("Received termination signal while running plan")?;
 
     Ok(PlanExecutionReport {
         plan_id: plan.id.clone(),
diff --git a/src/bin/scheduler/setup/directories.rs b/src/bin/scheduler/setup/directories.rs
index 4f9c33a9..410ec85e 100644
--- a/src/bin/scheduler/setup/directories.rs
+++ b/src/bin/scheduler/setup/directories.rs
@@ -14,7 +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::Terminate;
+use robotmk::termination::{ContextUnrecoverable, Terminate};
 use std::collections::HashSet;
 
 pub fn setup(
@@ -60,10 +60,11 @@ fn setup_managed_directory(managed_directory: &Utf8Path) -> AnyhowResult<()> {
     Ok(())
 }
 
-fn setup_results_directory(global_config: &GlobalConfig, plans: &[Plan]) -> AnyhowResult<()> {
+fn setup_results_directory(global_config: &GlobalConfig, plans: &[Plan]) -> Result<(), Terminate> {
     create_dir_all(&global_config.results_directory)?;
     create_dir_all(plan_results_directory(&global_config.results_directory))?;
-    clean_up_results_directory(global_config, plans).context("Failed to clean up results directory")
+    clean_up_results_directory(global_config, plans)
+        .context_unrecoverable("Failed to clean up results directory")
 }
 
 fn clean_up_results_directory(
diff --git a/src/termination.rs b/src/termination.rs
index f38739d0..5213c851 100644
--- a/src/termination.rs
+++ b/src/termination.rs
@@ -1,4 +1,5 @@
 use std::collections::HashSet;
+use std::fmt::Display;
 use std::future::Future;
 use std::time::Duration;
 use sysinfo::{Pid, Process, ProcessesToUpdate, System};
@@ -14,6 +15,24 @@ pub enum Terminate {
     Unrecoverable(#[from] anyhow::Error),
 }
 
+pub trait ContextUnrecoverable<T> {
+    fn context_unrecoverable<C>(self, context: C) -> Result<T, Terminate>
+    where
+        C: Display + Send + Sync + 'static;
+}
+
+impl<T> ContextUnrecoverable<T> for Result<T, Terminate> {
+    fn context_unrecoverable<C>(self, context: C) -> Result<T, Terminate>
+    where
+        C: Display + Send + Sync + 'static,
+    {
+        self.map_err(|err| match err {
+            Terminate::Unrecoverable(any) => Terminate::Unrecoverable(any.context(context)),
+            Terminate::Cancelled => Terminate::Cancelled,
+        })
+    }
+}
+
 #[derive(Error, Debug, Clone)]
 #[error("Terminated")]
 pub struct Cancelled;