use core::cmp::Ordering;
use std::collections::BinaryHeap;
use std::time::{Duration, Instant};
pub struct SleepTracker<T> {
heap: BinaryHeap<Sleeper<T>>,
}
struct Sleeper<T> {
wakeup: Instant,
data: T,
}
impl<T> PartialEq for Sleeper<T> {
fn eq(&self, other: &Sleeper<T>) -> bool {
self.wakeup == other.wakeup
}
}
impl<T> PartialOrd for Sleeper<T> {
fn partial_cmp(&self, other: &Sleeper<T>) -> Option<Ordering> {
Some(other.wakeup.cmp(&self.wakeup))
}
}
impl<T> Eq for Sleeper<T> {}
impl<T> Ord for Sleeper<T> {
fn cmp(&self, other: &Sleeper<T>) -> Ordering {
self.wakeup.cmp(&other.wakeup)
}
}
impl<T> SleepTracker<T> {
pub fn new() -> SleepTracker<T> {
SleepTracker {
heap: BinaryHeap::new(),
}
}
pub fn push(&mut self, sleep: u64, data: T) {
self.heap.push(Sleeper {
wakeup: Instant::now()
.checked_add(Duration::from_millis(sleep))
.expect("instant should not wrap"),
data,
});
}
pub fn len(&self) -> usize {
self.heap.len()
}
pub fn to_retry(&mut self) -> Vec<T> {
let now = Instant::now();
let mut result = Vec::new();
while let Some(next) = self.heap.peek() {
log::debug!("ERIC: now={now:?} next={:?}", next.wakeup);
if next.wakeup < now {
result.push(self.heap.pop().unwrap().data);
} else {
break;
}
}
result
}
pub fn time_to_next(&self) -> Option<Duration> {
self.heap
.peek()
.map(|s| s.wakeup.saturating_duration_since(Instant::now()))
}
}
#[test]
fn returns_in_order() {
let mut s = SleepTracker::new();
s.push(3, 3);
s.push(1, 1);
s.push(6, 6);
s.push(5, 5);
s.push(2, 2);
s.push(10000, 10000);
assert_eq!(s.len(), 6);
std::thread::sleep(Duration::from_millis(100));
assert_eq!(s.to_retry(), &[1, 2, 3, 5, 6]);
}