use std::{
future::Future,
pin::Pin,
sync::Arc,
task::{
Context,
Poll,
},
};
use super::TaskResult;
use super::task_completion::TaskCompletion;
use super::task_handle_inner::TaskHandleInner;
pub struct TaskHandle<R, E> {
pub(crate) inner: Arc<TaskHandleInner<R, E>>,
}
impl<R, E> TaskHandle<R, E> {
pub fn get(self) -> TaskResult<R, E> {
self.inner.state.wait_until(
|state| state.completed,
|state| {
state
.result
.take()
.expect("task handle completed without a result")
},
)
}
#[inline]
pub fn is_done(&self) -> bool {
self.inner.done.load()
}
#[inline]
pub fn cancel(&self) -> bool {
TaskCompletion {
inner: Arc::clone(&self.inner),
}
.cancel()
}
}
impl<R, E> Future for TaskHandle<R, E> {
type Output = TaskResult<R, E>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let result = self.inner.state.write(|state| {
if state.completed {
Some(
state
.result
.take()
.expect("task handle completed without a result"),
)
} else {
state.waker = Some(cx.waker().clone());
None
}
});
if let Some(result) = result {
Poll::Ready(result)
} else {
Poll::Pending
}
}
}