Skip to main content

agent_effects/
clock.rs

1//! Time source for leases, retry scheduling and settle delays.
2//!
3//! Injected rather than read from [`SystemTime::now`] directly, so tests can
4//! step time deterministically and stores can substitute their own clock.
5
6use std::sync::{Arc, Mutex};
7use std::time::{Duration, SystemTime};
8
9/// A source of wall-clock time.
10pub trait Clock: Send + Sync + 'static {
11    /// The current time.
12    fn now(&self) -> SystemTime;
13}
14
15/// Lets a test keep a handle to a clock it gave to a runtime.
16impl<C: Clock> Clock for Arc<C> {
17    fn now(&self) -> SystemTime {
18        C::now(self)
19    }
20}
21
22/// The operating system clock.
23#[derive(Clone, Copy, Debug, Default)]
24pub struct SystemClock;
25
26impl Clock for SystemClock {
27    fn now(&self) -> SystemTime {
28        SystemTime::now()
29    }
30}
31
32/// Wall-clock time that follows Tokio's clock.
33///
34/// Anchored to the system time when created, then advanced by
35/// [`tokio::time::Instant`]. Under `#[tokio::test(start_paused = true)]`,
36/// Tokio skips idle time, and this clock skips with it. Backoff sleeps,
37/// lease expiry and settle delays then all agree, and a test with minutes of
38/// retries finishes instantly.
39#[derive(Clone, Copy, Debug)]
40pub struct TokioClock {
41    system_base: SystemTime,
42    tokio_base: tokio::time::Instant,
43}
44
45impl TokioClock {
46    /// A clock anchored at the current system and Tokio time.
47    pub fn new() -> Self {
48        Self {
49            system_base: SystemTime::now(),
50            tokio_base: tokio::time::Instant::now(),
51        }
52    }
53}
54
55impl Default for TokioClock {
56    fn default() -> Self {
57        Self::new()
58    }
59}
60
61impl Clock for TokioClock {
62    fn now(&self) -> SystemTime {
63        self.system_base + self.tokio_base.elapsed()
64    }
65}
66
67/// A clock that only moves when told to. For tests.
68#[derive(Debug)]
69pub struct ManualClock {
70    now: Mutex<SystemTime>,
71}
72
73impl ManualClock {
74    /// A clock frozen at `start`.
75    pub const fn new(start: SystemTime) -> Self {
76        Self {
77            now: Mutex::new(start),
78        }
79    }
80
81    /// Moves the clock forward by `by`.
82    pub fn advance(&self, by: Duration) {
83        let mut now = self
84            .now
85            .lock()
86            .unwrap_or_else(std::sync::PoisonError::into_inner);
87        *now += by;
88    }
89
90    /// Sets the clock to `to`, which may be earlier, to simulate skew.
91    pub fn set(&self, to: SystemTime) {
92        *self
93            .now
94            .lock()
95            .unwrap_or_else(std::sync::PoisonError::into_inner) = to;
96    }
97}
98
99impl Default for ManualClock {
100    fn default() -> Self {
101        Self::new(SystemTime::UNIX_EPOCH + Duration::from_secs(1_700_000_000))
102    }
103}
104
105impl Clock for ManualClock {
106    fn now(&self) -> SystemTime {
107        *self
108            .now
109            .lock()
110            .unwrap_or_else(std::sync::PoisonError::into_inner)
111    }
112}
113
114#[cfg(test)]
115mod tests {
116    use super::*;
117
118    #[tokio::test(start_paused = true)]
119    async fn tokio_clock_follows_paused_time() {
120        let clock = TokioClock::new();
121        let start = clock.now();
122        tokio::time::sleep(Duration::from_secs(90)).await;
123        assert_eq!(clock.now(), start + Duration::from_secs(90));
124    }
125
126    #[test]
127    fn manual_clock_only_moves_when_told() {
128        let clock = ManualClock::default();
129        let start = clock.now();
130        assert_eq!(clock.now(), start);
131        clock.advance(Duration::from_secs(30));
132        assert_eq!(clock.now(), start + Duration::from_secs(30));
133    }
134}