use crate::{item::DelayItem, reserve::ReserveWaker, sleep::SleepLocal};
use async_condvar_fair::{Baton, BatonExt, Condvar};
use parking_lot::Mutex;
use std::{
collections::VecDeque,
future::Future,
pin::Pin,
sync::{atomic::Ordering, Arc},
task::Poll,
time::Instant,
};
type CondvarWaitFuture<'a> = dyn Future<Output = Option<Baton<'a>>> + Send + 'a;
pub struct DelayedPopFuture<'a, T> {
pub(super) size_condvar: &'a Condvar,
pub(super) reserve_condvar: &'a Condvar,
pub(super) queue: &'a Mutex<VecDeque<DelayItem<T>>>,
pub(super) condvar_future: Option<Pin<Box<CondvarWaitFuture<'a>>>>,
pub(super) sleep_future: Option<SleepLocal>,
pub(super) reserve_waker: Option<ReserveWaker<'a>>,
pub(super) future_id: u64,
}
impl<'a, T> Unpin for DelayedPopFuture<'a, T> {}
impl<'a, T: Send> Future for DelayedPopFuture<'a, T> {
type Output = T;
fn poll(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Self::Output> {
'main_loop: loop {
if let Some(condvar_future) = self.as_mut().condvar_future.as_mut() {
match condvar_future.as_mut().poll(cx) {
Poll::Pending => {
return Poll::Pending;
}
Poll::Ready(res) => {
res.dispose();
self.as_mut().condvar_future.take();
continue 'main_loop;
}
}
}
if let Some(sleep_future) = self.as_mut().sleep_future.as_mut() {
match Pin::new(sleep_future).poll(cx) {
Poll::Pending => {
return Poll::Pending;
}
Poll::Ready(_) => {
self.as_mut().sleep_future.take();
continue 'main_loop;
}
}
}
let mut lock = self.queue.lock();
if let Some(front_val) = lock.front_mut() {
let reserved_id = front_val.reserved.load(Ordering::Acquire);
if (reserved_id > 0) && (reserved_id != self.future_id) {
let wait_future = self.reserve_condvar.wait_no_relock(lock);
let prev = self.as_mut().condvar_future.replace(Box::pin(wait_future));
assert!(prev.is_none(), "Previous condvar wait should not exist");
continue 'main_loop;
}
else if front_val.pop_time > Instant::now() {
let prev = self.as_mut().sleep_future.replace(SleepLocal::new(
tokio::time::sleep_until(front_val.pop_time.into()),
));
assert!(prev.is_none(), "Sleep future should not exist");
front_val.reserved.store(self.future_id, Ordering::Release);
self.reserve_waker = Some(ReserveWaker {
condvar: self.reserve_condvar,
item_reserved_future: Arc::downgrade(&front_val.reserved),
});
drop(lock);
continue 'main_loop;
}
else {
let item = lock.pop_front().expect("First item should exist").item;
drop(lock);
self.size_condvar.notify_one();
return Poll::Ready(item);
}
} else {
let wait_future = self.size_condvar.wait_no_relock(lock);
let prev = self.as_mut().condvar_future.replace(Box::pin(wait_future));
assert!(prev.is_none(), "Previous condvar wait should not exist");
continue 'main_loop;
}
}
}
}