use crate::DecodeError;
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll, Waker};
type Outcome<T> = Result<T, DecodeError>;
type BoxedFuture<'a, T> = Pin<Box<dyn Future<Output = Outcome<T>> + 'a>>;
pub struct Pending<'a, T> {
future: Option<BoxedFuture<'a, T>>,
}
impl<'a, T> Pending<'a, T> {
pub(crate) fn from_future(future: impl Future<Output = Outcome<T>> + 'a) -> Self {
Self {
future: Some(Box::pin(future)),
}
}
pub fn try_take(&mut self) -> Option<Outcome<T>> {
let mut context = Context::from_waker(Waker::noop());
match self.poll_inner(&mut context) {
Poll::Ready(result) => Some(result),
Poll::Pending => None,
}
}
fn poll_inner(&mut self, context: &mut Context<'_>) -> Poll<Outcome<T>> {
let future = self
.future
.as_mut()
.expect("Pending polled after it completed");
let result = future.as_mut().poll(context);
if result.is_ready() {
self.future = None;
}
result
}
}
impl<T> Future for Pending<'_, T> {
type Output = Outcome<T>;
fn poll(self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Self::Output> {
self.get_mut().poll_inner(context)
}
}
impl<T> std::fmt::Debug for Pending<'_, T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Pending")
.field("completed", &self.future.is_none())
.finish()
}
}
#[cfg(all(feature = "worker", not(target_arch = "wasm32")))]
pub(crate) use shared::{reply, Reply};
#[cfg(all(feature = "worker", not(target_arch = "wasm32")))]
mod shared {
use super::*;
use std::sync::{Arc, Mutex};
enum State<T> {
Waiting(Option<Waker>),
Ready(Outcome<T>),
Taken,
Cancelled,
}
struct Slot<T>(Mutex<State<T>>);
impl<T> Slot<T> {
fn lock(&self) -> std::sync::MutexGuard<'_, State<T>> {
self.0
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
}
}
pub(crate) fn reply<T>() -> (impl Future<Output = Outcome<T>>, Reply<T>) {
let slot = Arc::new(Slot(Mutex::new(State::Waiting(None))));
(Awaiting(slot.clone()), Reply(slot))
}
struct Awaiting<T>(Arc<Slot<T>>);
impl<T> Future for Awaiting<T> {
type Output = Outcome<T>;
fn poll(self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Self::Output> {
let mut state = self.0.lock();
match &mut *state {
State::Waiting(waker) => {
*waker = Some(context.waker().clone());
Poll::Pending
}
State::Ready(_) => match std::mem::replace(&mut *state, State::Taken) {
State::Ready(result) => Poll::Ready(result),
_ => unreachable!("matched Ready under the same lock"),
},
State::Taken => panic!("decode result taken twice"),
State::Cancelled => unreachable!("only a dropped future cancels"),
}
}
}
impl<T> Drop for Awaiting<T> {
fn drop(&mut self) {
let mut state = self.0.lock();
if matches!(*state, State::Waiting(_)) {
*state = State::Cancelled;
}
}
}
pub(crate) struct Reply<T>(Arc<Slot<T>>);
impl<T> Reply<T> {
pub(crate) fn is_cancelled(&self) -> bool {
matches!(*self.0.lock(), State::Cancelled)
}
pub(crate) fn send(self, result: Outcome<T>) -> Result<(), Outcome<T>> {
deliver(&self.0, result)
}
}
impl<T> Drop for Reply<T> {
fn drop(&mut self) {
let _ = deliver(&self.0, Err(DecodeError::Closed));
}
}
fn deliver<T>(slot: &Slot<T>, result: Outcome<T>) -> Result<(), Outcome<T>> {
let mut state = slot.lock();
let State::Waiting(waker) = &mut *state else {
return Err(result);
};
let waker = waker.take();
*state = State::Ready(result);
drop(state);
if let Some(waker) = waker {
waker.wake();
}
Ok(())
}
}