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    pin::Pin,
9    task::{Context, Poll},
10};
11
12/// A wrapper around a [`tokio::task::JoinHandle`],
13/// which [aborts] the task when it is dropped.
14///
15/// [aborts]: tokio::task::JoinHandle::abort
16#[must_use = "Dropping the handle aborts the task immediately"]
17#[derive(Debug)]
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
51impl<T> Future for AbortOnDropHandle<T> {
52    type Output = Result<T, JoinError>;
53
54    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
55        Pin::new(&mut self.0).poll(cx)
56    }
57}
58
59impl<T> AsRef<JoinHandle<T>> for AbortOnDropHandle<T> {
60    fn as_ref(&self) -> &JoinHandle<T> {
61        &self.0
62    }
63}