use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll, Waker};
use std::time::{Duration, Instant};
use super::delayed_timer::TimerHandle;
struct SleepState {
fired: bool,
waker: Option<Waker>,
}
pub struct Sleep {
deadline: Instant,
timer: TimerHandle,
state: Arc<Mutex<SleepState>>,
armed: bool,
}
pub fn sleep(timer: &TimerHandle, dur: Duration) -> Sleep {
sleep_until(timer, Instant::now() + dur)
}
pub fn sleep_until(timer: &TimerHandle, deadline: Instant) -> Sleep {
Sleep {
deadline,
timer: timer.clone(),
state: Arc::new(Mutex::new(SleepState {
fired: false,
waker: None,
})),
armed: false,
}
}
impl Future for Sleep {
type Output = ();
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
let this = self.get_mut();
{
let mut st = this.state.lock().unwrap();
if st.fired {
return Poll::Ready(());
}
if Instant::now() >= this.deadline {
st.fired = true;
return Poll::Ready(());
}
st.waker = Some(cx.waker().clone());
}
if !this.armed {
let delay = this.deadline.saturating_duration_since(Instant::now());
let cb_state = Arc::clone(&this.state);
this.timer.schedule_wake(
delay,
Box::new(move || {
let mut st = cb_state.lock().unwrap();
st.fired = true;
if let Some(w) = st.waker.take() {
w.wake();
}
}),
);
this.armed = true;
}
Poll::Pending
}
}
impl Drop for Sleep {
fn drop(&mut self) {
self.state.lock().unwrap().waker = None;
}
}
pub struct TimerInterval {
timer: TimerHandle,
period: Duration,
next: Instant,
first: bool,
}
pub fn interval(timer: &TimerHandle, period: Duration) -> TimerInterval {
TimerInterval {
timer: timer.clone(),
period,
next: Instant::now() + period,
first: true,
}
}
impl TimerInterval {
pub async fn tick(&mut self) {
if self.first {
self.first = false;
return;
}
sleep_until(&self.timer, self.next).await;
self.next += self.period;
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::runtime::background::callback_executor::CallbackPool;
use crate::runtime::background::delayed_timer::DelayedTimer;
use crate::runtime::task::park_on_interruptible as drive;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::mpsc;
use std::task::Wake;
const T: Duration = Duration::from_secs(5);
struct CountWaker(Arc<AtomicUsize>);
impl Wake for CountWaker {
fn wake(self: Arc<Self>) {
self.0.fetch_add(1, Ordering::SeqCst);
}
fn wake_by_ref(self: &Arc<Self>) {
self.0.fetch_add(1, Ordering::SeqCst);
}
}
#[test]
fn sleep_completes_no_earlier_than_delay() {
let pool = CallbackPool::new();
let timer = DelayedTimer::new(pool.handle());
let delay = Duration::from_millis(60);
let start = Instant::now();
drive(sleep(&timer.handle(), delay), || false).unwrap();
let elapsed = start.elapsed();
assert!(
elapsed >= delay,
"sleep returned after {elapsed:?}, earlier than the {delay:?} delay"
);
}
#[test]
fn sleep_until_past_deadline_is_immediately_ready() {
let pool = CallbackPool::new();
let timer = DelayedTimer::new(pool.handle());
let past = Instant::now() - Duration::from_secs(1);
let count = Arc::new(AtomicUsize::new(0));
let waker = Waker::from(Arc::new(CountWaker(Arc::clone(&count))));
let mut cx = Context::from_waker(&waker);
let mut s = Box::pin(sleep_until(&timer.handle(), past));
assert!(s.as_mut().poll(&mut cx).is_ready());
assert_eq!(count.load(Ordering::SeqCst), 0);
}
#[test]
fn drop_before_deadline_wakes_nobody() {
let pool = CallbackPool::new();
let timer = DelayedTimer::new(pool.handle());
let count = Arc::new(AtomicUsize::new(0));
let waker = Waker::from(Arc::new(CountWaker(Arc::clone(&count))));
let mut cx = Context::from_waker(&waker);
let mut s = Box::pin(sleep(&timer.handle(), Duration::from_millis(60)));
assert!(s.as_mut().poll(&mut cx).is_pending());
drop(s);
std::thread::sleep(Duration::from_millis(140));
assert_eq!(
count.load(Ordering::SeqCst),
0,
"a dropped sleep must not wake a stale waker"
);
}
#[test]
fn concurrent_sleepers_complete_in_deadline_order() {
let pool = CallbackPool::new();
let timer = DelayedTimer::new(pool.handle());
let (tx, rx) = mpsc::channel();
let th_long = timer.handle();
let tx_long = tx.clone();
let long = std::thread::spawn(move || {
drive(sleep(&th_long, Duration::from_millis(150)), || false).unwrap();
tx_long.send("long").unwrap();
});
let th_short = timer.handle();
let short = std::thread::spawn(move || {
drive(sleep(&th_short, Duration::from_millis(30)), || false).unwrap();
tx.send("short").unwrap();
});
assert_eq!(rx.recv_timeout(T).unwrap(), "short");
assert_eq!(rx.recv_timeout(T).unwrap(), "long");
long.join().unwrap();
short.join().unwrap();
}
#[test]
fn interval_first_tick_immediate_then_periodic() {
let pool = CallbackPool::new();
let timer = DelayedTimer::new(pool.handle());
let period = Duration::from_millis(40);
let th = timer.handle();
let start = Instant::now();
drive(
async move {
let mut iv = interval(&th, period);
iv.tick().await; let after_first = start.elapsed();
assert!(
after_first < period,
"first tick should be immediate, was {after_first:?}"
);
iv.tick().await; iv.tick().await; },
|| false,
)
.unwrap();
assert!(
start.elapsed() >= 2 * period,
"two periodic ticks should take at least two periods, took {:?}",
start.elapsed()
);
}
#[test]
fn interval_bursts_to_catch_up_after_a_stall() {
let pool = CallbackPool::new();
let timer = DelayedTimer::new(pool.handle());
let period = Duration::from_millis(30);
let th = timer.handle();
drive(
async move {
let mut iv = interval(&th, period);
iv.tick().await; sleep(&th, Duration::from_millis(140)).await;
let t = Instant::now();
iv.tick().await; iv.tick().await; iv.tick().await; assert!(
t.elapsed() < period,
"overdue ticks must burst, three took {:?}",
t.elapsed()
);
},
|| false,
)
.unwrap();
}
}