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"]
18pub struct AbortOnDropHandle<T>(JoinHandle<T>);
19
20impl<T> Drop for AbortOnDropHandle<T> {
21    fn drop(&mut self) {
22        self.0.abort()
23    }
24}
25
26impl<T> AbortOnDropHandle<T> {
27    /// Create an [`AbortOnDropHandle`] from a [`JoinHandle`].
28    pub fn new(handle: JoinHandle<T>) -> Self {
29        Self(handle)
30    }
31
32    /// Abort the task associated with this handle,
33    /// equivalent to [`JoinHandle::abort`].
34    pub fn abort(&self) {
35        self.0.abort()
36    }
37
38    /// Checks if the task associated with this handle is finished,
39    /// equivalent to [`JoinHandle::is_finished`].
40    pub fn is_finished(&self) -> bool {
41        self.0.is_finished()
42    }
43
44    /// Returns a new [`AbortHandle`] that can be used to remotely abort this task,
45    /// equivalent to [`JoinHandle::abort_handle`].
46    pub fn abort_handle(&self) -> AbortHandle {
47        self.0.abort_handle()
48    }
49
50    /// Cancels aborting on drop and returns the original [`JoinHandle`].
51    pub fn detach(self) -> JoinHandle<T> {
52        // Avoid invoking `AbortOnDropHandle`'s `Drop` impl
53        let this = ManuallyDrop::new(self);
54        // SAFETY: `&this.0` is a reference, so it is certainly initialized, and
55        // it won't be double-dropped because it's in a `ManuallyDrop`
56        unsafe { std::ptr::read(&this.0) }
57    }
58}
59
60impl<T> std::fmt::Debug for AbortOnDropHandle<T> {
61    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
62        f.debug_struct("AbortOnDropHandle")
63            .field("id", &self.0.id())
64            .finish()
65    }
66}
67
68impl<T> Future for AbortOnDropHandle<T> {
69    type Output = Result<T, JoinError>;
70
71    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
72        Pin::new(&mut self.0).poll(cx)
73    }
74}
75
76impl<T> AsRef<JoinHandle<T>> for AbortOnDropHandle<T> {
77    fn as_ref(&self) -> &JoinHandle<T> {
78        &self.0
79    }
80}
81
82#[cfg(test)]
83mod tests {
84    use super::*;
85
86    /// A simple type that does not implement [`std::fmt::Debug`].
87    struct NotDebug;
88
89    fn is_debug<T: std::fmt::Debug>() {}
90
91    #[test]
92    fn assert_debug() {
93        is_debug::<AbortOnDropHandle<NotDebug>>();
94    }
95}
tokio_util/task/abort_on_drop.rs

tokio_util/task/
abort_on_drop.rs
