Skip to main content

nuclei/sync/cancellation/
guard.rs

1//! The guards that cancel a token when dropped.
2
3use super::CancellationToken;
4
5/// A wrapper for a [`CancellationToken`] which cancels it on drop, created by
6/// [`CancellationToken::drop_guard`].
7///
8/// ```
9/// use nuclei::sync::CancellationToken;
10///
11/// let token = CancellationToken::new();
12/// let guard = token.clone().drop_guard();
13/// let still_alive = guard.disarm();
14/// assert!(!token.is_cancelled());
15/// drop(still_alive);
16/// assert!(!token.is_cancelled());
17/// ```
18#[derive(Debug)]
19pub struct DropGuard {
20    pub(super) token: CancellationToken,
21    pub(super) armed: bool,
22}
23
24impl DropGuard {
25    /// Returns a reference to the token this guard wraps.
26    pub fn token(&self) -> &CancellationToken {
27        &self.token
28    }
29
30    /// Returns the token and disarms this guard: it no longer cancels the
31    /// token. Other guards of the same token are not affected.
32    pub fn disarm(mut self) -> CancellationToken {
33        self.armed = false;
34        self.token.clone()
35    }
36}
37
38impl Drop for DropGuard {
39    fn drop(&mut self) {
40        if self.armed {
41            self.token.cancel();
42        }
43    }
44}
45
46/// A wrapper for a borrowed [`CancellationToken`] which cancels it on drop,
47/// created by [`CancellationToken::drop_guard_ref`].
48#[derive(Debug)]
49pub struct DropGuardRef<'a> {
50    pub(super) token: &'a CancellationToken,
51    pub(super) armed: bool,
52}
53
54impl<'a> DropGuardRef<'a> {
55    /// Returns a reference to the token this guard wraps.
56    pub fn token(&self) -> &CancellationToken {
57        self.token
58    }
59
60    /// Returns the token and disarms this guard: it no longer cancels the
61    /// token. Other guards of the same token are not affected.
62    pub fn disarm(mut self) -> &'a CancellationToken {
63        self.armed = false;
64        self.token
65    }
66}
67
68impl Drop for DropGuardRef<'_> {
69    fn drop(&mut self) {
70        if self.armed {
71            self.token.cancel();
72        }
73    }
74}