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
9cfg_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
30cfg_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