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