sync: add DropGuardRef for CancellationToken (#7407)

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.

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();
+        }
+    }
+}

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;