1
//! Simple provider of simulated time
2
//!
3
//! See [`SimpleMockTimeProvider`]
4

            
5
use std::cmp::Reverse;
6
use std::future::Future;
7
use std::pin::Pin;
8
use std::sync::{Arc, Mutex, MutexGuard};
9
use std::task::{Context, Poll, Waker};
10
use web_time_compat::{Duration, Instant, InstantExt, SystemTime, SystemTimeExt};
11

            
12
use derive_more::AsMut;
13
use priority_queue::priority_queue::PriorityQueue;
14
use slotmap_careful::DenseSlotMap;
15

            
16
use tor_rtcompat::CoarseInstant;
17
use tor_rtcompat::CoarseTimeProvider;
18
use tor_rtcompat::SleepProvider;
19

            
20
use crate::time_core::MockTimeCore;
21

            
22
/// Simple provider of simulated time
23
///
24
/// Maintains a mocked view of the current [`Instant`] and [`SystemTime`].
25
///
26
/// The simulated time advances only when explicitly instructed,
27
/// by calling [`.advance()`](Provider::advance).
28
///
29
/// The wallclock time can be warped with
30
/// [`.jump_wallclock()`](Provider::jump_wallclock),
31
/// allowing simulation of wall clock non-monotonicity.
32
///
33
/// # Panics and aborts
34
///
35
/// Panics on time under/overflow.
36
///
37
/// May cause an abort if the [`SimpleMockTimeProvider`] implementation contains bugs.
38
#[derive(Clone, Debug)]
39
pub struct SimpleMockTimeProvider {
40
    /// The actual state
41
    state: Arc<Mutex<State>>,
42
}
43

            
44
/// Convenience abbreviation
45
pub(crate) use SimpleMockTimeProvider as Provider;
46

            
47
/// Identifier of a [`SleepFuture`]
48
type Id = slotmap_careful::DefaultKey;
49

            
50
/// Future for `sleep`
51
///
52
/// Iff this struct exists, there is an entry for `id` in `prov.futures`.
53
/// (It might contain `None`.)
54
pub struct SleepFuture {
55
    /// Reference to our state
56
    prov: Provider,
57

            
58
    /// Which `SleepFuture` are we
59
    id: Id,
60
}
61

            
62
/// Mutable state for a [`Provider`]
63
///
64
/// Each sleep ([`Id`], [`SleepFuture`]) is in one of the following states:
65
///
66
/// | state       | [`SleepFuture`]  | `futures`        | `unready`          |
67
/// |-------------|------------------|------------------|--------------------|
68
/// | UNPOLLED    | exists           | present, `None`  | present, `> now`   |
69
/// | WAITING     | exists           | present, `Some`  | present, `> now`   |
70
/// | READY       | exists           | present, `None`  | absent             |
71
/// | DROPPED     | dropped          | absent           | absent             |
72
#[derive(Debug, AsMut)]
73
struct State {
74
    /// Current time
75
    ///
76
    /// Contains an `Instant`, `SystemTime` and `CoarseTimeProvider`.
77
    core: MockTimeCore,
78

            
79
    /// Futures; record of every existing [`SleepFuture`], including any `Waker`
80
    ///
81
    /// Entry exists iff `SleepFuture` exists.
82
    ///
83
    /// Contains `None` if we haven't polled the future;
84
    /// `Some` if we have.
85
    ///
86
    /// We could use a `Vec` or `TiVec`
87
    /// but using a slotmap is more robust against bugs here.
88
    futures: DenseSlotMap<Id, Option<Waker>>,
89

            
90
    /// Priority queue
91
    ///
92
    /// Subset of `futures`.
93
    ///
94
    /// An entry is present iff the `Instant` is *strictly* after `State.now`,
95
    /// in which case that's when the future should be woken.
96
    ///
97
    /// `PriorityQueue` is a max-heap but we want earliest times, hence `Reverse`
98
    unready: PriorityQueue<Id, Reverse<Instant>>,
99
}
100

            
101
/// `Default` makes a `Provider` which starts at whatever the current real time is
102
impl Default for Provider {
103
83777
    fn default() -> Self {
104
83777
        Self::from_real()
105
83777
    }
106
}
107

            
108
impl Provider {
109
    /// Return a new mock time provider starting at a specified point in time
110
84769
    pub fn new(now: Instant, wallclock: SystemTime) -> Self {
111
84769
        let state = State {
112
84769
            core: MockTimeCore::new(now, wallclock),
113
84769
            futures: Default::default(),
114
84769
            unready: Default::default(),
115
84769
        };
116
84769
        Provider {
117
84769
            state: Arc::new(Mutex::new(state)),
118
84769
        }
119
84769
    }
120

            
121
    /// Return a new mock time provider starting at the current actual (non-mock) time
122
    ///
123
    /// Like any [`SimpleMockTimeProvider`], the time is frozen and only changes
124
    /// due to calls to `advance`.
125
83777
    pub fn from_real() -> Self {
126
83777
        Provider::from_wallclock(SystemTime::get())
127
83777
    }
128
    /// Return a new mock time provider starting at a specified wallclock time
129
    ///
130
    /// The monotonic time ([`Instant`]) starts at the current actual (non-mock) time.
131
    /// (Absolute values of the real monotonic time are not readily
132
    /// observable or distinguishable from Rust,
133
    /// nor can a fixed `Instant` be constructed,
134
    /// so this is usually sufficient for a reproducible test.)
135
84602
    pub fn from_wallclock(wallclock: SystemTime) -> Self {
136
84602
        Provider::new(Instant::get(), wallclock)
137
84602
    }
138

            
139
    /// Advance the simulated time by `d`
140
    ///
141
    /// This advances both the `Instant` (monotonic time)
142
    /// and `SystemTime` (wallclock time)
143
    /// by the same amount.
144
    ///
145
    /// Will wake sleeping [`SleepFuture`]s, as appropriate.
146
    ///
147
    /// Note that the tasks which were waiting on those now-expired `SleepFuture`s
148
    /// will only actually execute when they are next polled.
149
    /// `advance` does not yield to the executor or poll any futures.
150
    /// The executor will (presumably) poll those woken tasks, when it regains control.
151
    /// But the order in which the tasks run will depend on its scheduling policy,
152
    /// and might be different to the order implied by the futures' timeout values.
153
    ///
154
    /// To simulate normal time advancement, wakeups, and task activations,
155
    /// use [`MockExecutor::advance_*()`](crate::MockRuntime).
156
983760
    pub fn advance(&self, d: Duration) {
157
983760
        let mut state = self.lock();
158
983760
        state.core.advance(d);
159
983760
        state.wake_any();
160
983760
    }
161

            
162
    /// Warp the wallclock time
163
    ///
164
    /// This has no effect on any sleeping futures.
165
    /// It only affects the return value from [`.wallclock()`](Provider::wallclock).
166
1376
    pub fn jump_wallclock(&self, new_wallclock: SystemTime) {
167
1376
        self.lock().core.jump_wallclock(new_wallclock);
168
        // Really we ought to wake people up, here.
169
        // But absolutely every Rust API is wrong: none offer a way to sleep until a SystemTime.
170
        // (There might be some less-portable non-Rust APIs for that.)
171
1376
    }
172

            
173
    /// When will the next timeout occur?
174
    ///
175
    /// Returns the duration until the next [`SleepFuture`] should wake up.
176
    ///
177
    /// Advancing time by at least this amount will wake up that future,
178
    /// and any others with the same wakeup time.
179
    ///
180
    /// Will never return `Some(ZERO)`:
181
    /// any future that is supposed to wake up now (or earlier) has indeed already been woken,
182
    /// so it is no longer sleeping and isn't included in the calculation.
183
3061019
    pub fn time_until_next_timeout(&self) -> Option<Duration> {
184
3061019
        let state = self.lock();
185
3061019
        let Reverse(until) = state.unready.peek()?.1;
186
        // The invariant (see `State`) guarantees that entries in `unready` are always `> now`,
187
        // so we don't whether duration_since would panic or saturate.
188
1368599
        let d = until.duration_since(state.core.instant());
189
1368599
        Some(d)
190
3061019
    }
191

            
192
    /// Convenience function to lock the state
193
8646736
    fn lock(&self) -> MutexGuard<'_, State> {
194
8646736
        self.state.lock().expect("simple time state poisoned")
195
8646736
    }
