diff --git a/tokio-util/src/task/abort_on_drop.rs b/tokio-util/src/task/abort_on_drop.rs index 3739bbfa9..e0353ac85 100644 --- a/tokio-util/src/task/abort_on_drop.rs +++ b/tokio-util/src/task/abort_on_drop.rs @@ -5,6 +5,7 @@ use tokio::task::{AbortHandle, JoinError, JoinHandle}; use std::{ future::Future, + mem::ManuallyDrop, pin::Pin, task::{Context, Poll}, }; @@ -46,6 +47,15 @@ impl AbortOnDropHandle { pub fn abort_handle(&self) -> AbortHandle { self.0.abort_handle() } + + /// Cancels aborting on drop and returns the original [`JoinHandle`]. + pub fn detach(self) -> JoinHandle { + // Avoid invoking `AbortOnDropHandle`'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 Future for AbortOnDropHandle { diff --git a/tokio-util/tests/abort_on_drop.rs b/tokio-util/tests/abort_on_drop.rs index c7dcee35a..20634e88c 100644 --- a/tokio-util/tests/abort_on_drop.rs +++ b/tokio-util/tests/abort_on_drop.rs @@ -1,4 +1,4 @@ -use tokio::sync::oneshot; +use tokio::{sync::oneshot, task::yield_now}; use tokio_util::task::AbortOnDropHandle; #[tokio::test] @@ -25,3 +25,15 @@ async fn aborts_task_directly() { assert!(tx.is_closed()); assert!(handle.is_finished()); } + +#[tokio::test] +async fn does_not_abort_after_detach() { + let (tx, rx) = oneshot::channel::(); + let handle = tokio::spawn(async move { + let _ = rx.await; + }); + 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 +}