1 //! An [`AbortOnDropHandle`] is like a [`JoinHandle`], except that it 2 //! will abort the task as soon as it is dropped. 3 4 use tokio::task::{AbortHandle, JoinError, JoinHandle}; 5 6 use 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)] 18 pub struct AbortOnDropHandle<T>(JoinHandle<T>); 19 20 impl<T> Drop for AbortOnDropHandle<T> { drop(&mut self)21 fn drop(&mut self) { 22 self.0.abort() 23 } 24 } 25 26 impl<T> AbortOnDropHandle<T> { 27 /// Create an [`AbortOnDropHandle`] from a [`JoinHandle`]. new(handle: JoinHandle<T>) -> Self28 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`]. abort(&self)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`]. is_finished(&self) -> bool40 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`]. abort_handle(&self) -> AbortHandle46 pub fn abort_handle(&self) -> AbortHandle { 47 self.0.abort_handle() 48 } 49 } 50 51 impl<T> Future for AbortOnDropHandle<T> { 52 type Output = Result<T, JoinError>; 53 poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output>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 59 impl<T> AsRef<JoinHandle<T>> for AbortOnDropHandle<T> { as_ref(&self) -> &JoinHandle<T>60 fn as_ref(&self) -> &JoinHandle<T> { 61 &self.0 62 } 63 } 64