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}