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}