polkadot_node_clock/
mock.rs

1// Copyright (C) Parity Technologies (UK) Ltd.
2// This file is part of Polkadot.
3
4// Polkadot is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, either version 3 of the License, or
7// (at your option) any later version.
8
9// Polkadot is distributed in the hope that it will be useful,
10// but WITHOUT ANY WARRANTY; without even the implied warranty of
11// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12// GNU General Public License for more details.
13
14// You should have received a copy of the GNU General Public License
15// along with Polkadot.  If not, see <http://www.gnu.org/licenses/>.
16
17//! Deterministic, fully-virtual [`Clock`] implementation for tests.
18//!
19//! Virtualises all three time channels: `now()` (monotonic), `delay()` (driven by a wakeup
20//! queue rather than a real timer), and `duration_since_epoch()` (wall clock). Tests advance
21//! time explicitly via [`MockClock::advance`] / [`MockClock::advance_secs`]; pending
22//! [`Clock::delay`] futures resolve when their deadline is crossed.
23//!
24//! ```ignore
25//! let clock = Arc::new(MockClock::default());
26//! let mut delay = clock.delay(Duration::from_millis(500));
27//! clock.advance(Duration::from_millis(500)); // `delay` now resolves
28//! ```
29
30use crate::{BoxedDelay, Clock};
31use futures::channel::oneshot;
32use std::{
33	sync::{Arc, Mutex},
34	time::{Duration, Instant},
35};
36
37/// A deterministic clock backed by manually-advanced inner state.
38///
39/// `Clone` shares the same inner state — the same handle can be installed in a subsystem
40/// (`Arc<dyn Clock>`) and held by the test (`Arc<MockClock>`) to advance time.
41#[derive(Clone, Debug, Default)]
42pub struct MockClock {
43	inner: Arc<Mutex<MockClockInner>>,
44}
45
46impl MockClock {
47	/// Advance the clock by `dur`. All wakeups whose deadline is `<= self.now() + dur` resolve.
48	pub fn advance(&self, dur: Duration) {
49		let new_now = {
50			let mut inner = self.inner.lock().expect("MockClock mutex poisoned");
51			inner.now += dur;
52			inner.wall_clock_ms = inner.wall_clock_ms.saturating_add(dur.as_millis());
53			inner.now
54		};
55		self.inner.lock().expect("MockClock mutex poisoned").wakeup_up_to(new_now);
56	}
57
58	/// Advance the clock by `secs` seconds. Sugar for [`MockClock::advance`].
59	pub fn advance_secs(&self, secs: u64) {
60		self.advance(Duration::from_secs(secs));
61	}
62
63	/// Advance the clock to the deadline of the next pending wakeup, returning the elapsed
64	/// duration. Returns `None` when no wakeups are pending.
65	pub fn advance_to_next_wakeup(&self) -> Option<Duration> {
66		let next = self.inner.lock().expect("MockClock mutex poisoned").next_wakeup()?;
67		let now = self.inner.lock().expect("MockClock mutex poisoned").now;
68		let dur = next.saturating_duration_since(now);
69		self.advance(dur);
70		Some(dur)
71	}
72
73	/// Peek at the duration until the next pending wakeup, without advancing the clock.
74	/// Returns `None` when no wakeups are pending.
75	pub fn next_wakeup_in(&self) -> Option<Duration> {
76		let inner = self.inner.lock().expect("MockClock mutex poisoned");
77		let next = inner.next_wakeup()?;
78		Some(next.saturating_duration_since(inner.now))
79	}
80}
81
82impl Clock for MockClock {
83	fn now(&self) -> Instant {
84		self.inner.lock().expect("MockClock mutex poisoned").now
85	}
86
87	fn delay(&self, dur: Duration) -> BoxedDelay {
88		let deadline = {
89			let inner = self.inner.lock().expect("MockClock mutex poisoned");
90			inner.now + dur
91		};
92		let rx = self.inner.lock().expect("MockClock mutex poisoned").register_wakeup(deadline);
93
94		Box::pin(async move {
95			// `oneshot::Receiver::await` resolves with `Err` when the sender is dropped. That
96			// happens if the `MockClock` is dropped before the wakeup fires; in that case the
97			// surrounding subsystem is shutting down and the future returning here is fine.
98			let _ = rx.await;
99		})
100	}
101
102	fn duration_since_epoch(&self) -> Duration {
103		Duration::from_millis(
104			self.inner.lock().expect("MockClock mutex poisoned").wall_clock_ms as u64,
105		)
106	}
107}
108
109#[derive(Debug)]
110struct MockClockInner {
111	now: Instant,
112	wall_clock_ms: u128,
113	/// Pending wakeups, sorted by deadline.
114	wakeups: Vec<(Instant, oneshot::Sender<()>)>,
115}
116
117impl Default for MockClockInner {
118	fn default() -> Self {
119		Self { now: Instant::now(), wall_clock_ms: 0, wakeups: Vec::new() }
120	}
121}
122
123impl MockClockInner {
124	/// Resolve all wakeups whose deadline is `<= up_to`.
125	fn wakeup_up_to(&mut self, up_to: Instant) {
126		let drain_up_to = self.wakeups.partition_point(|w| w.0 <= up_to);
127		for (_, wakeup) in self.wakeups.drain(..drain_up_to) {
128			let _ = wakeup.send(());
129		}
130	}
131
132	/// Deadline of the next pending wakeup, if any.
133	fn next_wakeup(&self) -> Option<Instant> {
134		self.wakeups.first().map(|w| w.0)
135	}
136
137	/// Register a new wakeup. If `deadline <= now` resolves immediately.
138	fn register_wakeup(&mut self, deadline: Instant) -> oneshot::Receiver<()> {
139		let (tx, rx) = oneshot::channel();
140		let pos = self.wakeups.partition_point(|w| w.0 <= deadline);
141		self.wakeups.insert(pos, (deadline, tx));
142		// Resolve immediately if the deadline is already past.
143		self.wakeup_up_to(self.now);
144		rx
145	}
146}
147
148#[cfg(test)]
149mod tests {
150	use super::*;
151	use futures::FutureExt;
152
153	#[test]
154	fn now_advances() {
155		let clock = MockClock::default();
156		let start = clock.now();
157		clock.advance(Duration::from_secs(1));
158		assert_eq!(clock.now(), start + Duration::from_secs(1));
159	}
160
161	#[test]
162	fn duration_since_epoch_advances() {
163		let clock = MockClock::default();
164		assert_eq!(clock.duration_since_epoch().as_millis(), 0);
165		clock.advance(Duration::from_millis(123));
166		assert_eq!(clock.duration_since_epoch().as_millis(), 123);
167		clock.advance_secs(1);
168		assert_eq!(clock.duration_since_epoch().as_millis(), 1_123);
169	}
170
171	#[tokio::test]
172	async fn delay_resolves_on_advance() {
173		let clock = Arc::new(MockClock::default());
174		let mut delay = clock.delay(Duration::from_millis(100));
175		// Not yet ready.
176		assert!((&mut delay).now_or_never().is_none());
177		clock.advance(Duration::from_millis(50));
178		assert!((&mut delay).now_or_never().is_none());
179		clock.advance(Duration::from_millis(50));
180		assert!(delay.now_or_never().is_some());
181	}
182
183	#[tokio::test]
184	async fn advance_to_next_wakeup_jumps() {
185		let clock = Arc::new(MockClock::default());
186		let _delay_a = clock.delay(Duration::from_millis(200));
187		let _delay_b = clock.delay(Duration::from_millis(500));
188		let elapsed = clock.advance_to_next_wakeup().expect("a wakeup is pending");
189		assert_eq!(elapsed, Duration::from_millis(200));
190		let elapsed = clock.advance_to_next_wakeup().expect("a wakeup is pending");
191		assert_eq!(elapsed, Duration::from_millis(300));
192		assert!(clock.advance_to_next_wakeup().is_none());
193	}
194}
polkadot_node_clock/mock.rs

polkadot_node_clock/
mock.rs
