tower_http/on_early_drop/
guard.rs

1//! Drop guard that fires a callback when dropped unless marked completed.
2
3use crate::on_early_drop::traits::OnDropCallback;
4
5/// Runs a callback on drop unless [`completed`](Self::completed) is called
6/// first.
7///
8/// # Examples
9///
10/// ```
11/// use tower_http::on_early_drop::OnEarlyDropGuard;
12/// use std::sync::atomic::{AtomicUsize, Ordering};
13/// use std::sync::Arc;
14///
15/// let count = Arc::new(AtomicUsize::new(0));
16/// let count_for_guard = count.clone();
17/// {
18///     let _guard = OnEarlyDropGuard::new(move || {
19///         count_for_guard.fetch_add(1, Ordering::Relaxed);
20///     });
21/// }
22/// assert_eq!(count.load(Ordering::Relaxed), 1);
23/// ```
24///
25/// Marking the guard completed suppresses the callback:
26///
27/// ```
28/// use tower_http::on_early_drop::OnEarlyDropGuard;
29/// use std::sync::atomic::{AtomicUsize, Ordering};
30/// use std::sync::Arc;
31///
32/// let count = Arc::new(AtomicUsize::new(0));
33/// let count_for_guard = count.clone();
34/// {
35///     let mut guard = OnEarlyDropGuard::new(move || {
36///         count_for_guard.fetch_add(1, Ordering::Relaxed);
37///     });
38///     guard.completed();
39/// }
40/// assert_eq!(count.load(Ordering::Relaxed), 0);
41/// ```
42///
43/// [`OnEarlyDropLayer`]: super::OnEarlyDropLayer
44#[derive(Debug)]
45pub struct OnEarlyDropGuard<Callback: OnDropCallback> {
46    callback: Option<Callback>,
47}
48
49impl<Callback: OnDropCallback> OnEarlyDropGuard<Callback> {
50    /// Create a guard that will fire `callback` on drop.
51    pub fn new(callback: Callback) -> Self {
52        Self {
53            callback: Some(callback),
54        }
55    }
56
57    /// Mark the guard completed and drop the callback without firing it.
58    ///
59    /// Any resources captured by the callback are released immediately
60    /// rather than at guard drop time.
61    pub fn completed(&mut self) {
62        self.callback = None;
63    }
64}
65
66impl<Callback: OnDropCallback> Drop for OnEarlyDropGuard<Callback> {
67    fn drop(&mut self) {
68        // Panicking in Drop aborts the process if we are already unwinding,
69        // so avoid `expect` here.
70        if let Some(callback) = self.callback.take() {
71            callback.on_drop();
72        }
73    }
74}
75
76#[cfg(test)]
77mod tests {
78    use super::*;
79    use std::sync::atomic::{AtomicBool, Ordering};
80    use std::sync::Arc;
81
82    #[test]
83    fn fires_on_drop() {
84        let fired = Arc::new(AtomicBool::new(false));
85        let fired_for_guard = fired.clone();
86        {
87            let _guard = OnEarlyDropGuard::new(move || {
88                fired_for_guard.store(true, Ordering::Relaxed);
89            });
90        }
91        assert!(fired.load(Ordering::Relaxed));
92    }
93
94    #[test]
95    fn suppresses_when_completed() {
96        let fired = Arc::new(AtomicBool::new(false));
97        let fired_for_guard = fired.clone();
98        {
99            let mut guard = OnEarlyDropGuard::new(move || {
100                fired_for_guard.store(true, Ordering::Relaxed);
101            });
102            guard.completed();
103        }
104        assert!(!fired.load(Ordering::Relaxed));
105    }
106
107    #[test]
108    fn accepts_custom_on_drop_callback_impl() {
109        // Verify a named struct implementing OnDropCallback works through
110        // the guard, not just closures via the blanket impl.
111        struct Counter(Arc<AtomicBool>);
112        impl super::super::traits::OnDropCallback for Counter {
113            fn on_drop(self) {
114                self.0.store(true, Ordering::Relaxed);
115            }
116        }
117
118        let fired = Arc::new(AtomicBool::new(false));
119        {
120            let _guard = OnEarlyDropGuard::new(Counter(fired.clone()));
121        }
122        assert!(fired.load(Ordering::Relaxed));
123    }
124
125    // The guard must be Send and Sync whenever its callback is.
126    #[allow(dead_code)]
127    fn static_property_guard_is_send_sync() {
128        fn assert_send<T: Send>(_: &T) {}
129        fn assert_sync<T: Sync>(_: &T) {}
130
131        let guard = OnEarlyDropGuard::new(|| {});
132        assert_send(&guard);
133        assert_sync(&guard);
134    }
135}
tower_http/on_early_drop/guard.rs

tower_http/on_early_drop/
guard.rs
