Skip to main content

thruput_testkit/
clock.rs

1//! Virtual clock owned by the harness.
2//!
3//! Agents read time and sleep only through this clock, never through
4//! `std::time`. The clock only advances when a test advances it, so a whole
5//! scenario moves in virtual time.
6
7use std::sync::{Arc, Mutex, MutexGuard, PoisonError};
8
9use tokio::sync::oneshot;
10
11/// A logical clock measured in whole milliseconds of virtual time.
12#[derive(Clone, Debug, Default)]
13pub struct Clock {
14    inner: Arc<Mutex<ClockState>>,
15}
16
17#[derive(Debug, Default)]
18struct ClockState {
19    now_ms: u64,
20    waiters: Vec<Waiter>,
21}
22
23#[derive(Debug)]
24struct Waiter {
25    target_ms: u64,
26    wake: oneshot::Sender<()>,
27}
28
29impl Clock {
30    /// Create a clock at virtual time zero.
31    pub fn new() -> Self {
32        Self::default()
33    }
34
35    /// Current virtual time, in milliseconds since the clock started.
36    pub fn now_ms(&self) -> u64 {
37        self.lock().now_ms
38    }
39
40    /// Advance virtual time by `delta_ms`, waking every sleeper whose deadline
41    /// has now passed.
42    pub fn advance_ms(&self, delta_ms: u64) {
43        let mut state = self.lock();
44        state.now_ms = state.now_ms.saturating_add(delta_ms);
45        let now = state.now_ms;
46        let mut i = 0;
47        while i < state.waiters.len() {
48            if state.waiters[i].target_ms <= now {
49                let waiter = state.waiters.swap_remove(i);
50                // A dropped receiver just means that sleeper went away.
51                let _ = waiter.wake.send(());
52            } else {
53                i += 1;
54            }
55        }
56    }
57
58    /// Wait until virtual time has advanced by at least `dur_ms`.
59    ///
60    /// Returns immediately when `dur_ms` is zero. Otherwise the future resolves
61    /// only when another caller advances the clock past the deadline, so a
62    /// scenario that sleeps must arrange for time to advance.
63    pub async fn sleep_ms(&self, dur_ms: u64) {
64        let rx = {
65            let mut state = self.lock();
66            let target_ms = state.now_ms.saturating_add(dur_ms);
67            if state.now_ms >= target_ms {
68                return;
69            }
70            let (tx, rx) = oneshot::channel();
71            state.waiters.push(Waiter {
72                target_ms,
73                wake: tx,
74            });
75            rx
76        };
77        // A send error means the clock was torn down; either way, stop waiting.
78        let _ = rx.await;
79    }
80
81    fn lock(&self) -> MutexGuard<'_, ClockState> {
82        self.inner.lock().unwrap_or_else(PoisonError::into_inner)
83    }
84}