execution_policy/core/
test.rs

1use std::future::Future;
2use std::pin::Pin;
3use std::sync::{Arc, Mutex};
4use std::task::{Context, Poll, Waker};
5use std::time::{Duration, Instant};
6
7use super::rng::SplitMix64;
8use super::{BoxFuture, Core};
9
10struct ClockState {
11    base: Instant,
12    offset: Duration,
13    wakers: Vec<(Duration, Waker)>,
14}
15
16/// A virtual clock for deterministic tests. `now()` advances only via
17/// [`ManualClock::advance`]; sleeps registered on it resolve when crossed.
18#[derive(Clone)]
19pub struct ManualClock(Arc<Mutex<ClockState>>);
20
21impl std::fmt::Debug for ManualClock {
22    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
23        f.write_str("ManualClock")
24    }
25}
26
27impl Default for ManualClock {
28    fn default() -> Self {
29        Self::new()
30    }
31}
32
33impl ManualClock {
34    pub fn new() -> Self {
35        Self(Arc::new(Mutex::new(ClockState {
36            base: Instant::now(),
37            offset: Duration::ZERO,
38            wakers: Vec::new(),
39        })))
40    }
41
42    /// Advance virtual time, waking any sleeps whose deadline is now crossed.
43    pub fn advance(&self, dur: Duration) {
44        let mut st = self.0.lock().unwrap();
45        st.offset += dur;
46        let now = st.offset;
47        let mut ready = Vec::new();
48        st.wakers.retain(|(deadline, w)| {
49            if *deadline <= now {
50                ready.push(w.clone());
51                false
52            } else {
53                true
54            }
55        });
56        drop(st);
57        for w in ready {
58            w.wake();
59        }
60    }
61
62    fn now(&self) -> Instant {
63        let st = self.0.lock().unwrap();
64        st.base + st.offset
65    }
66
67    fn register(&self, deadline: Duration, waker: Waker) {
68        self.0.lock().unwrap().wakers.push((deadline, waker));
69    }
70
71    fn offset(&self) -> Duration {
72        self.0.lock().unwrap().offset
73    }
74}
75
76struct ManualSleep {
77    clock: ManualClock,
78    deadline: Duration,
79}
80
81impl Future for ManualSleep {
82    type Output = ();
83    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
84        if self.clock.offset() >= self.deadline {
85            Poll::Ready(())
86        } else {
87            self.clock.register(self.deadline, cx.waker().clone());
88            Poll::Pending
89        }
90    }
91}
92
93/// Deterministic [`Core`] for tests: virtual clock + seeded RNG, no real sleeps.
94#[derive(Debug)]
95pub struct TestCore {
96    clock: ManualClock,
97    rng: Mutex<SplitMix64>,
98}
99
100impl TestCore {
101    pub fn new(clock: ManualClock) -> Self {
102        Self::with_seed(clock, 0xDEAD_BEEF)
103    }
104
105    pub fn with_seed(clock: ManualClock, seed: u64) -> Self {
106        Self {
107            clock,
108            rng: Mutex::new(SplitMix64::new(seed)),
109        }
110    }
111}
112
113impl Core for TestCore {
114    fn now(&self) -> Instant {
115        self.clock.now()
116    }
117
118    fn sleep(&self, dur: Duration) -> BoxFuture<'_, ()> {
119        let deadline = self.clock.offset() + dur;
120        Box::pin(ManualSleep {
121            clock: self.clock.clone(),
122            deadline,
123        })
124    }
125
126    fn next_u64(&self) -> u64 {
127        self.rng.lock().unwrap().next_u64()
128    }
129}
130
131#[cfg(test)]
132mod tests {
133    use super::*;
134
135    #[test]
136    fn now_advances_only_on_advance() {
137        let clock = ManualClock::new();
138        let core = TestCore::new(clock.clone());
139        let t0 = core.now();
140        clock.advance(Duration::from_secs(5));
141        assert_eq!(core.now().duration_since(t0), Duration::from_secs(5));
142    }
143
144    #[tokio::test]
145    async fn sleep_resolves_when_clock_crosses_deadline() {
146        let clock = ManualClock::new();
147        let core = TestCore::new(clock.clone());
148        let mut fut = Box::pin(core.sleep(Duration::from_secs(10)));
149        let poll = noop_poll(&mut fut);
150        assert!(poll.is_pending());
151        clock.advance(Duration::from_secs(10));
152        fut.await;
153    }
154
155    fn noop_poll<F: Future + Unpin>(f: &mut F) -> Poll<F::Output> {
156        let waker = Waker::noop();
157        let mut cx = Context::from_waker(waker);
158        Pin::new(f).poll(&mut cx)
159    }
160}
execution_policy/core/test.rs

execution_policy/core/
test.rs
