subms_timer_wheel/features/

deadline_scheduler.rs

//! Absolute-deadline scheduling layer on top of the base wheel.
//! Callers schedule against wall-clock instants ("fire at t=...") and
//! drive the scheduler with `poll()` calls. The wheel itself stays
//! tick-counted; the layer translates between instant deltas and tick
//! deltas via an injected `Clock` so the workload is deterministic
//! under test.
//!
//! The clock abstraction (a trait, not a free-running `Instant`) is
//! deliberate: time-based tests that sleep are flaky and slow; tests
//! that mutate a `TestClock` finish in microseconds and are exact.
//!
//! Granularity is 1 ms per tick by default; a deadline of `now + 12
//! ms` lands twelve ticks out. Sub-ms deadlines round up to one tick.

use crate::TimerWheel;
use std::sync::OnceLock;
use std::time::{Duration, Instant};

/// Source of monotonic time. The deadline scheduler measures "from
/// now" deltas off this; production code injects [`MonotonicClock`]
/// and tests inject [`TestClock`].
pub trait Clock {
    /// Elapsed monotonic nanoseconds since the clock's origin. The
    /// origin doesn't matter; only deltas do.
    fn now_nanos(&self) -> u64;
}

#[derive(Default)]
pub struct MonotonicClock {
    /// Set once on first read (or eagerly by `new`). Persisting the
    /// origin is what keeps successive `now_nanos` reads monotonic;
    /// re-sampling `Instant::now` as the origin each call would make
    /// every read a fresh near-zero delta off an unrelated baseline.
    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
    }
}

/// Hand-stepped clock for deterministic tests. `advance(d)` moves
/// time forward by `d`; the scheduler then catches up via `poll()`.
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,
    /// Nanos consumed by previous ticks. Lets `poll()` advance
    /// `(elapsed - consumed) / tick_nanos` ticks atomically.
    consumed_nanos: u64,
}

impl<V, C: Clock> DeadlineScheduler<V, C> {
    /// Build a deadline scheduler with `num_slots` wheel slots and
    /// `tick` resolution (rounded up to 1 ns minimum).
    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
    }

    /// The injected clock. `TestClock::advance` takes `&self`, so a test or
    /// a demo can step time through this without owning the clock twice.
    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()
    }

    /// Schedule `value` to fire after `delay`. Equivalent to
    /// `schedule_at(now + delay, value)`.
    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)
    }

    /// Schedule `value` to fire at absolute deadline `when_nanos`
    /// (same epoch as `Clock::now_nanos`). If the deadline is in the
    /// past, the timer is queued for the next tick.
    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)
    }

    /// Push a pending timer out to a new deadline, keeping its id. This is
    /// the idle-timeout pattern: one timer per session, bumped on every
    /// inbound message rather than cancelled and re-armed.
    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)
    }

    /// Hand back every armed timer without firing it. The shutdown path.
    pub fn drain(&mut self) -> Vec<V> {
        self.wheel.drain()
    }

    /// Advance the wheel by however many ticks the clock has accrued
    /// since the last `poll`. Returns every fired value across the
    /// catch-up batch. Idempotent if called twice with no clock
    /// movement in between.
    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;