tokio_util/task/
abort_on_drop.rs

1//! An [`AbortOnDropHandle`] is like a [`JoinHandle`], except that it
2//! will abort the task as soon as it is dropped.
3
4use tokio::task::{AbortHandle, JoinError, JoinHandle};
5
6use std::{
7    future::Future,
8    mem::ManuallyDrop,
9    pin::Pin,
10    task::{Context, Poll},
11};
12
13/// A wrapper around a [`tokio::task::JoinHandle`],
14/// which [aborts] the task when it is dropped.
15///
16/// [aborts]: tokio::task::JoinHandle::abort
17#[must_use = "Dropping the handle aborts the task immediately"]
18#[derive(Debug)]
19pub struct AbortOnDropHandle<T>(JoinHandle<T>);
20
21impl<T> Drop for AbortOnDropHandle<T> {
22    fn drop(&mut self) {
23        self.0.abort()
24    }
25}
26
27impl<T> AbortOnDropHandle<T> {
28    /// Create an [`AbortOnDropHandle`] from a [`JoinHandle`].
29    pub fn new(handle: JoinHandle<T>) -> Self {
30        Self(handle)
31    }
32
33    /// Abort the task associated with this handle,
34    /// equivalent to [`JoinHandle::abort`].
35    pub fn abort(&self) {
36        self.0.abort()
37    }
38
39    /// Checks if the task associated with this handle is finished,
40    /// equivalent to [`JoinHandle::is_finished`].
41    pub fn is_finished(&self) -> bool {
42        self.0.is_finished()
43    }
44
45    /// Returns a new [`AbortHandle`] that can be used to remotely abort this task,
46    /// equivalent to [`JoinHandle::abort_handle`].
47    pub fn abort_handle(&self) -> AbortHandle {
48        self.0.abort_handle()
49    }
50
51    /// Cancels aborting on drop and returns the original [`JoinHandle`].
52    pub fn detach(self) -> JoinHandle<T> {
53        // Avoid invoking `AbortOnDropHandle`'s `Drop` impl
54        let this = ManuallyDrop::new(self);
55        // SAFETY: `&this.0` is a reference, so it is certainly initialized, and
56        // it won't be double-dropped because it's in a `ManuallyDrop`
57        unsafe { std::ptr::read(&this.0) }
58    }
59}
60
61impl<T> Future for AbortOnDropHandle<T> {
62    type Output = Result<T, JoinError>;
63
64    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
65        Pin::new(&mut self.0).poll(cx)
66    }
67}
68
69impl<T> AsRef<JoinHandle<T>> for AbortOnDropHandle<T> {
70    fn as_ref(&self) -> &JoinHandle<T> {
71        &self.0
72    }
73}
tokio_util/task/abort_on_drop.rs

tokio_util/task/
abort_on_drop.rs
