-
Notifications
You must be signed in to change notification settings - Fork 1.4k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
code motion: AbortOnDropJoinHandle lives in a mod just for now
at some point this will get factored to be tokio-specific but not just now
- Loading branch information
Showing
1 changed file
with
53 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,53 @@ | ||
use std::future::Future; | ||
use std::pin::Pin; | ||
use std::task::{Context, Poll}; | ||
/// Exactly like a [`tokio::task::JoinHandle`], except that it aborts the task when | ||
/// the handle is dropped. | ||
/// | ||
/// This behavior makes it easier to tie a worker task to the lifetime of a Resource | ||
/// by keeping this handle owned by the Resource. | ||
#[derive(Debug)] | ||
pub struct AbortOnDropJoinHandle<T>(tokio::task::JoinHandle<T>); | ||
impl<T> AbortOnDropJoinHandle<T> { | ||
/// Abort the task and wait for it to finish. Optionally returns the result | ||
/// of the task if it ran to completion prior to being aborted. | ||
pub(crate) async fn cancel(mut self) -> Option<T> { | ||
self.0.abort(); | ||
|
||
match (&mut self.0).await { | ||
Ok(value) => Some(value), | ||
Err(err) if err.is_cancelled() => None, | ||
Err(err) => std::panic::resume_unwind(err.into_panic()), | ||
} | ||
} | ||
} | ||
impl<T> Drop for AbortOnDropJoinHandle<T> { | ||
fn drop(&mut self) { | ||
self.0.abort() | ||
} | ||
} | ||
impl<T> std::ops::Deref for AbortOnDropJoinHandle<T> { | ||
type Target = tokio::task::JoinHandle<T>; | ||
fn deref(&self) -> &Self::Target { | ||
&self.0 | ||
} | ||
} | ||
impl<T> std::ops::DerefMut for AbortOnDropJoinHandle<T> { | ||
fn deref_mut(&mut self) -> &mut tokio::task::JoinHandle<T> { | ||
&mut self.0 | ||
} | ||
} | ||
impl<T> From<tokio::task::JoinHandle<T>> for AbortOnDropJoinHandle<T> { | ||
fn from(jh: tokio::task::JoinHandle<T>) -> Self { | ||
AbortOnDropJoinHandle(jh) | ||
} | ||
} | ||
impl<T> Future for AbortOnDropJoinHandle<T> { | ||
type Output = T; | ||
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { | ||
match Pin::new(&mut self.as_mut().0).poll(cx) { | ||
Poll::Pending => Poll::Pending, | ||
Poll::Ready(r) => Poll::Ready(r.expect("child task panicked")), | ||
} | ||
} | ||
} |