sync: add drop guard for cancellation token (#3839)

tokio-util/src/sync/cancellation_token.rs

@@ -1,5 +1,6 @@
 //! An asynchronously awaitable `CancellationToken`.
 //! The token allows to signal a cancellation request to one or more tasks.
+pub(crate) mod guard;
 
 use crate::loom::sync::atomic::AtomicUsize;
 use crate::loom::sync::Mutex;
@@ -11,6 +12,8 @@ use core::ptr::NonNull;
 use core::sync::atomic::Ordering;
 use core::task::{Context, Poll, Waker};
 
+use guard::DropGuard;
+
 /// A token which can be used to signal a cancellation request to one or more
 /// tasks.
 ///
@@ -275,6 +278,14 @@ impl CancellationToken {
         }
     }
 
+    /// Creates a `DropGuard` for this token.
+    ///
+    /// Returned guard will cancel this token (and all its children) on drop
+    /// unless disarmed.
+    pub fn drop_guard(self) -> DropGuard {
+        DropGuard { inner: Some(self) }
+    }
+
     unsafe fn register(
         &self,
         wait_node: &mut ListNode<WaitQueueEntry>,

tokio-util/src/sync/cancellation_token/guard.rs

@@ -0,0 +1,27 @@
+use crate::sync::CancellationToken;
+
+/// A wrapper for cancellation token which automatically cancels
+/// it on drop. It is created using `drop_guard` method on the `CancellationToken`.
+#[derive(Debug)]
+pub struct DropGuard {
+    pub(super) inner: Option<CancellationToken>,
+}
+
+impl DropGuard {
+    /// 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) -> CancellationToken {
+        self.inner
+            .take()
+            .expect("`inner` can be only None in a destructor")
+    }
+}
+
+impl Drop for DropGuard {
+    fn drop(&mut self) {
+        if let Some(inner) = &self.inner {
+            inner.cancel();
+        }
+    }
+}

tokio-util/src/sync/mod.rs

@@ -1,7 +1,7 @@
 //! Synchronization primitives
 
 mod cancellation_token;
-pub use cancellation_token::{CancellationToken, WaitForCancellationFuture};
+pub use cancellation_token::{guard::DropGuard, CancellationToken, WaitForCancellationFuture};
 
 mod intrusive_double_linked_list;