use crate::TimerWheel;
use std::sync::OnceLock;
use std::time::{Duration, Instant};
pub trait Clock {
fn now_nanos(&self) -> u64;
}
#[derive(Default)]
pub struct MonotonicClock {
origin: OnceLock<Instant>,
}
impl MonotonicClock {
pub fn new() -> Self {
let origin = OnceLock::new();
let _ = origin.set(Instant::now());
Self { origin }
}
}
impl Clock for MonotonicClock {
fn now_nanos(&self) -> u64 {
let origin = self.origin.get_or_init(Instant::now);
Instant::now().duration_since(*origin).as_nanos() as u64
}
}
pub struct TestClock {
now_nanos: std::cell::Cell<u64>,
}
impl Default for TestClock {
fn default() -> Self {
Self::new()
}
}
impl TestClock {
pub fn new() -> Self {
Self {
now_nanos: std::cell::Cell::new(0),
}
}
pub fn advance(&self, d: Duration) {
self.now_nanos
.set(self.now_nanos.get().saturating_add(d.as_nanos() as u64));
}
}
impl Clock for TestClock {
fn now_nanos(&self) -> u64 {
self.now_nanos.get()
}
}
pub struct DeadlineScheduler<V, C: Clock> {
wheel: TimerWheel<V>,
clock: C,
tick_nanos: u64,
consumed_nanos: u64,
}
impl<V, C: Clock> DeadlineScheduler<V, C> {
pub fn new(num_slots: usize, clock: C, tick: Duration) -> Self {
let tick_nanos = (tick.as_nanos() as u64).max(1);
Self {
wheel: TimerWheel::new(num_slots),
clock,
tick_nanos,
consumed_nanos: 0,
}
}
pub fn tick_nanos(&self) -> u64 {
self.tick_nanos
}
pub fn clock(&self) -> &C {
&self.clock
}
pub fn pending(&self) -> usize {
self.wheel.pending()
}
pub fn is_empty(&self) -> bool {
self.wheel.is_empty()
}
pub fn schedule_after(&mut self, delay: Duration, value: V) -> u64 {
let ticks = self.nanos_to_ticks(delay.as_nanos() as u64);
self.wheel.schedule(ticks, value)
}
pub fn schedule_at(&mut self, when_nanos: u64, value: V) -> u64 {
let now = self.clock.now_nanos();
let diff = when_nanos.saturating_sub(now);
let ticks = self.nanos_to_ticks(diff).max(1);
self.wheel.schedule(ticks, value)
}
pub fn cancel(&mut self, id: u64) -> bool {
self.wheel.cancel(id)
}
pub fn reschedule_at(&mut self, id: u64, when_nanos: u64) -> bool {
let now = self.clock.now_nanos();
let diff = when_nanos.saturating_sub(now);
let ticks = self.nanos_to_ticks(diff).max(1);
self.wheel.reschedule(id, ticks)
}
pub fn reschedule_after(&mut self, id: u64, delay: Duration) -> bool {
let ticks = self.nanos_to_ticks(delay.as_nanos() as u64).max(1);
self.wheel.reschedule(id, ticks)
}
pub fn drain(&mut self) -> Vec<V> {
self.wheel.drain()
}
pub fn poll(&mut self) -> Vec<V> {
let now = self.clock.now_nanos();
let pending = now.saturating_sub(self.consumed_nanos);
let ticks = (pending / self.tick_nanos) as usize;
self.consumed_nanos = self
.consumed_nanos
.saturating_add(ticks as u64 * self.tick_nanos);
let mut fired = Vec::new();
for _ in 0..ticks {
fired.extend(self.wheel.tick());
}
fired
}
fn nanos_to_ticks(&self, nanos: u64) -> usize {
nanos.div_ceil(self.tick_nanos) as usize
}
}
#[cfg(test)]
#[path = "deadline_scheduler_tests.rs"]
mod tests;