darkbio-clock 0.2.0

// clock-rs: virtual clock for testing blocking code
// Copyright 2026 Dark Bio AG. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

//! Test clocks, which move only when their owner advances them.

use crate::primitives::{Condvar, Mutex, MutexGuard};
use crate::{Clock, Signal};
#[cfg(feature = "crossbeam")]
use crossbeam_channel::{Receiver, RecvTimeoutError, Sender, TryRecvError};
use std::collections::BTreeMap;
use std::fmt;
use std::sync::{Arc, PoisonError, Weak};
use std::time::{Duration, Instant, SystemTime};

/// Owns a clock that moves only when advanced, for tests.
///
/// [`Self::clock`] hands out handles that read and sleep on its time. Only the
/// owner moves it, through methods that take `&mut self`, so each test clock
/// has one driver.
#[cfg_attr(docsrs, doc(cfg(feature = "test-clock")))]
pub struct TestClock {
    /// State shared with every clock handle.
    paused: Arc<Paused>,
}

impl TestClock {
    /// Creates a stopped clock at the current real monotonic and wall times.
    pub fn new() -> Self {
        let now = Instant::now();
        Self {
            paused: Arc::new(Paused {
                start: now,
                state: Mutex::new(PausedState {
                    now,
                    system_time: SystemTime::now(),
                    signals: BTreeMap::new(),
                    blocked: 0,
                    deadlines: BTreeMap::new(),
                    #[cfg(feature = "crossbeam")]
                    timers: BTreeMap::new(),
                    #[cfg(feature = "crossbeam")]
                    fired: Vec::new(),
                }),
                changed: Condvar::new(),
                #[cfg(test)]
                before_park: std::sync::Mutex::new(None),
                #[cfg(test)]
                before_rewait: std::sync::Mutex::new(None),
                #[cfg(all(test, feature = "crossbeam"))]
                after_timer_send: std::sync::Mutex::new(None),
                #[cfg(all(test, feature = "crossbeam"))]
                after_timer_receive: std::sync::Mutex::new(None),
            }),
        }
    }

    /// Returns a handle that reads and sleeps on this clock.
    pub fn clock(&self) -> Clock {
        Clock {
            paused: Some(self.paused.clone()),
        }
    }

    /// Moves both times forward by `by`, waking the sleeps and deadline waits
    /// it reaches and firing its due timers.
    ///
    /// Returns once the waits are notified and the timers hold their messages,
    /// without waiting for any thread to act. A zero advance does nothing.
    ///
    /// # Panics
    ///
    /// Panics if either time would overflow, before changing either one.
    pub fn advance(&mut self, by: Duration) {
        self.advance_with(|now| {
            now.checked_add(by)
                .expect("clock advance overflows Instant")
        });
    }

    /// Moves monotonic time to `target` and wall time by the same amount,
    /// waking the sleeps and deadline waits it reaches and firing its due timers.
    ///
    /// Returns once the waits are notified and the timers hold their messages,
    /// without waiting for any thread to act. Advancing to the current time
    /// does nothing.
    ///
    /// # Panics
    ///
    /// Panics if `target` is before now or wall time would overflow.
    /// Neither time changes after a panic.
    pub fn advance_to(&mut self, target: Instant) {
        self.advance_with(|now| {
            assert!(target >= now, "clock cannot go backwards");
            target
        });
    }

    /// Sets wall time forwards or backwards without moving monotonic time.
    ///
    /// This wakes no waits and fires no timers, since deadlines use monotonic time.
    pub fn set_system_time(&mut self, time: SystemTime) {
        self.paused.lock().system_time = time;
    }

    /// Blocks until at least `count` threads are parked in this clock's sleeps
    /// and condvar waits.
    ///
    /// Threads blocked in crossbeam receives and selects do not count. A thread
    /// parked earlier counts too, so the count proves no progress on its own.
    /// Nothing bounds the wait, so the test runner ends a hang.
    pub fn wait_blocked(&self, count: usize) {
        let mut state = self.paused.lock();
        while state.blocked < count {
            state = self
                .paused
                .changed
                .wait(state)
                .unwrap_or_else(PoisonError::into_inner);
        }
    }

    /// Blocks until at least `count` timers are armed on this clock and unfired.
    ///
    /// A waiting receive's timer counts, and so does a timer whose receiver was
    /// dropped. A timer armed earlier counts too, so the count proves no
    /// progress on its own. Nothing bounds the wait, so the test runner ends a
    /// hang.
    #[cfg(feature = "crossbeam")]
    #[cfg_attr(docsrs, doc(cfg(all(feature = "test-clock", feature = "crossbeam"))))]
    pub fn wait_timers(&self, count: usize) {
        let mut state = self.paused.lock();
        while state.timers.len() < count {
            state = self
                .paused
                .changed
                .wait(state)
                .unwrap_or_else(PoisonError::into_inner);
        }
    }

