use std::cmp::Reverse;
use std::collections::{BinaryHeap, HashMap};
use std::{task::Waker, time::Instant};
use dambi::LocalCell;
pub(crate) struct TimerHeap {
inner: LocalCell<TimerHeapInner>,
}
struct TimerHeapInner {
next_id: u64,
heap: BinaryHeap<Reverse<(Instant, u64)>>,
entries: HashMap<u64, TimerEntry>,
}
struct TimerEntry {
deadline: Instant,
waker: Option<Waker>,
}
impl TimerHeap {
pub(crate) fn new() -> Self {
Self {
inner: LocalCell::new(TimerHeapInner {
next_id: 1,
heap: BinaryHeap::new(),
entries: HashMap::new(),
}),
}
}
pub(crate) fn register(&self, id: Option<u64>, deadline: Instant, waker: &Waker) -> u64 {
let inner = self.inner.get_mut();
if let Some(id) = id
&& let Some(ent) = inner.entries.get_mut(&id)
{
ent.deadline = deadline;
ent.waker = Some(waker.clone());
inner.heap.push(Reverse((deadline, id)));
return id;
}
let id = inner.next_id;
inner.next_id = inner.next_id.wrapping_add(1);
inner.entries.insert(
id,
TimerEntry {
deadline,
waker: Some(waker.clone()),
},
);
inner.heap.push(Reverse((deadline, id)));
id
}
pub(crate) fn cancel(&self, id: u64) {
let inner = self.inner.get_mut();
let _ = inner.entries.remove(&id);
}
pub(crate) fn wake_due(&self) {
let inner = self.inner.get_mut();
if inner.heap.is_empty() {
return;
}
let now = Instant::now();
while let Some(Reverse((deadline, id))) = inner.heap.peek().copied() {
if deadline > now {
break;
}
let _ = inner.heap.pop();
let Some(ent) = inner.entries.get(&id) else {
continue;
};
if ent.deadline != deadline {
continue;
}
let ent = inner.entries.remove(&id);
if let Some(ent) = ent
&& let Some(w) = ent.waker
{
w.wake();
}
}
}
pub(crate) fn next_deadline(&self) -> Option<Instant> {
let inner = self.inner.get_mut();
loop {
let Reverse((deadline, id)) = inner.heap.peek().copied()?;
let Some(ent) = inner.entries.get(&id) else {
let _ = inner.heap.pop();
continue;
};
if ent.deadline == deadline {
return Some(deadline);
}
let _ = inner.heap.pop();
}
}
}