1use crate::{Clock, Timer};
2use futures::channel::oneshot;
3use jiff::Timestamp;
4use std::pin::Pin;
5use std::sync::{Arc, Mutex, PoisonError};
6use std::time::Duration;
7
8#[derive(Debug)]
12pub struct FrozenClock {
13 instant: Timestamp,
14}
15
16impl FrozenClock {
17 pub fn at(rfc3339: &str) -> Self {
20 let instant = rfc3339.parse().unwrap_or(Timestamp::UNIX_EPOCH);
21 Self { instant }
22 }
23
24 pub fn at_instant(instant: Timestamp) -> Self {
26 Self { instant }
27 }
28}
29
30impl Clock for FrozenClock {
31 fn now(&self) -> Timestamp {
32 self.instant
33 }
34}
35
36type PendingSleeps = Vec<(Timestamp, oneshot::Sender<()>)>;
39
40#[derive(Debug)]
42pub struct MockClock {
43 inner: Arc<Mutex<Timestamp>>,
44 pending: Arc<Mutex<PendingSleeps>>,
45}
46
47impl MockClock {
48 pub fn new(start: Timestamp) -> Self {
50 Self {
51 inner: Arc::new(Mutex::new(start)),
52 pending: Arc::new(Mutex::new(Vec::new())),
53 }
54 }
55
56 pub fn set_now(&self, instant: Timestamp) {
59 let mut guard = self.inner.lock().unwrap_or_else(PoisonError::into_inner);
60 *guard = instant;
61 drop(guard);
62 self.fire_pending();
63 }
64
65 pub fn advance_by(&self, dur: Duration) {
68 let new_now = {
69 let mut guard = self.inner.lock().unwrap_or_else(PoisonError::into_inner);
70 *guard = guard.checked_add(dur).unwrap_or(*guard);
71 *guard
72 };
73 self.fire_ready(new_now);
74 }
75
76 pub fn advance_and_fire(&self, dur: Duration) {
79 self.advance_by(dur);
80 }
81
82 fn current(&self) -> Timestamp {
83 *self.inner.lock().unwrap_or_else(PoisonError::into_inner)
84 }
85
86 fn fire_pending(&self) {
87 let now = self.current();
88 self.fire_ready(now);
89 }
90
91 fn fire_ready(&self, now: Timestamp) {
92 let to_fire: PendingSleeps = {
95 let mut pending = self.pending.lock().unwrap_or_else(PoisonError::into_inner);
96 let (fired, remaining): (PendingSleeps, PendingSleeps) = std::mem::take(&mut *pending)
97 .into_iter()
98 .partition(|(wake_at, _)| *wake_at <= now);
99 *pending = remaining;
100 fired
101 };
102 for (_, sender) in to_fire {
103 let _ = sender.send(());
105 }
106 }
107
108 pub(crate) fn register_sleep(&self, wake_at: Timestamp) -> oneshot::Receiver<()> {
109 let (tx, rx) = oneshot::channel();
110 let now = self.current();
111 if wake_at <= now {
112 let _ = tx.send(());
114 } else {
115 let mut pending = self.pending.lock().unwrap_or_else(PoisonError::into_inner);
116 pending.push((wake_at, tx));
117 }
118 rx
119 }
120}
121
122impl Clock for MockClock {
123 fn now(&self) -> Timestamp {
124 self.current()
125 }
126}
127
128#[derive(Debug)]
133pub struct AdvanceableTimer {
134 clock: Arc<MockClock>,
135}
136
137impl AdvanceableTimer {
138 pub fn new(clock: Arc<MockClock>) -> Self {
140 Self { clock }
141 }
142}
143
144impl Timer for AdvanceableTimer {
145 fn sleep(&self, dur: Duration) -> Pin<Box<dyn Future<Output = ()> + Send>> {
146 let now = self.clock.current();
147 let wake_at = now.checked_add(dur).unwrap_or(now);
148 let rx = self.clock.register_sleep(wake_at);
149 Box::pin(async move {
150 let _ = rx.await;
152 })
153 }
154
155 fn next_tick(&self) -> Pin<Box<dyn Future<Output = ()> + Send>> {
156 Box::pin(futures::future::ready(()))
157 }
158}
159
160pub fn frozen_at(rfc3339: &str) -> Arc<dyn Clock> {
164 Arc::new(FrozenClock::at(rfc3339))
165}
166
167pub fn mock_pair(start: Timestamp) -> (Arc<MockClock>, Arc<AdvanceableTimer>) {
169 let clock = Arc::new(MockClock::new(start));
170 let timer = Arc::new(AdvanceableTimer::new(Arc::clone(&clock)));
171 (clock, timer)
172}
173
174#[cfg(test)]
175mod tests {
176 #![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
177 use super::*;
178
179 fn t0() -> Timestamp {
180 "2026-05-10T00:00:00Z".parse().expect("valid RFC 3339")
181 }
182
183 #[test]
186 fn frozen_at_parses_rfc3339() {
187 let clock = FrozenClock::at("2026-05-10T00:00:00Z");
188 assert_eq!(clock.now(), t0());
189 }
190
191 #[test]
192 fn frozen_clock_at_instant() {
193 let clock = FrozenClock::at_instant(t0());
194 assert_eq!(clock.now(), t0());
195 }
196
197 #[test]
198 fn frozen_clock_never_advances() {
199 let clock = FrozenClock::at_instant(t0());
200 let a = clock.now();
201 let b = clock.now();
202 assert_eq!(a, b);
203 }
204
205 #[test]
206 fn frozen_clock_unix_millis() {
207 let clock = FrozenClock::at("1970-01-01T00:00:01Z");
208 assert_eq!(clock.now_unix_millis(), 1_000);
209 }
210
211 #[test]
212 fn frozen_at_helper_trait_dispatch() {
213 let c: Arc<dyn Clock> = frozen_at("2026-05-10T00:00:00Z");
214 assert_eq!(c.now(), t0());
215 }
216
217 #[test]
220 fn mock_clock_starts_at_given_instant() {
221 let clock = MockClock::new(t0());
222 assert_eq!(clock.now(), t0());
223 }
224
225 #[test]
226 fn mock_clock_advance_by_increases_time() {
227 let clock = MockClock::new(t0());
228 clock.advance_by(Duration::from_secs(1));
229 assert_eq!(clock.now(), t0() + jiff::SignedDuration::from_secs(1));
230 }
231
232 #[test]
233 fn mock_clock_set_now_updates_instant() {
234 let clock = MockClock::new(t0());
235 let new_now = t0() + jiff::SignedDuration::from_hours(5);
236 clock.set_now(new_now);
237 assert_eq!(clock.now(), new_now);
238 }
239
240 #[test]
241 fn mock_clock_trait_dispatch() {
242 let c: Arc<dyn Clock> = Arc::new(MockClock::new(t0()));
243 assert_eq!(c.now(), t0());
244 }
245
246 #[tokio::test]
249 async fn advanceable_timer_sleep_resolves_after_advance() {
250 let (clock, timer) = mock_pair(t0());
251 let sleep_fut = timer.sleep(Duration::from_secs(60));
252 let clock2 = Arc::clone(&clock);
254 tokio::spawn(async move {
255 clock2.advance_and_fire(Duration::from_secs(60));
256 })
257 .await
258 .ok();
259 sleep_fut.await; }
261
262 #[tokio::test]
263 async fn advanceable_timer_partial_advance_fires_only_due_sleeps() {
264 let (clock, timer) = mock_pair(t0());
265 let sleep_30 = timer.sleep(Duration::from_secs(30));
266 let sleep_60 = timer.sleep(Duration::from_secs(60));
267
268 clock.advance_and_fire(Duration::from_secs(45));
269
270 sleep_30.await;
272
273 clock.advance_and_fire(Duration::from_secs(15));
275 sleep_60.await;
276 }
277
278 #[tokio::test]
279 async fn advanceable_timer_next_tick_resolves_immediately() {
280 let (_, timer) = mock_pair(t0());
281 timer.next_tick().await;
282 }
283
284 #[test]
285 fn mock_pair_helper_returns_linked_pair() {
286 let (clock, timer) = mock_pair(t0());
287 let _: Arc<MockClock> = clock;
288 let _: Arc<AdvanceableTimer> = timer;
289 }
290
291 #[test]
292 fn frozen_clock_send_sync() {
293 fn assert_send_sync<T: Send + Sync>() {}
294 assert_send_sync::<FrozenClock>();
295 }
296
297 #[test]
298 fn mock_clock_send_sync() {
299 fn assert_send_sync<T: Send + Sync>() {}
300 assert_send_sync::<MockClock>();
301 assert_send_sync::<AdvanceableTimer>();
302 }
303}