    /// Returns the earliest deadline among this clock's parked sleeps, deadline
    /// waits and unfired timers, or `None` when there is none.
    ///
    /// A wait is listed while parked, and one woken by an advance stays listed
    /// until its thread runs, so await an advance's effect before reading the
    /// next deadline. A timer leaves the list when it fires.
    pub fn next_deadline(&self) -> Option<Instant> {
        // Compare the earliest parked wait with the earliest unfired timer
        let state = self.paused.lock();
        let deadline = state.deadlines.keys().next().copied();
        #[cfg(feature = "crossbeam")]
        let deadline = deadline
            .into_iter()
            .chain(state.timers.keys().next().map(|&(deadline, _)| deadline))
            .min();
        deadline
    }

    /// Validates both new times before updating state and notifying waiters.
    fn advance_with(&mut self, next: impl FnOnce(Instant) -> Instant) {
        // Compute and check the new times outside the lock, so a failed check's
        // panic runs no hook under it. Only the owner advances, so nothing
        // changes the times in between.
        let (now, system_time) = {
            let state = self.paused.lock();
            (state.now, state.system_time)
        };
        let next = next(now);
        if next == now {
            return;
        }
        let system_time = system_time
            .checked_add(next - now)
            .expect("clock advance overflows SystemTime");

        // Publish both times and collect live waiters and due timers under one lock
        let mut state = self.paused.lock();
        let signals: Vec<_> = state.signals.values().filter_map(Weak::upgrade).collect();
        state.now = next;
        state.system_time = system_time;
        #[cfg(feature = "crossbeam")]
        let timers = {
            let mut timers = Vec::new();
            while state
                .timers
                .first_key_value()
                .is_some_and(|(&(deadline, _), _)| deadline <= next)
            {
                timers.push(state.timers.pop_first().expect("due timer exists").1);
            }
            self.paused.changed.notify_all();
            timers
        };
        drop(state);

        // Deliver in deadline order with the new times visible and no lock held
        #[cfg(feature = "crossbeam")]
        for timer in timers {
            self.paused.fire(&timer);
        }

        // Wake outside the clock lock, since parking takes the signal lock first
        for signal in signals {
            signal.advance();
        }
    }
}

impl Default for TestClock {
    /// Creates a stopped clock at the current real monotonic and wall times.
    fn default() -> Self {
        Self::new()
    }
}

impl fmt::Debug for TestClock {
    /// Shows the advance, wall time, parked threads and armed timers.
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // Snapshot the counts before calling the formatter's writer
        let state = self.paused.lock();
        let advanced = state.now - self.paused.start;
        let system_time = state.system_time;
        let blocked = state.blocked;
        #[cfg(feature = "crossbeam")]
        let timers = state.timers.len();
        drop(state);

        // Format without the clock lock, since the writer may read this clock
        let mut debug = f.debug_struct("TestClock");
        debug
            .field("advanced", &advanced)
            .field("system_time", &system_time)
            .field("blocked", &blocked);
        #[cfg(feature = "crossbeam")]
        debug.field("timers", &timers);
        debug.finish()
    }
}

/// A test clock's state, shared by its owner, handles and waiters.
pub(crate) struct Paused {
    /// Time the clock started at, to show how far it has advanced.
    start: Instant,
    /// Current times, live signals, parked thread count, deadlines and timers.
    pub(crate) state: Mutex<PausedState>,
    /// Wakes drivers when parked threads or armed timers change.
    pub(crate) changed: Condvar,
    /// One-shot test hook, run before a park's first wait with no clock or
    /// signal lock held.
    #[cfg(test)]
    pub(crate) before_park: std::sync::Mutex<Option<BeforePark>>,
    /// One-shot test hook, run like `before_park` before a park waits again
    /// after a spurious wakeup.
    #[cfg(test)]
    pub(crate) before_rewait: std::sync::Mutex<Option<BeforePark>>,
    /// One-shot test hook, run after a timer's send with no crate lock held.
    #[cfg(all(test, feature = "crossbeam"))]
    after_timer_send: std::sync::Mutex<Option<TimerHook>>,
    /// One-shot test hook, run when a receive's timer wins, before the receiver
    /// is checked again.
    #[cfg(all(test, feature = "crossbeam"))]
    after_timer_receive: std::sync::Mutex<Option<TimerHook>>,
}

