tor_rtmock/
sleep_runtime.rs

Help
1//! Declare MockSleepRuntime.
2
3use pin_project::pin_project;
4use tracing::trace;
5
6use crate::time::MockSleepProvider;
7
8use crate::util::impl_runtime_prelude::*;
9
10/// A deprecated wrapper Runtime that overrides SleepProvider for the
11/// underlying runtime.
12///
13/// ### Deprecated
14///
15/// The [`MockSleepProvider`] used here has some limitations.
16/// See its documentation for more information.
17/// Use [`MockRuntime`](crate::MockRuntime) for new tests.
18#[derive(Clone, Debug, Deftly)]
19#[derive_deftly(SomeMockRuntime)]
20// TODO #1885, see MockSleepProvider's cfg_attr deprecated
21#[cfg_attr(not(test), deprecated(since = "0.29.0"))]
22pub struct MockSleepRuntime<R: Runtime> {
23    /// The underlying runtime. Most calls get delegated here.
24    #[deftly(mock(task, net))]
25    #[deftly(mock(toplevel_where = "R: ToplevelBlockOn"))]
26    runtime: R,
27    /// A MockSleepProvider.  Time-related calls get delegated here.
28    #[deftly(mock(sleep))]
29    sleep: MockSleepProvider,
30}
31
32impl<R: Runtime> MockSleepRuntime<R> {
33    /// Create a new runtime that wraps `runtime`, but overrides
34    /// its view of time with a [`MockSleepProvider`].
35    pub fn new(runtime: R) -> Self {
36        let sleep = MockSleepProvider::new(SystemTime::now());
37        MockSleepRuntime { runtime, sleep }
38    }
39
40    /// Return a reference to the underlying runtime.
41    pub fn inner(&self) -> &R {
42        &self.runtime
43    }
44
45    /// Return a reference to the [`MockSleepProvider`]
46    pub fn mock_sleep(&self) -> &MockSleepProvider {
47        &self.sleep
48    }
49
50    /// See [`MockSleepProvider::advance()`]
51    pub async fn advance(&self, dur: Duration) {
52        self.sleep.advance(dur).await;
53    }
54    /// See [`MockSleepProvider::jump_to()`]
55    pub fn jump_to(&self, new_wallclock: SystemTime) {
56        self.sleep.jump_to(new_wallclock);
57    }
58    /// Run a future under mock time, advancing time forward where necessary until it completes.
59    /// Users of this function should read the whole of this documentation before using!
60    ///
61    /// **NOTE** Instead of using this, consider [`MockRuntime`](crate::MockRuntime),
62    /// which will fully isolate the test case
63    /// (albeit at the cost of demanding manual management of the simulated time).
64    ///
65    /// The returned future will run `fut`, expecting it to create `Sleeping` futures (as returned
66    /// by `MockSleepProvider::sleep()` and similar functions). When all such created futures have
67    /// been polled (indicating the future is waiting on them), time will be advanced in order that
68    /// the first (or only) of said futures returns `Ready`. This process then repeats until `fut`
69    /// returns `Ready` itself (as in, the returned wrapper future will wait for all created
70    /// `Sleeping` futures to be polled, and advance time again).
71    ///
72    /// **Note:** The above described algorithm interacts poorly with futures that spawn
73    /// asynchronous background tasks, or otherwise expect work to complete in the background
74    /// before time is advanced. These futures will need to make use of the
75    /// `SleepProvider::block_advance` (and similar) APIs in order to prevent time advancing while
76    /// said tasks complete; see the documentation for those APIs for more detail.
77    ///
78    /// # Panics
79    ///
80    /// Panics if another `WaitFor` future is already running. (If two ran simultaneously, they
81    /// would both try and advance the same mock time clock, which would be bad.)
82    pub fn wait_for<F: futures::Future>(&self, fut: F) -> WaitFor<F> {
83        assert!(
84            !self.sleep.has_waitfor_waker(),
85            "attempted to call MockSleepRuntime::wait_for while another WaitFor is active"
86        );
87        WaitFor {
88            sleep: self.sleep.clone(),
89            fut,
90        }
91    }
92}
93
94/// A future that advances time until another future is ready to complete.
95#[pin_project]
96pub struct WaitFor<F> {
97    /// A reference to the sleep provider that's simulating time for us.
98    #[pin]
99    sleep: MockSleepProvider,
100    /// The future that we're waiting for.
101    #[pin]
102    fut: F,
103}
104
105use std::pin::Pin;
106use std::task::{Context, Poll};
107
108impl<F: Future> Future for WaitFor<F> {
109    type Output = F::Output;
110
111    #[allow(clippy::cognitive_complexity)]
112    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
113        trace!("waitfor poll");
114        let mut this = self.project();
115        this.sleep.register_waitfor_waker(cx.waker().clone());
116
117        if let Poll::Ready(r) = this.fut.poll(cx) {
118            trace!("waitfor done!");
119            this.sleep.clear_waitfor_waker();
120            return Poll::Ready(r);
121        }
122        trace!("waitfor poll complete");
123
124        if this.sleep.should_advance() {
125            if let Some(duration) = this.sleep.time_until_next_timeout() {
126                trace!("Advancing by {:?}", duration);
127                this.sleep.advance_noyield(duration);
128            } else {
129                // If we get here, something's probably wedged and the test isn't going to complete
130                // anyway: we were expecting to advance in order to make progress, but we can't.
131                // If we don't panic, the test will just run forever, which is really annoying, so
132                // just panic and fail quickly.
133                panic!("WaitFor told to advance, but didn't have any duration to advance by");
134            }
135        } else {
136            trace!("waiting for sleepers to advance");
137        }
138        Poll::Pending
139    }
140}
tor_rtmock/sleep_runtime.rs

tor_rtmock/
sleep_runtime.rs
