1 | #![cfg_attr (not(feature = "rt" ), allow(dead_code))] |
2 | |
3 | //! Source of time abstraction. |
4 | //! |
5 | //! By default, `std::time::Instant::now()` is used. However, when the |
6 | //! `test-util` feature flag is enabled, the values returned for `now()` are |
7 | //! configurable. |
8 | |
9 | cfg_not_test_util! { |
10 | use crate::time::{Instant}; |
11 | |
12 | #[derive(Debug, Clone)] |
13 | pub(crate) struct Clock {} |
14 | |
15 | pub(crate) fn now() -> Instant { |
16 | Instant::from_std(std::time::Instant::now()) |
17 | } |
18 | |
19 | impl Clock { |
20 | pub(crate) fn new(_enable_pausing: bool, _start_paused: bool) -> Clock { |
21 | Clock {} |
22 | } |
23 | |
24 | pub(crate) fn now(&self) -> Instant { |
25 | now() |
26 | } |
27 | } |
28 | } |
29 | |
30 | cfg_test_util! { |
31 | use crate::time::{Duration, Instant}; |
32 | use crate::loom::sync::Mutex; |
33 | use crate::loom::sync::atomic::Ordering; |
34 | use std::sync::atomic::AtomicBool as StdAtomicBool; |
35 | |
36 | cfg_rt! { |
37 | #[track_caller ] |
38 | fn with_clock<R>(f: impl FnOnce(Option<&Clock>) -> Result<R, &'static str>) -> R { |
39 | use crate::runtime::Handle; |
40 | |
41 | let res = match Handle::try_current() { |
42 | Ok(handle) => f(Some(handle.inner.driver().clock())), |
43 | Err(ref e) if e.is_missing_context() => f(None), |
44 | Err(_) => panic!("{}" , crate::util::error::THREAD_LOCAL_DESTROYED_ERROR), |
45 | }; |
46 | |
47 | match res { |
48 | Ok(ret) => ret, |
49 | Err(msg) => panic!("{}" , msg), |
50 | } |
51 | } |
52 | } |
53 | |
54 | cfg_not_rt! { |
55 | #[track_caller ] |
56 | fn with_clock<R>(f: impl FnOnce(Option<&Clock>) -> Result<R, &'static str>) -> R { |
57 | match f(None) { |
58 | Ok(ret) => ret, |
59 | Err(msg) => panic!("{}" , msg), |
60 | } |
61 | } |
62 | } |
63 | |
64 | /// A handle to a source of time. |
65 | #[derive(Debug)] |
66 | pub(crate) struct Clock { |
67 | inner: Mutex<Inner>, |
68 | } |
69 | |
70 | // Used to track if the clock was ever paused. This is an optimization to |
71 | // avoid touching the mutex if `test-util` was accidentally enabled in |
72 | // release mode. |
73 | // |
74 | // A static is used so we can avoid accessing the thread-local as well. The |
75 | // `std` AtomicBool is used directly because loom does not support static |
76 | // atomics. |
77 | static DID_PAUSE_CLOCK: StdAtomicBool = StdAtomicBool::new(false); |
78 | |
79 | #[derive(Debug)] |
80 | struct Inner { |
81 | /// True if the ability to pause time is enabled. |
82 | enable_pausing: bool, |
83 | |
84 | /// Instant to use as the clock's base instant. |
85 | base: std::time::Instant, |
86 | |
87 | /// Instant at which the clock was last unfrozen. |
88 | unfrozen: Option<std::time::Instant>, |
89 | |
90 | /// Number of `inhibit_auto_advance` calls still in effect. |
91 | auto_advance_inhibit_count: usize, |
92 | } |
93 | |
94 | /// Pauses time. |
95 | /// |
96 | /// The current value of `Instant::now()` is saved and all subsequent calls |
97 | /// to `Instant::now()` will return the saved value. The saved value can be |
98 | /// changed by [`advance`] or by the time auto-advancing once the runtime |
99 | /// has no work to do. This only affects the `Instant` type in Tokio, and |
100 | /// the `Instant` in std continues to work as normal. |
101 | /// |
102 | /// Pausing time requires the `current_thread` Tokio runtime. This is the |
103 | /// default runtime used by `#[tokio::test]`. The runtime can be initialized |
104 | /// with time in a paused state using the `Builder::start_paused` method. |
105 | /// |
106 | /// For cases where time is immediately paused, it is better to pause |
107 | /// the time using the `main` or `test` macro: |
108 | /// ``` |
109 | /// #[tokio::main(flavor = "current_thread", start_paused = true)] |
110 | /// async fn main() { |
111 | /// println!("Hello world"); |
112 | /// } |
113 | /// ``` |
114 | /// |
115 | /// # Panics |
116 | /// |
117 | /// Panics if time is already frozen or if called from outside of a |
118 | /// `current_thread` Tokio runtime. |
119 | /// |
120 | /// # Auto-advance |
121 | /// |
122 | /// If time is paused and the runtime has no work to do, the clock is |
123 | /// auto-advanced to the next pending timer. This means that [`Sleep`] or |
124 | /// other timer-backed primitives can cause the runtime to advance the |
125 | /// current time when awaited. |
126 | /// |
127 | /// [`Sleep`]: crate::time::Sleep |
128 | /// [`advance`]: crate::time::advance |
129 | #[track_caller ] |
130 | pub fn pause() { |
131 | with_clock(|maybe_clock| { |
132 | match maybe_clock { |
133 | Some(clock) => clock.pause(), |
134 | None => Err("time cannot be frozen from outside the Tokio runtime" ), |
135 | } |
136 | }); |
137 | } |
138 | |
139 | /// Resumes time. |
140 | /// |
141 | /// Clears the saved `Instant::now()` value. Subsequent calls to |
142 | /// `Instant::now()` will return the value returned by the system call. |
143 | /// |
144 | /// # Panics |
145 | /// |
146 | /// Panics if time is not frozen or if called from outside of the Tokio |
147 | /// runtime. |
148 | #[track_caller ] |
149 | pub fn resume() { |
150 | with_clock(|maybe_clock| { |
151 | let clock = match maybe_clock { |
152 | Some(clock) => clock, |
153 | None => return Err("time cannot be frozen from outside the Tokio runtime" ), |
154 | }; |
155 | |
156 | let mut inner = clock.inner.lock(); |
157 | |
158 | if inner.unfrozen.is_some() { |
159 | return Err("time is not frozen" ); |
160 | } |
161 | |
162 | inner.unfrozen = Some(std::time::Instant::now()); |
163 | Ok(()) |
164 | }); |
165 | } |
166 | |
167 | /// Advances time. |
168 | /// |
169 | /// Increments the saved `Instant::now()` value by `duration`. Subsequent |
170 | /// calls to `Instant::now()` will return the result of the increment. |
171 | /// |
172 | /// This function will make the current time jump forward by the given |
173 | /// duration in one jump. This means that all `sleep` calls with a deadline |
174 | /// before the new time will immediately complete "at the same time", and |
175 | /// the runtime is free to poll them in any order. Additionally, this |
176 | /// method will not wait for the `sleep` calls it advanced past to complete. |
177 | /// If you want to do that, you should instead call [`sleep`] and rely on |
178 | /// the runtime's auto-advance feature. |
179 | /// |
180 | /// Note that calls to `sleep` are not guaranteed to complete the first time |
181 | /// they are polled after a call to `advance`. For example, this can happen |
182 | /// if the runtime has not yet touched the timer driver after the call to |
183 | /// `advance`. However if they don't, the runtime will poll the task again |
184 | /// shortly. |
185 | /// |
186 | /// # Panics |
187 | /// |
188 | /// Panics if time is not frozen or if called from outside of the Tokio |
189 | /// runtime. |
190 | /// |
191 | /// # Auto-advance |
192 | /// |
193 | /// If the time is paused and there is no work to do, the runtime advances |
194 | /// time to the next timer. See [`pause`](pause#auto-advance) for more |
195 | /// details. |
196 | /// |
197 | /// [`sleep`]: fn@crate::time::sleep |
198 | pub async fn advance(duration: Duration) { |
199 | with_clock(|maybe_clock| { |
200 | let clock = match maybe_clock { |
201 | Some(clock) => clock, |
202 | None => return Err("time cannot be frozen from outside the Tokio runtime" ), |
203 | }; |
204 | |
205 | clock.advance(duration) |
206 | }); |
207 | |
208 | crate::task::yield_now().await; |
209 | } |
210 | |
211 | /// Returns the current instant, factoring in frozen time. |
212 | pub(crate) fn now() -> Instant { |
213 | if !DID_PAUSE_CLOCK.load(Ordering::Acquire) { |
214 | return Instant::from_std(std::time::Instant::now()); |
215 | } |
216 | |
217 | with_clock(|maybe_clock| { |
218 | Ok(if let Some(clock) = maybe_clock { |
219 | clock.now() |
220 | } else { |
221 | Instant::from_std(std::time::Instant::now()) |
222 | }) |
223 | }) |
224 | } |
225 | |
226 | impl Clock { |
227 | /// Returns a new `Clock` instance that uses the current execution context's |
228 | /// source of time. |
229 | pub(crate) fn new(enable_pausing: bool, start_paused: bool) -> Clock { |
230 | let now = std::time::Instant::now(); |
231 | |
232 | let clock = Clock { |
233 | inner: Mutex::new(Inner { |
234 | enable_pausing, |
235 | base: now, |
236 | unfrozen: Some(now), |
237 | auto_advance_inhibit_count: 0, |
238 | }), |
239 | }; |
240 | |
241 | if start_paused { |
242 | if let Err(msg) = clock.pause() { |
243 | panic!("{}" , msg); |
244 | } |
245 | } |
246 | |
247 | clock |
248 | } |
249 | |
250 | pub(crate) fn pause(&self) -> Result<(), &'static str> { |
251 | let mut inner = self.inner.lock(); |
252 | |
253 | if !inner.enable_pausing { |
254 | drop(inner); // avoid poisoning the lock |
255 | return Err("`time::pause()` requires the `current_thread` Tokio runtime. \ |
256 | This is the default Runtime used by `#[tokio::test]." ); |
257 | } |
258 | |
259 | // Track that we paused the clock |
260 | DID_PAUSE_CLOCK.store(true, Ordering::Release); |
261 | |
262 | let elapsed = match inner.unfrozen.as_ref() { |
263 | Some(v) => v.elapsed(), |
264 | None => return Err("time is already frozen" ) |
265 | }; |
266 | inner.base += elapsed; |
267 | inner.unfrozen = None; |
268 | |
269 | Ok(()) |
270 | } |
271 | |
272 | /// Temporarily stop auto-advancing the clock (see `tokio::time::pause`). |
273 | pub(crate) fn inhibit_auto_advance(&self) { |
274 | let mut inner = self.inner.lock(); |
275 | inner.auto_advance_inhibit_count += 1; |
276 | } |
277 | |
278 | pub(crate) fn allow_auto_advance(&self) { |
279 | let mut inner = self.inner.lock(); |
280 | inner.auto_advance_inhibit_count -= 1; |
281 | } |
282 | |
283 | pub(crate) fn can_auto_advance(&self) -> bool { |
284 | let inner = self.inner.lock(); |
285 | inner.unfrozen.is_none() && inner.auto_advance_inhibit_count == 0 |
286 | } |
287 | |
288 | pub(crate) fn advance(&self, duration: Duration) -> Result<(), &'static str> { |
289 | let mut inner = self.inner.lock(); |
290 | |
291 | if inner.unfrozen.is_some() { |
292 | return Err("time is not frozen" ); |
293 | } |
294 | |
295 | inner.base += duration; |
296 | Ok(()) |
297 | } |
298 | |
299 | pub(crate) fn now(&self) -> Instant { |
300 | let inner = self.inner.lock(); |
301 | |
302 | let mut ret = inner.base; |
303 | |
304 | if let Some(unfrozen) = inner.unfrozen { |
305 | ret += unfrozen.elapsed(); |
306 | } |
307 | |
308 | Instant::from_std(ret) |
309 | } |
310 | } |
311 | } |
312 | |