use std::{pin::Pin, task::Poll};
pub use web_async::time::{Duration, Instant};
use crate::Waiter;
pub struct Deadline {
at: Option<Instant>,
sleep: Option<Pin<Box<web_async::time::Sleep>>>,
}
impl Deadline {
pub fn new() -> Self {
Self { at: None, sleep: None }
}
pub fn at(at: Instant) -> Self {
Self {
at: Some(at),
sleep: None,
}
}
pub fn after(duration: Duration) -> Self {
Self {
at: Instant::now().checked_add(duration),
sleep: None,
}
}
pub fn set(&mut self, at: Option<Instant>) {
if self.at == at {
return;
}
self.at = at;
if let (Some(at), Some(sleep)) = (at, &mut self.sleep) {
sleep.as_mut().reset(at);
}
}
pub fn deadline(&self) -> Option<Instant> {
self.at
}
pub fn poll(&mut self, waiter: &Waiter) -> Poll<()> {
let Some(at) = self.at else { return Poll::Pending };
let sleep = self
.sleep
.get_or_insert_with(|| Box::pin(web_async::time::sleep_until(at)));
if sleep.is_elapsed() {
return Poll::Ready(());
}
waiter.poll_future(sleep.as_mut())
}
pub async fn wait(&mut self) {
crate::wait(|waiter| self.poll(waiter)).await
}
}
impl Default for Deadline {
fn default() -> Self {
Self::new()
}
}
impl std::fmt::Debug for Deadline {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Deadline").field("at", &self.at).finish()
}
}
#[cfg(test)]
mod tests {
use std::task::Waker;
use super::*;
fn poll_once(deadline: &mut Deadline) -> Poll<()> {
let waiter = Waiter::new(Waker::noop().clone());
deadline.poll(&waiter)
}
#[tokio::test(start_paused = true)]
async fn fires_at_its_deadline() {
let at = Instant::now() + Duration::from_secs(5);
let mut deadline = Deadline::at(at);
deadline.wait().await;
assert!(Instant::now() >= at, "returned before the deadline");
}
#[tokio::test(start_paused = true)]
async fn an_unrepresentable_duration_disarms_instead_of_panicking() {
let mut deadline = Deadline::after(Duration::MAX);
assert_eq!(deadline.deadline(), None);
assert!(poll_once(&mut deadline).is_pending());
}
#[tokio::test(start_paused = true)]
async fn disarmed_never_fires() {
let mut deadline = Deadline::new();
assert!(poll_once(&mut deadline).is_pending());
let res = tokio::time::timeout(Duration::from_secs(60), deadline.wait()).await;
assert!(res.is_err(), "a disarmed deadline fired");
}
#[tokio::test(start_paused = true)]
async fn stays_ready_once_elapsed() {
let mut deadline = Deadline::after(Duration::from_secs(1));
deadline.wait().await;
assert!(poll_once(&mut deadline).is_ready());
assert!(poll_once(&mut deadline).is_ready());
}
#[tokio::test(start_paused = true)]
async fn re_arming_to_the_same_instant_does_not_restart() {
let at = Instant::now() + Duration::from_secs(1);
let mut deadline = Deadline::at(at);
deadline.wait().await;
deadline.set(Some(at));
assert!(poll_once(&mut deadline).is_ready());
}
#[tokio::test(start_paused = true)]
async fn re_arming_later_defers_the_fire() {
let start = Instant::now();
let mut deadline = Deadline::after(Duration::from_secs(1));
assert!(poll_once(&mut deadline).is_pending());
let later = start + Duration::from_secs(10);
deadline.set(Some(later));
deadline.wait().await;
assert!(Instant::now() >= later, "fired at the original deadline");
}
#[tokio::test(start_paused = true)]
async fn disarming_a_live_countdown_stops_it() {
let mut deadline = Deadline::after(Duration::from_secs(1));
assert!(poll_once(&mut deadline).is_pending());
deadline.set(None);
assert_eq!(deadline.deadline(), None);
let res = tokio::time::timeout(Duration::from_secs(60), deadline.wait()).await;
assert!(res.is_err(), "a disarmed deadline fired");
}
#[tokio::test(start_paused = true)]
async fn re_arming_after_disarm_fires_again() {
let mut deadline = Deadline::after(Duration::from_secs(1));
assert!(poll_once(&mut deadline).is_pending());
deadline.set(None);
let at = Instant::now() + Duration::from_secs(3);
deadline.set(Some(at));
deadline.wait().await;
assert!(Instant::now() >= at, "returned before the re-armed deadline");
}
}