196
}
197

            
198
impl SleepProvider for Provider {
199
    type SleepFuture = SleepFuture;
200

            
201
85383
    fn sleep(&self, d: Duration) -> SleepFuture {
202
85383
        let mut state = self.lock();
203
85383
        let until = state.core.instant() + d;
204

            
205
85383
        let id = state.futures.insert(None);
206
85383
        state.unready.push(id, Reverse(until));
207

            
208
85383
        let fut = SleepFuture {
209
85383
            id,
210
85383
            prov: self.clone(),
211
85383
        };
212

            
213
        // This sleep is now UNPOLLED, except that its time might be `<= now`:
214

            
215
        // Possibly, `until` isn't *strictly* after than `state.now`, since d might be 0.
216
        // If so, .wake_any() will restore the invariant by immediately waking.
217
85383
        state.wake_any();
218

            
219
        // This sleep is now UNPOLLED or READY, according to whether duration was 0.
220

            
221
85383
        fut
222
85383
    }
223

            
224
3935097
    fn now(&self) -> Instant {
225
3935097
        self.lock().core.instant()
226
3935097
    }
227
128923
    fn wallclock(&self) -> SystemTime {
228
128923
        self.lock().core.wallclock()
229
128923
    }
230
}
231

            
232
impl CoarseTimeProvider for Provider {
233
288310
    fn now_coarse(&self) -> CoarseInstant {
234
288310
        self.lock().core.coarse().now_coarse()
235
288310
    }
236
}
237

            
238
impl Future for SleepFuture {
239
    type Output = ();
240

            
241
82655
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
242
82655
        let mut state = self.prov.lock();
243
82655
        if let Some((_, Reverse(scheduled))) = state.unready.get(&self.id) {
244
            // Presence of this entry implies scheduled > now: we are UNPOLLED or WAITING
245
62558
            assert!(*scheduled > state.core.instant());
246
62558
            let waker = Some(cx.waker().clone());
247
            // Make this be WAITING.  (If we're re-polled, we simply drop any previous waker.)
248
62558
            *state
249
62558
                .futures
250
62558
                .get_mut(self.id)
251
62558
                .expect("polling futures entry") = waker;
252
62558
            Poll::Pending
253
        } else {
254
            // Absence implies scheduled (no longer stored) <= now: we are READY
255
20097
            Poll::Ready(())
256
        }
257
82655
    }
