diff --git a/tokio-util/src/sync/cancellation_token.rs b/tokio-util/src/sync/cancellation_token.rs index 85305de47..64cff2199 100644 --- a/tokio-util/src/sync/cancellation_token.rs +++ b/tokio-util/src/sync/cancellation_token.rs @@ -1,6 +1,7 @@ //! An asynchronously awaitable `CancellationToken`. //! The token allows to signal a cancellation request to one or more tasks. pub(crate) mod guard; +pub(crate) mod guard_ref; mod tree_node; use crate::loom::sync::Arc; @@ -10,6 +11,7 @@ use core::pin::Pin; use core::task::{Context, Poll}; use guard::DropGuard; +use guard_ref::DropGuardRef; use pin_project_lite::pin_project; /// A token which can be used to signal a cancellation request to one or more @@ -242,6 +244,14 @@ impl CancellationToken { DropGuard { inner: Some(self) } } + /// Creates a `DropGuardRef` for this token. + /// + /// Returned guard will cancel this token (and all its children) on drop + /// unless disarmed. + pub fn drop_guard_ref(&self) -> DropGuardRef<'_> { + DropGuardRef { inner: Some(self) } + } + /// Runs a future to completion and returns its result wrapped inside of an `Option` /// unless the `CancellationToken` is cancelled. In that case the function returns /// `None` and the future gets dropped. diff --git a/tokio-util/src/sync/cancellation_token/guard_ref.rs b/tokio-util/src/sync/cancellation_token/guard_ref.rs new file mode 100644 index 000000000..b48fd19de --- /dev/null +++ b/tokio-util/src/sync/cancellation_token/guard_ref.rs @@ -0,0 +1,29 @@ +use crate::sync::CancellationToken; + +/// A wrapper for cancellation token which automatically cancels +/// it on drop. It is created using `drop_guard_ref` method on the `CancellationToken`. +/// +/// This is a `ref` version of `DropGuard` +#[derive(Debug)] +pub struct DropGuardRef<'a> { + pub(super) inner: Option<&'a CancellationToken>, +} + +impl<'a> DropGuardRef<'a> { + /// Returns stored cancellation token and removes this drop guard instance + /// (i.e. it will no longer cancel token). Other guards for this token + /// are not affected. + pub fn disarm(mut self) -> &'a CancellationToken { + self.inner + .take() + .expect("`inner` can be only None in a destructor") + } +} + +impl Drop for DropGuardRef<'_> { + fn drop(&mut self) { + if let Some(inner) = self.inner { + inner.cancel(); + } + } +} diff --git a/tokio-util/src/sync/mod.rs b/tokio-util/src/sync/mod.rs index 1ae1b7ee4..28f2fbd2d 100644 --- a/tokio-util/src/sync/mod.rs +++ b/tokio-util/src/sync/mod.rs @@ -2,7 +2,8 @@ mod cancellation_token; pub use cancellation_token::{ - guard::DropGuard, CancellationToken, WaitForCancellationFuture, WaitForCancellationFutureOwned, + guard::DropGuard, guard_ref::DropGuardRef, CancellationToken, WaitForCancellationFuture, + WaitForCancellationFutureOwned, }; mod mpsc;