pgdo/coordinate/
finally.rs

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