258
}
259

            
260
impl State {
261
    /// Restore the invariant for `unready` after `now` has been increased
262
    ///
263
    /// Ie, ensures that any sleeps which are
264
    /// WAITING/UNPOLLED except that they are `<= now`,
265
    /// are moved to state READY.
266
1069143
    fn wake_any(&mut self) {
267
        loop {
268
1105630
            match self.unready.peek() {
269
                // Keep picking off entries with scheduled <= now
270
808946
                Some((_, Reverse(scheduled))) if *scheduled <= self.core.instant() => {
271
36487
                    let (id, _) = self.unready.pop().expect("vanished");
272
                    // We can .take() the waker since this can only ever run once
273
                    // per sleep future (since it happens when we pop it from unready).
274
36487
                    let futures_entry = self.futures.get_mut(id).expect("stale unready entry");
275
36487
                    if let Some(waker) = futures_entry.take() {
276
26696
                        waker.wake();
277
26696
                    }
278
                }
279
1069143
                _ => break,
280
            }
281
        }
282
1069143
    }
283
}
284

            
285
impl Drop for SleepFuture {
286
80213
    fn drop(&mut self) {
287
80213
        let mut state = self.prov.lock();
288
80213
        let _: Option<Waker> = state.futures.remove(self.id).expect("entry vanished");
289
80213
        let _: Option<(Id, Reverse<Instant>)> = state.unready.remove(&self.id);
290
        // Now it is DROPPED.
291
80213
    }
292
}
293

            
294
#[cfg(test)]
295
mod test {
296
    // @@ begin test lint list maintained by maint/add_warning @@
297
    #![allow(clippy::bool_assert_comparison)]
298
    #![allow(clippy::clone_on_copy)]
299
    #![allow(clippy::dbg_macro)]
300
    #![allow(clippy::mixed_attributes_style)]
301
    #![allow(clippy::print_stderr)]
302
    #![allow(clippy::print_stdout)]
303
    #![allow(clippy::single_char_pattern)]
304
    #![allow(clippy::unwrap_used)]
305
    #![allow(clippy::unchecked_time_subtraction)]
306
    #![allow(clippy::useless_vec)]
307
    #![allow(clippy::needless_pass_by_value)]
308
    #![allow(clippy::string_slice)] // See arti#2571
309
    //! <!-- @@ end test lint list maintained by maint/add_warning @@ -->
310
    use super::*;
311
    use crate::task::MockExecutor;
312
    use Poll::*;
313
    use futures::poll;
314
    use humantime::parse_rfc3339;
315
    use tor_rtcompat::ToplevelBlockOn as _;
316

            
317
    fn ms(ms: u64) -> Duration {
318
        Duration::from_millis(ms)
319
    }
320

            
321
    fn run_test<FUT>(f: impl FnOnce(Provider, MockExecutor) -> FUT)
322
    where
323
        FUT: Future<Output = ()>,
324
    {
325
        let sp = Provider::new(
326
            Instant::get(), // it would have been nice to make this fixed for the test
327
            parse_rfc3339("2000-01-01T00:00:00Z").unwrap(),
328
        );
329
        let exec = MockExecutor::new();
330
        exec.block_on(f(sp, exec.clone()));
331
    }
332

            
333
    #[test]
334
    fn simple() {
335
        run_test(|sp, _exec| async move {
336
            let n1 = sp.now();
337
            let w1 = sp.wallclock();
338
            let mut f1 = sp.sleep(ms(500));
339
            let mut f2 = sp.sleep(ms(1500));
340
            assert_eq!(poll!(&mut f1), Pending);
341
            sp.advance(ms(200));
342
            assert_eq!(n1 + ms(200), sp.now());
343
            assert_eq!(w1 + ms(200), sp.wallclock());
344
            assert_eq!(poll!(&mut f1), Pending);
345
            assert_eq!(poll!(&mut f2), Pending);
346
            drop(f2);
347
            sp.jump_wallclock(w1 + ms(10_000));
348
            sp.advance(ms(300));
349
            assert_eq!(n1 + ms(500), sp.now());
350
            assert_eq!(w1 + ms(10_300), sp.wallclock());
351
            assert_eq!(poll!(&mut f1), Ready(()));
352
            let mut f0 = sp.sleep(ms(0));
353
            assert_eq!(poll!(&mut f0), Ready(()));
354
        });
355
    }
356

            
357
    #[test]
358
    fn task() {
359
        run_test(|sp, exec| async move {
360
            let st = Arc::new(Mutex::new(0_i8));
361

            
362
            exec.spawn_identified("test task", {
363
                let st = st.clone();
364
                let sp = sp.clone();
365
                async move {
366
                    *st.lock().unwrap() = 1;
367
                    sp.sleep(ms(500)).await;
368
                    *st.lock().unwrap() = 2;
369
                    sp.sleep(ms(300)).await;
370
                    *st.lock().unwrap() = 3;
371
                }
372
            });
373

            
374
            let st = move || *st.lock().unwrap();
375

            
376
            assert_eq!(st(), 0);
377
            exec.progress_until_stalled().await;
378
            assert_eq!(st(), 1);
379
            assert_eq!(sp.time_until_next_timeout(), Some(ms(500)));
380

            
381
            sp.advance(ms(500));
382

            
383
            assert_eq!(st(), 1);
384
            assert_eq!(sp.time_until_next_timeout(), None);
385
            exec.progress_until_stalled().await;
386
            assert_eq!(st(), 2);
387
            assert_eq!(sp.time_until_next_timeout(), Some(ms(300)));
388

            
389
            sp.advance(ms(500));
390
            assert_eq!(st(), 2);
391
            assert_eq!(sp.time_until_next_timeout(), None);
392
            exec.progress_until_stalled().await;
393
            assert_eq!(sp.time_until_next_timeout(), None);
394
            assert_eq!(st(), 3);
395
        });
396
    }
397
}