/// Mutable part of a test clock.
pub(crate) struct PausedState {
    /// Current monotonic time.
    now: Instant,
    /// Current wall time, independent of monotonic time when set explicitly.
    system_time: SystemTime,
    /// Live waiters and condvars indexed by signal address, removed on drop.
    pub(crate) signals: BTreeMap<usize, Weak<Signal>>,
    /// Threads committed to parking while holding their signal's lock.
    pub(crate) blocked: usize,
    /// Parked timed waits counted by deadline, each unlisted when its park ends.
    deadlines: BTreeMap<Instant, usize>,
    /// Unfired timers by deadline, with each timer's address telling equal
    /// deadlines apart.
    #[cfg(feature = "crossbeam")]
    timers: BTreeMap<(Instant, usize), Arc<Timer>>,
    /// Senders of delivered public timers, kept so their channels stay connected
    /// while the clock lives.
    #[cfg(feature = "crossbeam")]
    fired: Vec<Sender<Instant>>,
}

impl Paused {
    /// Returns armed timers and retained senders for the bookkeeping model.
    #[cfg(all(test, feature = "crossbeam", not(loom)))]
    pub(crate) fn timer_counts(&self) -> (usize, usize) {
        let state = self.lock();
        (state.timers.len(), state.fired.len())
    }

    /// Arms a public timer, whose channel stays connected after delivery.
    #[cfg(feature = "crossbeam")]
    pub(crate) fn at(&self, deadline: Instant) -> Receiver<Instant> {
        self.arm_timer(deadline, true).1
    }

    /// Receives until the clock reaches `deadline`, where a message or a
    /// disconnection wins over expiry, as in crossbeam.
    #[cfg(feature = "crossbeam")]
    pub(crate) fn recv_deadline<T>(
        &self,
        receiver: &Receiver<T>,
        deadline: Instant,
    ) -> Result<T, RecvTimeoutError> {
        // Prefer a ready receiver even when the deadline has already passed
        match receiver.try_recv() {
            Ok(value) => return Ok(value),
            Err(TryRecvError::Disconnected) => return Err(RecvTimeoutError::Disconnected),
            Err(TryRecvError::Empty) => {}
        }
        if self.now() >= deadline {
            return Err(RecvTimeoutError::Timeout);
        }

        // Arm before selecting, and disarm on every return from the adapter
        let (timer, timeout) = self.arm_timer(deadline, false);
        let _registration = ReceiveTimer {
            paused: self,
            timer,
        };

        // Wait on the receiver and the clock's timer, with no real timeout
        crossbeam_channel::select! {
            recv(receiver) -> result => result.map_err(RecvTimeoutError::from),
            recv(timeout) -> _ => {
                // Let tests make the receiver ready after the timer has won
                #[cfg(test)]
                {
                    let hook = self.after_timer_receive.lock().unwrap().take();
                    if let Some(hook) = hook {
                        hook();
                    }
                }

                // Check the receiver again, since select picks randomly among ready arms
                receiver.try_recv().map_err(|err| match err {
                    TryRecvError::Empty => RecvTimeoutError::Timeout,
                    TryRecvError::Disconnected => RecvTimeoutError::Disconnected,
                })
            }
        }
    }

    /// Registers an unfired timer or delivers it immediately if already due.
    #[cfg(feature = "crossbeam")]
    fn arm_timer(&self, deadline: Instant, retain: bool) -> (Arc<Timer>, Receiver<Instant>) {
        // Give the timer a stable address and room for its only message
        let (sender, receiver) = crossbeam_channel::bounded(1);
        let timer = Arc::new(Timer {
            deadline,
            sender,
            retain,
        });

        // Check and register under one lock so an advance cannot miss this timer
        let mut state = self.lock();
        if deadline > state.now {
            state.timers.insert(timer.key(), timer.clone());
            self.changed.notify_all();
            drop(state);
        } else {
            drop(state);
            self.fire(&timer);
        }
        (timer, receiver)
    }

    /// Sends a timer's deadline with no clock lock held, and keeps the sender of
    /// a delivered public timer.
    #[cfg(feature = "crossbeam")]
    fn fire(&self, timer: &Timer) {
        // The capacity-1 channel has never been sent to, so this cannot block
        let delivered = timer.sender.send(timer.deadline).is_ok();

        // Let tests observe publication and delivery order before the next send
        #[cfg(test)]
        {
            let hook = self.after_timer_send.lock().unwrap().take();
            if let Some(hook) = hook {
                hook();
            }
        }

        // Failed deliveries and adapter timers keep no sender in the clock
        if delivered && timer.retain {
            self.lock().fired.push(timer.sender.clone());
        }
    }

    /// Returns the clock's current monotonic time.
    pub(crate) fn now(&self) -> Instant {
        self.lock().now
    }

    /// Returns the clock's current wall time.
    pub(crate) fn system_time(&self) -> SystemTime {
        self.lock().system_time
    }

