pgdo/coordinate/
cleanup.rs

1use std::panic::{catch_unwind, resume_unwind};
2
3/// Perform `task`. On error or panic, perform `cleanup` and return the original
4/// error, or continue to panic.
5///
6/// This is resilient to errors and panics in `cleanup` too: they will be
7/// logged, but ultimately the errors and panics from `task` will be propagated.
8pub fn with_cleanup<TASK, T, E, CLEANUP, CT, CE>(cleanup: CLEANUP, task: TASK) -> Result<T, E>
9where
10    TASK: std::panic::UnwindSafe + FnOnce() -> Result<T, E>,
11    CLEANUP: std::panic::UnwindSafe + FnOnce() -> Result<CT, CE>,
12    E: std::fmt::Display,
13    CE: std::fmt::Display + Into<E>,
14{
15    match catch_unwind(task) {
16        Ok(Ok(t)) => Ok(t),
17        Ok(Err(e)) => match catch_unwind(cleanup) {
18            Ok(Ok(_)) => Err(e),
19            Ok(Err(ce)) => {
20                log::error!("Task failed & cleaning-up also failed: {ce}");
21                Err(e)
22            }
23            Err(_) => {
24                log::error!("Task failed & cleaning-up panicked (suppressed)");
25                Err(e)
26            }
27        },
28        Err(panic) => match catch_unwind(cleanup) {
29            Ok(Ok(_)) => resume_unwind(panic),
30            Ok(Err(ce)) => {
31                log::error!("Task panicked & cleaning-up failed: {ce}");
32                resume_unwind(panic)
33            }
34            Err(_) => {
35                log::error!("Task panicked & cleaning-up also panicked (suppressed)");
36                resume_unwind(panic)
37            }
38        },
39    }
40}
41
42#[cfg(test)]
43mod tests {
44    use super::with_cleanup;
45
46    #[test]
47    fn test_with_cleanup() {
48        let result: Result<&'static str, &'static str> = with_cleanup(
49            || Ok::<&'static str, &'static str>("Ok/cleanup"),
50            || Ok("Ok/task"),
51        );
52        assert!(matches!(result, Ok("Ok/task")));
53    }
54
55    #[test]
56    fn test_with_cleanup_error_in_task() {
57        let result: Result<(), &'static str> =
58            with_cleanup(|| Ok::<(), &'static str>(()), || Err("Err/task")?);
59        assert!(matches!(result, Err("Err/task")));
60    }
61
62    #[test]
63    #[should_panic(expected = "Panic/task")]
64    fn test_with_cleanup_panic_in_task() {
65        let _result: Result<(), &'static str> =
66            with_cleanup(|| Ok::<(), &'static str>(()), || panic!("Panic/task"));
67    }
68
69    #[test]
70    fn test_with_cleanup_error_in_cleanup() {
71        let result: Result<(), &'static str> =
72            with_cleanup(|| Err::<(), &'static str>("Err/cleanup"), || Ok(()));
73        assert!(matches!(result, Ok(())));
74    }
75
76    #[test]
77    fn test_with_cleanup_panic_in_cleanup() {
78        let result: Result<(), &'static str> = with_cleanup(
79            || -> Result<(), &'static str> { panic!("Panic/cleanup") },
80            || Ok(()),
81        );
82        assert!(matches!(result, Ok(())));
83    }
84
85    #[test]
86    fn test_with_cleanup_error_in_task_and_cleanup() {
87        let result: Result<(), &'static str> = with_cleanup(
88            || Err::<(), &'static str>("Err/cleanup"),
89            || Err("Err/task")?,
90        );
91        assert!(matches!(result, Err("Err/task")));
92    }
93
94    #[test]
95    #[should_panic(expected = "Panic/task")]
96    fn test_with_cleanup_panic_in_task_and_cleanup() {
97        let _result: Result<(), &'static str> = with_cleanup(
98            || -> Result<(), &'static str> { panic!("Panic/cleanup") },
99            || panic!("Panic/task"),
100        );
101    }
102}
pgdo/coordinate/cleanup.rs

pgdo/coordinate/
cleanup.rs
