spate_core/clock/
tokio.rs

1//! Injectable time on the tokio timer: [`Clock`] reads
2//! [`tokio::time::Instant`] and sleeps until a deadline.
3//!
4//! [`SystemClock`] follows the runtime's clock, so paused tokio time applies
5//! to it. `TestClock` moves only when a test advances it, for code on a
6//! multi-thread runtime, where tokio time cannot be paused.
7
8use std::future::Future;
9use std::pin::Pin;
10// Only `TestClock` needs `Duration`; the trait and `SystemClock` do not.
11#[cfg(any(test, feature = "testing"))]
12use std::time::Duration;
13use tokio::time::Instant;
14
15/// A future that resolves when a [`Clock`] reaches some instant.
16///
17/// Boxed and `'static`: implementations must not borrow the clock, so a
18/// caller can hold the future across a `select!` arm whose body needs
19/// `&mut self`.
20pub type Sleep = Pin<Box<dyn Future<Output = ()> + Send + 'static>>;
21
22/// A monotonic clock on [`tokio::time::Instant`].
23pub trait Clock: Send + Sync + std::fmt::Debug {
24    /// The current instant on this clock's timeline.
25    fn now(&self) -> Instant;
26
27    /// Resolve once this clock reaches `deadline`.
28    ///
29    /// Returns immediately if `deadline` is already past. On a frozen
30    /// clock this parks until the test advances past `deadline`. It must
31    /// never fall back to real time, or the freeze leaks.
32    fn sleep_until(&self, deadline: Instant) -> Sleep;
33}
34
35/// The production clock, on the runtime's tokio time.
36#[derive(Debug, Default, Clone, Copy)]
37pub struct SystemClock;
38
39impl Clock for SystemClock {
40    fn now(&self) -> Instant {
41        Instant::now()
42    }
43
44    fn sleep_until(&self, deadline: Instant) -> Sleep {
45        Box::pin(tokio::time::sleep_until(deadline))
46    }
47}
48
49/// A clock that only moves when a test moves it.
50///
51/// Wired into production code, it stops every timer drawn from it.
52#[cfg(any(test, feature = "testing"))]
53#[doc(hidden)]
54#[derive(Debug)]
55pub struct TestClock {
56    /// The instant this clock was created at. Fixed for its lifetime.
57    base: Instant,
58    /// Nanoseconds elapsed on this clock's timeline. The watch channel is
59    /// the single source of truth for "now" *and* the wakeup for parked
60    /// [`Clock::sleep_until`] futures, so the two can never disagree.
61    offset_nanos: tokio::sync::watch::Sender<u64>,
62}
63
64#[cfg(any(test, feature = "testing"))]
65impl TestClock {
66    /// A clock frozen at construction. Build it after the runtime exists
67    /// so `tokio::time::Instant::now()` reads the runtime's clock.
68    #[must_use]
69    pub fn frozen() -> std::sync::Arc<TestClock> {
70        std::sync::Arc::new(TestClock {
71            base: Instant::now(),
72            offset_nanos: tokio::sync::watch::Sender::new(0),
73        })
74    }
75
76    /// Move the clock forward, waking everything parked on a deadline this
77    /// jump passes.
78    ///
79    /// One jump fires every deadline it passes at the same instant. Use
80    /// [`advance_stepped`](TestClock::advance_stepped) when code under test
81    /// must react between them.
82    pub fn advance(&self, by: Duration) {
83        let nanos = u64::try_from(by.as_nanos()).expect("test clock advance fits in u64 nanos");
84        self.offset_nanos.send_modify(|n| *n += nanos);
85    }
86
87    /// Advance by `total` in `step`-sized increments, running `between`
88    /// after each one.
89    ///
90    /// `between` is where the test lets the code under test react before
91    /// the next step. Panics if `step` is zero.
92    pub fn advance_stepped(&self, total: Duration, step: Duration, mut between: impl FnMut()) {
93        assert!(!step.is_zero(), "advance_stepped needs a non-zero step");
94        let mut moved = Duration::ZERO;
95        while moved < total {
96            let chunk = step.min(total - moved);
97            self.advance(chunk);
98            moved += chunk;
99            between();
100        }
101    }
102}
103
104#[cfg(any(test, feature = "testing"))]
105impl Clock for TestClock {
106    fn now(&self) -> Instant {
107        self.base + Duration::from_nanos(*self.offset_nanos.borrow())
108    }
109
110    fn sleep_until(&self, deadline: Instant) -> Sleep {
111        let base = self.base;
112        let mut rx = self.offset_nanos.subscribe();
113        Box::pin(async move {
114            loop {
115                if base + Duration::from_nanos(*rx.borrow_and_update()) >= deadline {
116                    return;
117                }
118                if rx.changed().await.is_err() {
119                    // The clock is gone, so it can never reach `deadline`.
120                    // Park rather than return: waking would fire a timer
121                    // whose time never came.
122                    std::future::pending::<()>().await;
123                }
124            }
125        })
126    }
127}
spate_core/clock/tokio.rs

spate_core/clock/
tokio.rs
