gpu_async_reactor/core/
manual.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! A virtual clock: the host for tests, simulation and replay. Time moves only when the owner
4//! moves it, and every [`ManualClock::timer`] sleep fires when it does, so a wait or an executor
5//! run is reproducible to the nanosecond on any machine.
6
7use alloc::collections::BTreeMap;
8use alloc::rc::Rc;
9use alloc::vec::Vec;
10use core::cell::RefCell;
11use core::future::Future;
12use core::pin::Pin;
13use core::task::{Context, Poll, Waker};
14
15use crate::core::time::{Clock, DurNs, MonoNs, Timer};
16
17#[derive(Default)]
18struct State {
19    mono: MonoNs,
20    next_id: u64,
21    timers: BTreeMap<(MonoNs, u64), Waker>,
22}
23
24/// A clock that moves only when told to; clones share it.
25#[derive(Clone, Default)]
26pub struct ManualClock {
27    state: Rc<RefCell<State>>,
28}
29
30/// The [`Timer`] of one [`ManualClock`]; it fires only when that clock moves.
31#[derive(Clone)]
32pub struct ManualTimer {
33    state: Rc<RefCell<State>>,
34}
35
36impl ManualClock {
37    /// A clock at 0 ns with no timers.
38    #[must_use]
39    pub fn new() -> Self {
40        Self::default()
41    }
42
43    /// Another handle on the same clock.
44    #[must_use]
45    pub fn handle(&self) -> Self {
46        self.clone()
47    }
48
49    /// The timer that fires as this clock moves.
50    #[must_use]
51    pub fn timer(&self) -> ManualTimer {
52        ManualTimer {
53            state: Rc::clone(&self.state),
54        }
55    }
56
57    /// Move time forward by `step`, waking every due timer.
58    pub fn advance(&self, step: DurNs) {
59        let to = self.monotonic_ns().after(step);
60        self.set_monotonic(to);
61    }
62
63    /// Set time to `at` (never backwards: an earlier `at` leaves the clock where it is), waking
64    /// every due timer, earliest first.
65    pub fn set_monotonic(&self, at: MonoNs) {
66        let due: Vec<Waker> = {
67            let mut s = self.state.borrow_mut();
68            if at > s.mono {
69                s.mono = at;
70            }
71            let now = s.mono;
72            let pending = match now.0.checked_add(1) {
73                Some(next) => s.timers.split_off(&(MonoNs(next), 0)),
74                None => BTreeMap::new(),
75            };
76            core::mem::replace(&mut s.timers, pending)
77                .into_values()
78                .collect()
79        };
80        for w in due {
81            w.wake();
82        }
83    }
84
85    /// Sleeps that have been polled and are not yet due.
86    #[must_use]
87    pub fn pending_timers(&self) -> usize {
88        self.state.borrow().timers.len()
89    }
90
91    /// The earliest pending sleep's deadline, always later than now.
92    #[must_use]
93    pub fn next_deadline(&self) -> Option<MonoNs> {
94        self.state.borrow().timers.keys().next().map(|&(at, _)| at)
95    }
96}
97
98impl Clock for ManualClock {
99    fn monotonic_ns(&self) -> MonoNs {
100        self.state.borrow().mono
101    }
102}
103
104impl Timer for ManualTimer {
105    async fn sleep_until(&self, at: MonoNs) {
106        Sleep {
107            state: &self.state,
108            at,
109            id: None,
110        }
111        .await;
112    }
113}
114
115struct Sleep<'a> {
116    state: &'a Rc<RefCell<State>>,
117    at: MonoNs,
118    id: Option<u64>,
119}
120
121impl Future for Sleep<'_> {
122    type Output = ();
123
124    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
125        let this = self.get_mut();
126        let mut s = this.state.borrow_mut();
127        if s.mono >= this.at {
128            if let Some(id) = this.id.take() {
129                s.timers.remove(&(this.at, id));
130            }
131            return Poll::Ready(());
132        }
133        let id = *this.id.get_or_insert_with(|| {
134            s.next_id += 1;
135            s.next_id
136        });
137        s.timers.insert((this.at, id), cx.waker().clone());
138        Poll::Pending
139    }
140}
141
142impl Drop for Sleep<'_> {
143    fn drop(&mut self) {
144        if let Some(id) = self.id {
145            self.state.borrow_mut().timers.remove(&(self.at, id));
146        }
147    }
148}
gpu_async_reactor/core/manual.rs

gpu_async_reactor/core/
manual.rs