    /// Returns how far the clock has advanced since its creation.
    pub(crate) fn advanced(&self) -> Duration {
        self.lock().now - self.start
    }

    /// Registers a waiter's signal before its first generation check.
    pub(crate) fn register(&self, signal: &Arc<Signal>) {
        self.lock()
            .signals
            .insert(Arc::as_ptr(signal) as usize, Arc::downgrade(signal));
    }

    /// Removes a waiter's signal without retaining storage for dead waiters.
    pub(crate) fn unregister(&self, signal: &Arc<Signal>) {
        self.lock().signals.remove(&(Arc::as_ptr(signal) as usize));
    }

    /// Counts a park, and lists its deadline, until the guard drops.
    ///
    /// The caller holds its signal lock until it parks, so an advance by a
    /// driver that saw the count still wakes the park.
    pub(crate) fn block(&self, deadline: Option<Instant>) -> Blocked<'_> {
        // Count the park and list its deadline
        let mut state = self.lock();
        state.blocked += 1;
        if let Some(deadline) = deadline {
            *state.deadlines.entry(deadline).or_default() += 1;
        }

        // Wake drivers waiting for the count to grow
        self.changed.notify_all();
        Blocked {
            paused: self,
            deadline,
        }
    }

    /// Takes the one-shot test hook for a park's first wait, or for a wait
    /// `again` after a spurious wakeup, for the caller to run with no lock held.
    #[cfg(test)]
    pub(crate) fn take_hook(&self, again: bool) -> Option<BeforePark> {
        let hook = if again {
            &self.before_rewait
        } else {
            &self.before_park
        };
        hook.lock().unwrap().take()
    }

    /// Locks the state, recovering it from poisoning, since no update under the
    /// lock can stop halfway.
    fn lock(&self) -> MutexGuard<'_, PausedState> {
        self.state.lock().unwrap_or_else(PoisonError::into_inner)
    }
}

/// A timer's one message, and whether its channel outlives delivery.
#[cfg(feature = "crossbeam")]
struct Timer {
    /// Deadline sent as the message, even when an advance overshoots it.
    deadline: Instant,
    /// Sender of the timer's capacity-1 channel.
    sender: Sender<Instant>,
    /// Whether the clock keeps the sender after delivery, as for public timers.
    retain: bool,
}

#[cfg(feature = "crossbeam")]
impl Timer {
    /// Identifies this allocation among timers at the same deadline.
    fn key(&self) -> (Instant, usize) {
        (self.deadline, self as *const Self as usize)
    }
}

/// Removes a receive adapter's timer when it returns, even before expiry.
#[cfg(feature = "crossbeam")]
struct ReceiveTimer<'a> {
    /// Clock holding the armed timer, if it has not fired yet.
    paused: &'a Paused,
    /// Keeps the registration's address unique until it is removed.
    timer: Arc<Timer>,
}

#[cfg(feature = "crossbeam")]
impl Drop for ReceiveTimer<'_> {
    /// Unlists an unfired timer without retaining its sender.
    fn drop(&mut self) {
        self.paused.lock().timers.remove(&self.timer.key());
        self.paused.changed.notify_all();
    }
}

/// Counts one parked thread until its signal lock is retaken after waking.
pub(crate) struct Blocked<'a> {
    /// Clock whose parked count includes this thread.
    paused: &'a Paused,
    /// Deadline registered for this park, if it is timed.
    deadline: Option<Instant>,
}

impl Drop for Blocked<'_> {
    /// Removes this thread and its deadline from the clock's parked state.
    fn drop(&mut self) {
        // Uncount the park and unlist its deadline
        let mut state = self.paused.lock();
        state.blocked -= 1;
        if let Some(deadline) = self.deadline {
            let count = state
                .deadlines
                .get_mut(&deadline)
                .expect("parked deadline exists");
            *count -= 1;
            if *count == 0 {
                state.deadlines.remove(&deadline);
            }
        }

        // Wake watchers of the count, such as tests waiting for it to drop
        self.paused.changed.notify_all();
    }
}

/// Pauses a test wait before parking and exposes its real timer, if any.
#[cfg(test)]
pub(crate) type BeforePark = Box<dyn FnOnce(Option<Instant>) + Send>;

/// Observes a timer race with no crate lock held.
#[cfg(all(test, feature = "crossbeam"))]
type TimerHook = Box<dyn FnOnce() + Send>;

// The timer tests live in src/tests, loaded from here so that they keep this
// module's private items in reach
#[cfg(all(test, feature = "crossbeam", not(loom)))]
#[cfg_attr(coverage_nightly, coverage(off))]
#[path = "tests/paused.rs"]
mod tests;