Skip to main content

lsp_max/service/
cancellation.rs

1//! Utilities for ensuring cancellation safety and resource cleanup.
2
3/// A guard that executes a closure when dropped, unless disarmed.
4///
5/// This is useful for ensuring that resources are cleaned up even if a future
6/// is cancelled (dropped) before it completes.
7///
8/// See the [cancellation safety documentation](../../docs/CANCELLATION_SAFETY.md) for more details.
9#[derive(Debug)]
10pub struct CancellationGuard<F: FnOnce()> {
11    on_drop: Option<F>,
12}
13
14impl<F: FnOnce()> CancellationGuard<F> {
15    /// Creates a new guard that will execute the given closure when dropped.
16    pub fn new(on_drop: F) -> Self {
17        Self {
18            on_drop: Some(on_drop),
19        }
20    }
21
22    /// Disarms the guard, preventing the closure from being executed when dropped.
23    pub fn disarm(mut self) {
24        self.on_drop = None;
25    }
26}
27
28impl<F: FnOnce()> Drop for CancellationGuard<F> {
29    fn drop(&mut self) {
30        if let Some(on_drop) = self.on_drop.take() {
31            on_drop();
32        }
33    }
34}
35
36/// A trait for types that can be wrapped in a [`CancellationGuard`].
37pub trait OnDrop: Sized {
38    /// Wraps this value in a [`CancellationGuard`] that will call the given
39    /// closure when dropped.
40    fn on_drop<F: FnOnce()>(self, on_drop: F) -> (Self, CancellationGuard<F>) {
41        (self, CancellationGuard::new(on_drop))
42    }
43}
44
45impl<T> OnDrop for T {}
46
47#[cfg(test)]
48mod tests {
49    use std::sync::atomic::{AtomicBool, Ordering};
50    use std::sync::Arc;
51
52    use super::*;
53
54    #[test]
55    fn executes_on_drop() {
56        let executed = Arc::new(AtomicBool::new(false));
57        let executed_clone = executed.clone();
58
59        {
60            let _guard = CancellationGuard::new(move || {
61                executed_clone.store(true, Ordering::SeqCst);
62            });
63        }
64
65        assert!(executed.load(Ordering::SeqCst));
66    }
67
68    #[test]
69    fn does_not_execute_when_disarmed() {
70        let executed = Arc::new(AtomicBool::new(false));
71        let executed_clone = executed.clone();
72
73        {
74            let guard = CancellationGuard::new(move || {
75                executed_clone.store(true, Ordering::SeqCst);
76            });
77            guard.disarm();
78        }
79
80        assert!(!executed.load(Ordering::SeqCst));
81    }
82
83    #[test]
84    fn on_drop_helper() {
85        let executed = Arc::new(AtomicBool::new(false));
86        let executed_clone = executed.clone();
87
88        {
89            let (_val, _guard) = "resource".on_drop(move || {
90                executed_clone.store(true, Ordering::SeqCst);
91            });
92        }
93
94        assert!(executed.load(Ordering::SeqCst));
95    }
96}