From 8167f871378246ab2108d1ea144257cf95934e4f Mon Sep 17 00:00:00 2001 From: Alex H Date: Mon, 9 Feb 2026 04:23:46 -0800 Subject: [PATCH] task: add `AbortOnDrop` (#7855) --- tokio-util/src/task/abort_on_drop.rs | 64 +++++++++++++++++++++++++++- tokio-util/src/task/mod.rs | 2 +- tokio-util/tests/abort_on_drop.rs | 41 +++++++++++++++++- 3 files changed, 103 insertions(+), 4 deletions(-) diff --git a/tokio-util/src/task/abort_on_drop.rs b/tokio-util/src/task/abort_on_drop.rs index 191731cb9..d06ebd5fa 100644 --- a/tokio-util/src/task/abort_on_drop.rs +++ b/tokio-util/src/task/abort_on_drop.rs @@ -1,5 +1,8 @@ //! An [`AbortOnDropHandle`] is like a [`JoinHandle`], except that it //! will abort the task as soon as it is dropped. +//! +//! Correspondingly, an [`AbortOnDrop`] is like a [`AbortHandle`] that will abort +//! the task as soon as it is dropped. use tokio::task::{AbortHandle, JoinError, JoinHandle}; @@ -19,7 +22,7 @@ pub struct AbortOnDropHandle(JoinHandle); impl Drop for AbortOnDropHandle { fn drop(&mut self) { - self.0.abort() + self.abort() } } @@ -31,12 +34,14 @@ impl AbortOnDropHandle { /// Abort the task associated with this handle, /// equivalent to [`JoinHandle::abort`]. + #[inline] pub fn abort(&self) { self.0.abort() } /// Checks if the task associated with this handle is finished, /// equivalent to [`JoinHandle::is_finished`]. + #[inline] pub fn is_finished(&self) -> bool { self.0.is_finished() } @@ -79,6 +84,62 @@ impl AsRef> for AbortOnDropHandle { } } +/// A wrapper around a [`tokio::task::AbortHandle`], +/// which [aborts] the task when it is dropped. +/// +/// Unlike [`AbortOnDropHandle`], [`AbortOnDrop`] cannot be `.await`ed for a result. +/// +/// It has no generic parameter, making it suitable when you only need to keep +/// a task handle in a struct and do not care about the output. +/// +/// [aborts]: tokio::task::AbortHandle::abort +#[must_use = "Dropping the handle aborts the task immediately"] +pub struct AbortOnDrop(AbortHandle); + +impl Drop for AbortOnDrop { + fn drop(&mut self) { + self.abort() + } +} + +impl AbortOnDrop { + /// Create an [`AbortOnDrop`] from a [`AbortHandle`]. + pub fn new(handle: AbortHandle) -> Self { + Self(handle) + } + + /// Abort the task associated with this handle, + /// equivalent to [`AbortHandle::abort`]. + #[inline] + pub fn abort(&self) { + self.0.abort() + } + + /// Checks if the task associated with this handle is finished, + /// equivalent to [`AbortHandle::is_finished`]. + #[inline] + pub fn is_finished(&self) -> bool { + self.0.is_finished() + } + + /// Cancels aborting on drop and returns the original [`AbortHandle`]. + pub fn detach(self) -> AbortHandle { + // Avoid invoking `AbortOnDrop`'s `Drop` impl + let this = ManuallyDrop::new(self); + // SAFETY: `&this.0` is a reference, so it is certainly initialized, and + // it won't be double-dropped because it's in a `ManuallyDrop` + unsafe { std::ptr::read(&this.0) } + } +} + +impl std::fmt::Debug for AbortOnDrop { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AbortOnDrop") + .field("id", &self.0.id()) + .finish() + } +} + #[cfg(test)] mod tests { use super::*; @@ -90,6 +151,7 @@ mod tests { #[test] fn assert_debug() { + is_debug::(); is_debug::>(); } } diff --git a/tokio-util/src/task/mod.rs b/tokio-util/src/task/mod.rs index 3c35ebff4..87870e21f 100644 --- a/tokio-util/src/task/mod.rs +++ b/tokio-util/src/task/mod.rs @@ -12,7 +12,7 @@ cfg_rt! { pub use task_tracker::TaskTracker; mod abort_on_drop; - pub use abort_on_drop::AbortOnDropHandle; + pub use abort_on_drop::{AbortOnDrop, AbortOnDropHandle}; mod join_queue; pub use join_queue::JoinQueue; diff --git a/tokio-util/tests/abort_on_drop.rs b/tokio-util/tests/abort_on_drop.rs index 20634e88c..2f97ab4bd 100644 --- a/tokio-util/tests/abort_on_drop.rs +++ b/tokio-util/tests/abort_on_drop.rs @@ -1,5 +1,5 @@ use tokio::{sync::oneshot, task::yield_now}; -use tokio_util::task::AbortOnDropHandle; +use tokio_util::task::{AbortOnDrop, AbortOnDropHandle}; #[tokio::test] async fn aborts_task_on_drop() { @@ -35,5 +35,42 @@ async fn does_not_abort_after_detach() { let handle = AbortOnDropHandle::new(handle); handle.detach(); // returns and drops the original join handle yield_now().await; - assert!(!tx.is_closed()); // task is still live + assert!(!tx.is_closed()); // the task is still alive +} + +#[tokio::test] +async fn handle_aborts_task_on_drop() { + let (mut tx, rx) = oneshot::channel::(); + let handle = tokio::spawn(async move { + let _ = rx.await; + }); + let handle = AbortOnDrop::new(handle.abort_handle()); + drop(handle); + tx.closed().await; + assert!(tx.is_closed()); +} + +#[tokio::test] +async fn handle_aborts_task_directly() { + let (mut tx, rx) = oneshot::channel::(); + let handle = tokio::spawn(async move { + let _ = rx.await; + }); + let handle = AbortOnDrop::new(handle.abort_handle()); + handle.abort(); + tx.closed().await; + assert!(tx.is_closed()); + assert!(handle.is_finished()); +} + +#[tokio::test] +async fn handle_does_not_abort_after_detach() { + let (tx, rx) = oneshot::channel::(); + let handle = tokio::spawn(async move { + let _ = rx.await; + }); + let handle = AbortOnDrop::new(handle.abort_handle()); + handle.detach(); // returns and drops the original abort handle + yield_now().await; + assert!(!tx.is_closed()); // the task is still alive }