thruput-testkit 0.43.0

//! Virtual clock owned by the harness.
//!
//! Agents read time and sleep only through this clock, never through
//! `std::time`. The clock only advances when a test advances it, so a whole
//! scenario moves in virtual time.

use std::sync::{Arc, Mutex, MutexGuard, PoisonError};

use tokio::sync::oneshot;

/// A logical clock measured in whole milliseconds of virtual time.
#[derive(Clone, Debug, Default)]
pub struct Clock {
    inner: Arc<Mutex<ClockState>>,
}

#[derive(Debug, Default)]
struct ClockState {
    now_ms: u64,
    waiters: Vec<Waiter>,
}

#[derive(Debug)]
struct Waiter {
    target_ms: u64,
    wake: oneshot::Sender<()>,
}

impl Clock {
    /// Create a clock at virtual time zero.
    pub fn new() -> Self {
        Self::default()
    }

    /// Current virtual time, in milliseconds since the clock started.
    pub fn now_ms(&self) -> u64 {
        self.lock().now_ms
    }

    /// Advance virtual time by `delta_ms`, waking every sleeper whose deadline
    /// has now passed.
    pub fn advance_ms(&self, delta_ms: u64) {
        let mut state = self.lock();
        state.now_ms = state.now_ms.saturating_add(delta_ms);
        let now = state.now_ms;
        let mut i = 0;
        while i < state.waiters.len() {
            if state.waiters[i].target_ms <= now {
                let waiter = state.waiters.swap_remove(i);
                // A dropped receiver just means that sleeper went away.
                let _ = waiter.wake.send(());
            } else {
                i += 1;
            }
        }
    }

    /// Wait until virtual time has advanced by at least `dur_ms`.
    ///
    /// Returns immediately when `dur_ms` is zero. Otherwise the future resolves
    /// only when another caller advances the clock past the deadline, so a
    /// scenario that sleeps must arrange for time to advance.
    pub async fn sleep_ms(&self, dur_ms: u64) {
        let rx = {
            let mut state = self.lock();
            let target_ms = state.now_ms.saturating_add(dur_ms);
            if state.now_ms >= target_ms {
                return;
            }
            let (tx, rx) = oneshot::channel();
            state.waiters.push(Waiter {
                target_ms,
                wake: tx,
            });
            rx
        };
        // A send error means the clock was torn down; either way, stop waiting.
        let _ = rx.await;
    }

    fn lock(&self) -> MutexGuard<'_, ClockState> {
        self.inner.lock().unwrap_or_else(PoisonError::into_inner)
    }
}