use std::time::Duration;
use tokio::sync::watch;
use crate::bus::LogicalTime;
use crate::participant::spec::MissedTick;
pub(crate) fn duration_to_nanos_saturating(duration: Duration) -> u64 {
u64::try_from(duration.as_nanos()).unwrap_or(u64::MAX)
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct SchedulerTick {
pub fired_at: LogicalTime,
pub missed_ticks: u32,
}
#[allow(async_fn_in_trait)]
pub trait StepScheduler: Send + Sync + 'static {
async fn wait_until(&self, target: LogicalTime) -> SchedulerTick;
fn now(&self) -> LogicalTime;
}
pub struct RealScheduler {
missed_tick: MissedTick,
period: Option<Duration>,
epoch: u64,
started_at: tokio::time::Instant,
started_ns: u64,
}
impl RealScheduler {
pub fn new(missed_tick: MissedTick, period: Option<Duration>, now: LogicalTime) -> Self {
RealScheduler {
missed_tick,
period,
epoch: now.epoch(),
started_at: tokio::time::Instant::now(),
started_ns: now.time_ns(),
}
}
fn instant_for(&self, target: LogicalTime) -> tokio::time::Instant {
let delta_ns = target.time_ns().saturating_sub(self.started_ns);
self.started_at + Duration::from_nanos(delta_ns)
}
fn logical_for(&self, instant: tokio::time::Instant) -> LogicalTime {
let elapsed = instant.saturating_duration_since(self.started_at);
LogicalTime::new(
self.epoch,
self.started_ns
.saturating_add(duration_to_nanos_saturating(elapsed)),
)
}
}
impl StepScheduler for RealScheduler {
async fn wait_until(&self, target: LogicalTime) -> SchedulerTick {
let deadline = self.instant_for(target);
tokio::time::sleep_until(deadline).await;
let mut missed_ticks = 0u32;
if self.missed_tick == MissedTick::Collapse {
if let Some(period) = self.period.filter(|period| !period.is_zero()) {
let real_now = tokio::time::Instant::now();
let mut next_deadline = deadline;
while next_deadline + period <= real_now {
next_deadline += period;
missed_ticks = missed_ticks.saturating_add(1);
}
}
}
SchedulerTick {
fired_at: self.logical_for(tokio::time::Instant::now()),
missed_ticks,
}
}
fn now(&self) -> LogicalTime {
self.logical_for(tokio::time::Instant::now())
}
}
pub struct SimulationScheduler {
missed_tick: MissedTick,
period: Option<Duration>,
_tx_keepalive: watch::Sender<LogicalTime>,
rx: watch::Receiver<LogicalTime>,
paused: std::sync::Arc<std::sync::atomic::AtomicBool>,
}
#[derive(Clone)]
pub struct SimulationClockHandle {
tx: watch::Sender<LogicalTime>,
paused: std::sync::Arc<std::sync::atomic::AtomicBool>,
}
impl SimulationClockHandle {
pub fn advance(&self, time: LogicalTime) {
self.tx.send_if_modified(|current| {
if time > *current {
*current = time;
true
} else {
false
}
});
}
pub fn pause(&self) {
use std::sync::atomic::Ordering;
self.paused.store(true, Ordering::SeqCst);
self.tx.send_modify(|_| {});
}
pub fn resume(&self) {
use std::sync::atomic::Ordering;
self.paused.store(false, Ordering::SeqCst);
self.tx.send_modify(|_| {});
}
}
impl SimulationScheduler {
pub fn new(
missed_tick: MissedTick,
period: Option<Duration>,
start: LogicalTime,
) -> (Self, SimulationClockHandle) {
let (tx, rx) = watch::channel(start);
let paused = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
let scheduler = SimulationScheduler {
missed_tick,
period,
_tx_keepalive: tx.clone(),
rx,
paused: std::sync::Arc::clone(&paused),
};
let handle = SimulationClockHandle { tx, paused };
(scheduler, handle)
}
fn is_paused(&self) -> bool {
self.paused.load(std::sync::atomic::Ordering::SeqCst)
}
fn missed_ticks(&self, target: LogicalTime, current: LogicalTime) -> u32 {
let (MissedTick::Collapse, Some(period)) = (self.missed_tick, self.period) else {
return 0;
};
if current <= target {
return 0;
}
let period_ns = duration_to_nanos_saturating(period);
if period_ns == 0 {
return 0;
}
let overrun_ns = current.time_ns().saturating_sub(target.time_ns());
u32::try_from(overrun_ns / period_ns).unwrap_or(u32::MAX)
}
}
impl StepScheduler for SimulationScheduler {
async fn wait_until(&self, target: LogicalTime) -> SchedulerTick {
let mut rx = self.rx.clone();
loop {
let current = *rx.borrow_and_update();
if current >= target && !self.is_paused() {
let missed_ticks = self.missed_ticks(target, current);
return SchedulerTick {
fired_at: current,
missed_ticks,
};
}
if rx.changed().await.is_err() {
return SchedulerTick {
fired_at: *rx.borrow_and_update(),
missed_ticks: 0,
};
}
}
}
fn now(&self) -> LogicalTime {
*self.rx.borrow()
}
}
pub enum AnyStepScheduler {
Real(RealScheduler),
Simulation(SimulationScheduler),
}
impl StepScheduler for AnyStepScheduler {
async fn wait_until(&self, target: LogicalTime) -> SchedulerTick {
match self {
AnyStepScheduler::Real(scheduler) => scheduler.wait_until(target).await,
AnyStepScheduler::Simulation(scheduler) => scheduler.wait_until(target).await,
}
}
fn now(&self) -> LogicalTime {
match self {
AnyStepScheduler::Real(scheduler) => scheduler.now(),
AnyStepScheduler::Simulation(scheduler) => scheduler.now(),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn lt(ns: u64) -> LogicalTime {
LogicalTime::new(0, ns)
}
const PERIOD: Duration = Duration::from_millis(10);
const SIM_PERIOD: Duration = Duration::from_nanos(10);
#[tokio::test(start_paused = true)]
async fn real_scheduler_wakes_at_target_and_reports_no_miss_when_on_time() {
let start = tokio::time::Instant::now();
let scheduler = RealScheduler::new(MissedTick::Collapse, Some(PERIOD), lt(0));
let period_ns = PERIOD.as_nanos() as u64;
let tick = scheduler.wait_until(lt(period_ns)).await;
assert_eq!(tick.missed_ticks, 0);
assert!(tick.fired_at.time_ns() >= period_ns);
assert!(start.elapsed() >= PERIOD);
}
#[tokio::test(start_paused = true)]
async fn real_scheduler_collapses_a_missed_tick_instead_of_bursting() {
let scheduler = RealScheduler::new(MissedTick::Collapse, Some(PERIOD), lt(0));
tokio::time::advance(Duration::from_millis(50)).await;
let period_ns = PERIOD.as_nanos() as u64;
let tick = scheduler.wait_until(lt(period_ns)).await;
assert_eq!(
tick.missed_ticks, 4,
"a multi-period overrun collapses to one tick, reporting the skipped count"
);
}
#[tokio::test]
async fn simulation_scheduler_releases_ticks_in_order_deterministically() {
let (scheduler, handle) =
SimulationScheduler::new(MissedTick::Collapse, Some(SIM_PERIOD), lt(0));
let mut fired = Vec::new();
for step in 1..=5u64 {
let target = lt(step * 10);
let handle = handle.clone();
let advance_target = target;
let advancer = tokio::spawn(async move { handle.advance(advance_target) });
let tick = scheduler.wait_until(target).await;
advancer.await.unwrap();
fired.push(tick.fired_at.time_ns());
}
assert_eq!(fired, vec![10, 20, 30, 40, 50]);
}
#[tokio::test]
async fn simulation_scheduler_never_sleeps_on_a_wall_clock_timer() {
let (scheduler, handle) =
SimulationScheduler::new(MissedTick::Collapse, Some(SIM_PERIOD), lt(0));
handle.advance(lt(100));
let started = std::time::Instant::now();
let tick = scheduler.wait_until(lt(100)).await;
let elapsed = started.elapsed();
assert_eq!(tick.fired_at, lt(100));
assert!(
elapsed < Duration::from_millis(50),
"simulation scheduler should resolve immediately once time is already due, took {elapsed:?}"
);
}
#[tokio::test]
async fn simulation_scheduler_releases_when_advance_happened_before_wait() {
let (scheduler, handle) =
SimulationScheduler::new(MissedTick::Collapse, Some(SIM_PERIOD), lt(0));
handle.advance(lt(30));
let tick = tokio::time::timeout(Duration::from_millis(50), scheduler.wait_until(lt(20)))
.await
.expect("advance-before-wait should release without waiting for another change");
assert_eq!(tick.fired_at, lt(30));
assert_eq!(tick.missed_ticks, 1);
}
#[tokio::test]
async fn simulation_scheduler_does_not_miss_racing_advance_after_pending_poll() {
let (scheduler, handle) =
SimulationScheduler::new(MissedTick::Collapse, Some(SIM_PERIOD), lt(0));
let wait = scheduler.wait_until(lt(10));
tokio::pin!(wait);
assert!(
poll_once(wait.as_mut()).is_none(),
"wait should pend before logical time reaches the target"
);
handle.advance(lt(10));
let tick = tokio::time::timeout(Duration::from_secs(1), &mut wait)
.await
.expect("advance after a pending poll should wake the waiter");
assert_eq!(tick.fired_at, lt(10));
assert_eq!(tick.missed_ticks, 0);
}
#[tokio::test]
async fn simulation_scheduler_keeps_waiting_if_external_handle_is_dropped() {
let (scheduler, handle) =
SimulationScheduler::new(MissedTick::Collapse, Some(SIM_PERIOD), lt(0));
drop(handle);
let wait = scheduler.wait_until(lt(10));
tokio::pin!(wait);
assert!(
poll_once(wait.as_mut()).is_none(),
"dropping the external handle must not close the scheduler feed and release a stale tick"
);
}
#[tokio::test]
async fn simulation_scheduler_collapses_a_multi_period_jump() {
let (scheduler, handle) =
SimulationScheduler::new(MissedTick::Collapse, Some(SIM_PERIOD), lt(0));
handle.advance(lt(40));
let tick = scheduler.wait_until(lt(10)).await;
assert_eq!(
tick.missed_ticks, 3,
"a jump past the target collapses to one released tick, reporting all 3 skipped periods"
);
assert_eq!(tick.fired_at, lt(40));
}
#[tokio::test(start_paused = true)]
async fn simulation_scheduler_pause_blocks_release_until_resumed() {
let (scheduler, handle) =
SimulationScheduler::new(MissedTick::Collapse, Some(SIM_PERIOD), lt(0));
handle.pause();
handle.advance(lt(10));
let wait = scheduler.wait_until(lt(10));
tokio::pin!(wait);
assert_eq!(
scheduler.now(),
lt(10),
"paused scheduler still reports the latest observed logical time"
);
assert!(
poll_once(wait.as_mut()).is_none(),
"paused scheduler must not release a due tick"
);
handle.resume();
let tick = tokio::time::timeout(Duration::from_secs(1), &mut wait)
.await
.expect("resumed scheduler should release the due tick promptly");
assert_eq!(tick.fired_at, lt(10));
}
fn poll_once<F: Future>(fut: std::pin::Pin<&mut F>) -> Option<F::Output> {
use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};
fn noop(_: *const ()) {}
fn clone(_: *const ()) -> RawWaker {
RawWaker::new(std::ptr::null(), &VTABLE)
}
static VTABLE: RawWakerVTable = RawWakerVTable::new(clone, noop, noop, noop);
let waker = unsafe { Waker::from_raw(RawWaker::new(std::ptr::null(), &VTABLE)) };
let mut cx = Context::from_waker(&waker);
match fut.poll(&mut cx) {
Poll::Ready(output) => Some(output),
Poll::Pending => None,
}
}
}