use super::{
mailbox::{Cancel, Forward},
registration::{Key, Observation, Registration},
runtime::{Deferred, Local},
slab::{Id, Slab},
timeout::TimeoutWheel,
};
use crate::Error;
use commonware_utils::channel::oneshot;
use std::{
cmp::Reverse,
collections::BinaryHeap,
future::Future,
mem,
pin::Pin,
sync::Arc,
task::{Context, Poll, Waker},
time::{Duration, Instant, SystemTime},
};
enum State {
Unregistered {
deadline: Instant,
},
Registered(Registration<TimerId>),
Done,
}
pub struct Sleep {
state: State,
}
impl Sleep {
pub fn new(duration: Duration) -> Self {
if duration.is_zero() {
return Self { state: State::Done };
}
let deadline = Instant::now()
.checked_add(duration.min(TimeoutWheel::MAX_TIMEOUT))
.expect("sleep deadline clamped to TimeoutWheel::MAX_TIMEOUT is not representable");
Self {
state: State::Unregistered { deadline },
}
}
pub fn until(deadline: SystemTime) -> Self {
let duration = deadline
.duration_since(SystemTime::now())
.unwrap_or_default();
Self::new(duration)
}
}
impl Future for Sleep {
type Output = ();
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
let this = self.get_mut();
match &mut this.state {
State::Done => Poll::Ready(()),
State::Registered(registration) => {
std::task::ready!(Pin::new(registration).poll(cx))
.expect("io_uring sleep polled after its worker closed");
this.state = State::Done;
Poll::Ready(())
}
State::Unregistered { deadline } => {
if *deadline <= Instant::now() {
this.state = State::Done;
return Poll::Ready(());
}
let owner = Local::current().expect("io_uring sleep requires a current worker");
assert!(
!owner.borrow().closing,
"io_uring sleep polled after its worker closed"
);
let waker = cx.waker().clone();
let mut local = owner.borrow_mut();
let id = local.timers.insert(*deadline, waker);
this.state =
State::Registered(Registration::new(Arc::downgrade(&local.mailbox), id));
Poll::Pending
}
}
}
}
#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)]
pub struct TimerId(Id);
impl Key for TimerId {
type Output = ();
fn observe(self, local: &mut Local, waker: &Waker) -> Observation<()> {
local.timers.observe(self, waker)
}
fn refresh(self, local: &mut Local, waker: Waker) -> Option<Waker> {
Some(
local
.timers
.refresh(self, waker)
.expect("live io_uring sleeper registration missing"),
)
}
fn forward(self, sender: oneshot::Sender<Result<(), Error>>) -> Forward {
Forward::Timer(self, sender)
}
fn cancel(self) -> Cancel {
Cancel::Timer(self)
}
}
enum Observer {
Local(Waker),
Forwarded(oneshot::Sender<Result<(), Error>>),
}
impl Observer {
fn complete(self, result: Result<(), Error>, deferred: &mut Deferred) {
match self {
Self::Local(waker) => deferred.wakes.push(waker),
Self::Forwarded(sender) => deferred.completions.push((sender, result)),
}
}
fn release(self, deferred: &mut Deferred) {
match self {
Self::Local(waker) => deferred.drops.push(waker),
Self::Forwarded(sender) => deferred.completions.push((sender, Err(Error::Closed))),
}
}
}
#[derive(Default)]
pub struct Timers {
entries: Slab<Observer>,
deadlines: BinaryHeap<Reverse<(Instant, TimerId)>>,
}
impl Timers {
fn insert(&mut self, deadline: Instant, waker: Waker) -> TimerId {
let id = TimerId(self.entries.insert(Observer::Local(waker)));
self.deadlines.push(Reverse((deadline, id)));
id
}
fn observe(&self, id: TimerId, waker: &Waker) -> Observation<()> {
match self.entries.get(id.0) {
None => Observation::Ready(()),
Some(Observer::Local(current)) if current.will_wake(waker) => Observation::Pending,
Some(Observer::Local(_)) => Observation::Refresh,
Some(Observer::Forwarded(_)) => panic!("timer has no local observer"),
}
}
fn refresh(&mut self, id: TimerId, waker: Waker) -> Result<Waker, Waker> {
match self.entries.get_mut(id.0) {
Some(Observer::Local(current)) => Ok(mem::replace(current, waker)),
None => Err(waker),
Some(Observer::Forwarded(_)) => panic!("timer has no local observer"),
}
}
pub fn forward(
&mut self,
id: TimerId,
sender: oneshot::Sender<Result<(), Error>>,
deferred: &mut Deferred,
) {
let Some(observer) = self.entries.get_mut(id.0) else {
deferred.completions.push((sender, Ok(())));
return;
};
let Observer::Local(waker) = mem::replace(observer, Observer::Forwarded(sender)) else {
panic!("timer forwarded twice");
};
deferred.drops.push(waker);
}
fn remove(&mut self, id: TimerId) -> Option<Observer> {
let observer = self.entries.remove(id.0)?;
self.compact();
Some(observer)
}
pub fn cancel(&mut self, id: TimerId, deferred: &mut Deferred) {
if let Some(observer) = self.remove(id) {
observer.release(deferred);
}
}
pub fn expire(&mut self, now: Instant, deferred: &mut Deferred) {
while let Some(deadline) = self.next_deadline() {
if deadline > now {
break;
}
let Reverse((_, id)) = self.deadlines.pop().unwrap();
self.entries
.remove(id.0)
.unwrap()
.complete(Ok(()), deferred);
}
self.compact();
}
pub fn next_deadline(&mut self) -> Option<Instant> {
while let Some(&Reverse((deadline, id))) = self.deadlines.peek() {
if self.entries.get(id.0).is_some() {
return Some(deadline);
}
self.deadlines.pop();
}
None
}
pub fn clear(&mut self, deferred: &mut Deferred) {
for index in 0..self.entries.slots() {
if let Some(id) = self.entries.id_at(index) {
self.entries
.remove(id)
.unwrap()
.complete(Err(Error::Closed), deferred);
}
}
self.deadlines.clear();
}
fn compact(&mut self) {
let stale = self.deadlines.len() - self.entries.len();
if stale <= 64 || stale <= self.entries.len() {
return;
}
let entries = &self.entries;
self.deadlines
.retain(|&Reverse((_, id))| entries.get(id.0).is_some());
self.deadlines.shrink_to_fit();
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{
Clock as _, Runner as _,
iouring::{
Runner,
operation::tests::{Reentrant, poll_on_foreign_thread},
},
utils::{extract_panic_message, reschedule},
};
use futures::{FutureExt as _, poll};
use std::{
panic::{AssertUnwindSafe, catch_unwind},
sync::atomic::{AtomicUsize, Ordering},
task::Wake,
thread,
};
impl Timers {
pub(crate) fn insert_due(&mut self, waker: Waker) -> TimerId {
self.insert(Instant::now(), waker)
}
}
#[derive(Default)]
struct Counter(AtomicUsize);
impl Wake for Counter {
fn wake(self: Arc<Self>) {
self.0.fetch_add(1, Ordering::Relaxed);
}
}
#[test]
fn test_elapsed_sleeps_are_ready_without_a_worker() {
let mut cx = Context::from_waker(Waker::noop());
for mut sleep in [
Sleep::new(Duration::ZERO),
Sleep::until(SystemTime::UNIX_EPOCH),
Sleep {
state: State::Unregistered {
deadline: Instant::now(),
},
},
] {
assert!(sleep.poll_unpin(&mut cx).is_ready());
assert!(matches!(sleep.state, State::Done));
}
}
#[test]
fn test_far_future_sleep_clamps_the_deadline() {
let before = Instant::now();
let sleep = Sleep::new(Duration::MAX);
let after = Instant::now();
let State::Unregistered { deadline } = sleep.state else {
panic!("positive sleep must retain a deadline");
};
assert!(deadline >= before + TimeoutWheel::MAX_TIMEOUT);
assert!(deadline <= after + TimeoutWheel::MAX_TIMEOUT);
}
#[test]
fn test_expiry_removes_due_timers_before_waking() {
let now = Instant::now();
let later = now + Duration::from_secs(1);
let counter = Arc::new(Counter::default());
let mut timers = Timers::default();
assert_eq!(timers.next_deadline(), None);
let future = timers.insert(later, Waker::from(counter.clone()));
let due = timers.insert(now, Waker::from(counter.clone()));
let also_due = timers.insert(now, Waker::from(counter.clone()));
let mut deferred = Deferred::default();
timers.expire(now, &mut deferred);
assert_eq!(deferred.wakes.len(), 2);
assert!(timers.remove(due).is_none());
assert!(timers.remove(also_due).is_none());
assert!(timers.entries.get(future.0).is_some());
assert_eq!(timers.next_deadline(), Some(later));
assert_eq!(counter.0.load(Ordering::Relaxed), 0);
timers.expire(now, &mut deferred);
assert_eq!(deferred.wakes.len(), 2);
for waker in deferred.wakes.drain(..) {
waker.wake();
}
assert_eq!(counter.0.load(Ordering::Relaxed), 2);
timers.expire(later, &mut deferred);
assert_eq!(timers.next_deadline(), None);
assert_eq!(timers.entries.len(), 0);
assert_eq!(deferred.wakes.len(), 1);
deferred.wakes.pop().unwrap().wake();
assert_eq!(counter.0.load(Ordering::Relaxed), 3);
}
#[test]
fn test_refresh_replaces_only_the_waker() {
let now = Instant::now();
let first = Arc::new(Counter::default());
let second = Arc::new(Counter::default());
let mut timers = Timers::default();
let id = timers.insert(now, Waker::from(first.clone()));
let replacement = Waker::from(second.clone());
assert!(matches!(
timers.observe(id, &replacement),
Observation::Refresh
));
let displaced = timers.refresh(id, replacement.clone()).unwrap();
assert_eq!(Arc::strong_count(&first), 2);
assert!(matches!(
timers.observe(id, &replacement),
Observation::Pending
));
assert_eq!(timers.entries.len(), 1);
assert_eq!(timers.deadlines.len(), 1);
assert_eq!(timers.next_deadline(), Some(now));
drop(displaced);
assert_eq!(Arc::strong_count(&first), 1);
let mut deferred = Deferred::default();
timers.expire(now, &mut deferred);
deferred.wakes.pop().unwrap().wake();
assert_eq!(first.0.load(Ordering::Relaxed), 0);
assert_eq!(second.0.load(Ordering::Relaxed), 1);
}
#[test]
fn test_stale_ids_and_deadlines_preserve_reused_slots() {
let now = Instant::now();
let later = now + Duration::from_secs(1);
let mut timers = Timers::default();
let old = timers.insert(now, Waker::noop().clone());
assert!(timers.remove(old).is_some());
let current_waker = Waker::from(Arc::new(Counter::default()));
let current = timers.insert(later, current_waker.clone());
assert_eq!(old.0.index, current.0.index);
assert_ne!(old.0.generation, current.0.generation);
assert!(timers.remove(old).is_none());
assert!(matches!(
timers.observe(old, ¤t_waker),
Observation::Ready(())
));
let waker = Waker::from(Arc::new(Counter::default()));
let rejected = timers.refresh(old, waker.clone()).unwrap_err();
assert!(rejected.will_wake(&waker));
assert!(matches!(
timers.observe(current, ¤t_waker),
Observation::Pending
));
assert_eq!(timers.next_deadline(), Some(later));
let mut deferred = Deferred::default();
timers.expire(now, &mut deferred);
assert!(deferred.wakes.is_empty());
assert!(timers.entries.get(current.0).is_some());
}
#[test]
fn test_compaction_bounds_stale_deadlines() {
let now = Instant::now();
let later = now + Duration::from_secs(1);
let mut timers = Timers::default();
let oldest = timers.insert(now, Waker::noop().clone());
for _ in 0..1000 {
let id = timers.insert(later, Waker::noop().clone());
assert!(timers.remove(id).is_some());
assert!(timers.deadlines.len() <= timers.entries.len() + 64);
}
assert_eq!(timers.entries.slots(), 2);
assert_eq!(timers.next_deadline(), Some(now));
assert!(timers.remove(oldest).is_some());
assert_eq!(timers.next_deadline(), None);
for _ in 0..128 {
timers.insert(now, Waker::noop().clone());
}
let future = timers.insert(later, Waker::noop().clone());
for _ in 0..65 {
let id = timers.insert(later + Duration::from_secs(1), Waker::noop().clone());
assert!(timers.remove(id).is_some());
}
assert_eq!(timers.deadlines.len(), 128 + 1 + 65);
let mut deferred = Deferred::default();
timers.expire(now, &mut deferred);
assert_eq!(deferred.wakes.len(), 128);
assert_eq!(timers.deadlines.len(), 1);
assert!(timers.entries.get(future.0).is_some());
}
#[test]
fn test_clear_defers_wakes_and_preserves_reuse() {
let now = Instant::now();
let later = now + Duration::from_secs(1);
let counter = Arc::new(Counter::default());
let mut timers = Timers::default();
let cancelled = timers.insert(now, Waker::noop().clone());
let id = timers.insert(later, Waker::from(counter.clone()));
assert!(timers.remove(cancelled).is_some());
let mut deferred = Deferred::default();
deferred.wakes.push(Waker::noop().clone());
timers.clear(&mut deferred);
timers.clear(&mut deferred);
assert_eq!(deferred.wakes.len(), 2);
assert!(deferred.drops.is_empty());
assert!(timers.deadlines.is_empty());
assert_eq!(timers.next_deadline(), None);
assert_eq!(timers.entries.len(), 0);
assert_eq!(Arc::strong_count(&counter), 2);
assert_eq!(counter.0.load(Ordering::Relaxed), 0);
for waker in deferred.wakes.drain(..) {
waker.wake();
}
assert_eq!(counter.0.load(Ordering::Relaxed), 1);
assert_eq!(Arc::strong_count(&counter), 1);
let reused = timers.insert(now, Waker::noop().clone());
assert_eq!(reused.0.index, id.0.index);
assert!(timers.remove(id).is_none());
assert_eq!(timers.next_deadline(), Some(now));
}
#[test]
fn test_waker_refresh_and_clone_panic_preserve_cancellation() {
Runner::default().start(|_| async {
let mut sleep = Sleep::new(Duration::from_secs(60));
let first = Arc::new(Reentrant::default());
let first_waker = first.waker();
assert!(
sleep
.poll_unpin(&mut Context::from_waker(&first_waker))
.is_pending()
);
let State::Registered(registration) = &sleep.state else {
panic!("sleep did not register");
};
let timer_id = registration.key();
let owner = Local::current().unwrap();
let second = Arc::new(Reentrant::default());
let second_waker = second.waker();
assert!(
sleep
.poll_unpin(&mut Context::from_waker(&second_waker))
.is_pending()
);
assert!(matches!(
owner.borrow().timers.observe(timer_id, &second_waker),
Observation::Pending
));
assert_eq!(first.drops.load(Ordering::Relaxed), 0);
reschedule().await;
assert_eq!(first.drops.load(Ordering::Relaxed), 1);
first
.panic_callback
.store(Reentrant::CLONE, Ordering::Relaxed);
let panic = catch_unwind(AssertUnwindSafe(|| {
let _ = sleep.poll_unpin(&mut Context::from_waker(&first_waker));
}))
.expect_err("changed sleep observer must clone");
assert_eq!(extract_panic_message(&*panic), "waker callback panic 1");
assert!(matches!(
owner.borrow().timers.observe(timer_id, &second_waker),
Observation::Pending
));
drop(sleep);
assert!(owner.borrow().timers.entries.get(timer_id.0).is_none());
assert_eq!(second.drops.load(Ordering::Relaxed), 0);
reschedule().await;
assert_eq!(second.drops.load(Ordering::Relaxed), 1);
});
}
#[test]
fn test_foreign_drop_cancels_registration() {
for poll_first in [false, true] {
Runner::default().start(|_| async {
let mut sleep = Sleep::new(Duration::from_secs(60));
assert!(poll!(&mut sleep).is_pending());
let State::Registered(registration) = &sleep.state else {
panic!("sleep did not register");
};
let timer_id = registration.key();
let owner = Local::current().unwrap();
thread::spawn(move || {
if poll_first {
assert!(
sleep
.poll_unpin(&mut Context::from_waker(Waker::noop()))
.is_pending()
);
}
drop(sleep);
})
.join()
.unwrap();
assert!(owner.borrow().timers.entries.get(timer_id.0).is_some());
reschedule().await;
assert!(owner.borrow().timers.entries.get(timer_id.0).is_none());
});
}
}
#[test]
fn test_sleep_moves_between_threads() {
for expired in [false, true] {
for return_to_owner in [false, true] {
Runner::default().start(|context| async move {
let mut sleep = Sleep::new(Duration::from_millis(20));
let State::Unregistered { deadline } = sleep.state else {
unreachable!()
};
assert!(poll!(&mut sleep).is_pending());
if expired {
context.sleep(Duration::from_millis(30)).await;
}
let sleep = poll_on_foreign_thread(sleep, Waker::noop().clone());
let sleep = poll_on_foreign_thread(sleep, Waker::noop().clone());
if return_to_owner {
sleep.await;
} else {
let (sender, receiver) = oneshot::channel();
let thread = thread::spawn(move || {
futures::executor::block_on(sleep);
sender.send(()).unwrap();
});
receiver.await.unwrap();
thread.join().unwrap();
}
assert!(Instant::now() >= deadline);
});
}
}
}
#[test]
fn test_forwarding_and_cancellation_preserve_reused_timer() {
for promoted in [false, true] {
Runner::default().start(|_| async {
let mut sleep = Sleep::new(Duration::from_secs(60));
assert!(poll!(&mut sleep).is_pending());
let State::Registered(registration) = &sleep.state else {
unreachable!()
};
let old = registration.key();
let sleep = poll_on_foreign_thread(sleep, Waker::noop().clone());
if promoted {
reschedule().await;
}
drop(sleep);
let mut replacement = Sleep::new(Duration::from_secs(60));
assert!(poll!(&mut replacement).is_pending());
let State::Registered(registration) = &replacement.state else {
unreachable!()
};
let current = registration.key();
assert_eq!(old.0.index, current.0.index);
assert_ne!(old.0.generation, current.0.generation);
reschedule().await;
assert!(poll!(&mut replacement).is_pending());
let owner = Local::current().unwrap();
assert!(owner.borrow().timers.entries.get(old.0).is_none());
assert!(owner.borrow().timers.entries.get(current.0).is_some());
drop(replacement);
assert!(owner.borrow().timers.entries.get(current.0).is_none());
});
}
}
#[test]
fn test_forwarded_timer_closure_and_callbacks() {
for promoted in [false, true] {
let callbacks = Arc::new(Reentrant::default());
let (mut sleep,) = Runner::default().start(|_| async {
let mut sleep = Sleep::new(Duration::from_secs(60));
assert!(
sleep
.poll_unpin(&mut Context::from_waker(&callbacks.waker()))
.is_pending()
);
let sleep = poll_on_foreign_thread(sleep, callbacks.waker());
if promoted {
reschedule().await;
}
(sleep,)
});
if promoted {
assert!(callbacks.wakes.load(Ordering::Relaxed) > 0);
}
let result = catch_unwind(AssertUnwindSafe(|| {
sleep.poll_unpin(&mut Context::from_waker(Waker::noop()))
}));
assert!(
extract_panic_message(&*result.unwrap_err())
.starts_with("io_uring sleep polled after its worker closed")
);
drop(sleep);
assert_eq!(Arc::strong_count(&callbacks), 1);
}
}
#[test]
fn test_worker_closure_releases_escaped_sleep() {
let counter = Arc::new(Counter::default());
let (mut sleep,) = Runner::default().start(|_| async {
let mut sleep = Sleep::new(Duration::from_secs(60));
let waker = Waker::from(counter.clone());
assert!(
sleep
.poll_unpin(&mut Context::from_waker(&waker))
.is_pending()
);
(sleep,)
});
assert_eq!(Arc::strong_count(&counter), 1);
assert_eq!(counter.0.load(Ordering::Relaxed), 1);
let panic = catch_unwind(AssertUnwindSafe(|| {
sleep.poll_unpin(&mut Context::from_waker(Waker::noop()))
}))
.expect_err("registered sleep must reject polling after closure");
assert!(
extract_panic_message(&*panic)
.starts_with("io_uring sleep polled after its worker closed")
);
drop(sleep);
}
}