From d35ff7064f2e0f90764487f45606dba57f3beaba Mon Sep 17 00:00:00 2001
From: Mikail Bagishov <bagishov.mikail@yandex.ru>
Date: Mon, 28 Jun 2021 13:34:58 +0300
Subject: [PATCH] sync: add drop guard for cancellation token (#3839)

---
 tokio-util/src/sync/cancellation_token.rs     | 11 ++++++++
 .../src/sync/cancellation_token/guard.rs      | 27 +++++++++++++++++++
 tokio-util/src/sync/mod.rs                    |  2 +-
 3 files changed, 39 insertions(+), 1 deletion(-)
 create mode 100644 tokio-util/src/sync/cancellation_token/guard.rs

diff --git a/tokio-util/src/sync/cancellation_token.rs b/tokio-util/src/sync/cancellation_token.rs
index a239a7716..2488cbf02 100644
--- a/tokio-util/src/sync/cancellation_token.rs
+++ b/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>,
diff --git a/tokio-util/src/sync/cancellation_token/guard.rs b/tokio-util/src/sync/cancellation_token/guard.rs
new file mode 100644
index 000000000..54ed7ea2e
--- /dev/null
+++ b/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();
+        }
+    }
+}
diff --git a/tokio-util/src/sync/mod.rs b/tokio-util/src/sync/mod.rs
index 34493c2cb..0b78a156c 100644
--- a/tokio-util/src/sync/mod.rs
+++ b/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;