time: prevent time::advance from going too far (#3712)

Previously, `time::advance` would set the mocked clock forward the
requested amount, then yield. However, if there was no work ready to
perform immediately, this would result in advancing to the next expiring
sleep.

Now, `time::advance(...)` will unblock at the requested time. The
difference between `time::advance(...)` and `time::sleep(...)` is a bit
fuzzy. The main difference is `time::sleep(...)` operates on the current
task and `time::advance(...)` operates at the runtime level.

Fixes #3710

tokio/src/time/clock.rs

@@ -120,22 +120,11 @@ cfg_test_util! {
     /// Panics if time is not frozen or if called from outside of the Tokio
     /// runtime.
     pub async fn advance(duration: Duration) {
-        use crate::future::poll_fn;
-        use std::task::Poll;
-
         let clock = clock().expect("time cannot be frozen from outside the Tokio runtime");
+        let until = clock.now() + duration;
         clock.advance(duration);
 
-        let mut yielded = false;
-        poll_fn(|cx| {
-            if yielded {
-                Poll::Ready(())
-            } else {
-                yielded = true;
-                cx.waker().wake_by_ref();
-                Poll::Pending
-            }
-        }).await;
+        crate::time::sleep_until(until).await;
     }
 
     /// Return the current instant, factoring in frozen time.

tokio/tests/time_pause.rs

@@ -3,8 +3,29 @@
 
 use rand::SeedableRng;
 use rand::{rngs::StdRng, Rng};
-use tokio::time::{self, Duration, Instant};
-use tokio_test::assert_err;
+use tokio::time::{self, Duration, Instant, Sleep};
+use tokio_test::{assert_err, assert_pending, assert_ready_eq, task};
+
+use std::{
+    future::Future,
+    pin::Pin,
+    task::{Context, Poll},
+};
+
+macro_rules! assert_elapsed {
+    ($now:expr, $ms:expr) => {{
+        let elapsed = $now.elapsed();
+        let lower = ms($ms);
+
+        // Handles ms rounding
+        assert!(
+            elapsed >= lower && elapsed <= lower + ms(1),
+            "actual = {:?}, expected = {:?}",
+            elapsed,
+            lower
+        );
+    }};
+}
 
 #[tokio::test]
 async fn pause_time_in_main() {
@@ -57,3 +78,162 @@ async fn paused_time_stress_run() -> Vec<Duration> {
 
     times
 }
+
+#[tokio::test(start_paused = true)]
+async fn advance_after_poll() {
+    time::sleep(ms(1)).await;
+
+    let start = Instant::now();
+
+    let mut sleep = task::spawn(time::sleep_until(start + ms(300)));
+
+    assert_pending!(sleep.poll());
+
+    let before = Instant::now();
+    time::advance(ms(100)).await;
+    assert_elapsed!(before, 100);
+
+    assert_pending!(sleep.poll());
+}
+
+#[tokio::test(start_paused = true)]
+async fn sleep_no_poll() {
+    let start = Instant::now();
+
+    // TODO: Skip this
+    time::advance(ms(1)).await;
+
+    let mut sleep = task::spawn(time::sleep_until(start + ms(300)));
+
+    let before = Instant::now();
+    time::advance(ms(100)).await;
+    assert_elapsed!(before, 100);
+
+    assert_pending!(sleep.poll());
+}
+
+enum State {
+    Begin,
+    AwaitingAdvance(Pin<Box<dyn Future<Output = ()>>>),
+    AfterAdvance,
+}
+
+struct Tester {
+    sleep: Pin<Box<Sleep>>,
+    state: State,
+    before: Option<Instant>,
+    poll: bool,
+}
+
+impl Future for Tester {
+    type Output = ();
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        match &mut self.state {
+            State::Begin => {
+                if self.poll {
+                    assert_pending!(self.sleep.as_mut().poll(cx));
+                }
+                self.before = Some(Instant::now());
+                let advance_fut = Box::pin(time::advance(ms(100)));
+                self.state = State::AwaitingAdvance(advance_fut);
+                self.poll(cx)
+            }
+            State::AwaitingAdvance(ref mut advance_fut) => match advance_fut.as_mut().poll(cx) {
+                Poll::Pending => Poll::Pending,
+                Poll::Ready(()) => {
+                    self.state = State::AfterAdvance;
+                    self.poll(cx)
+                }
+            },
+            State::AfterAdvance => {
+                assert_elapsed!(self.before.unwrap(), 100);
+
+                assert_pending!(self.sleep.as_mut().poll(cx));
+
+                Poll::Ready(())
+            }
+        }
+    }
+}
+
+#[tokio::test(start_paused = true)]
+async fn sleep_same_task() {
+    let start = Instant::now();
+
+    // TODO: Skip this
+    time::advance(ms(1)).await;
+
+    let sleep = Box::pin(time::sleep_until(start + ms(300)));
+
+    Tester {
+        sleep,
+        state: State::Begin,
+        before: None,
+        poll: true,
+    }
+    .await;
+}
+
+#[tokio::test(start_paused = true)]
+async fn sleep_same_task_no_poll() {
+    let start = Instant::now();
+
+    // TODO: Skip this
+    time::advance(ms(1)).await;
+
+    let sleep = Box::pin(time::sleep_until(start + ms(300)));
+
+    Tester {
+        sleep,
+        state: State::Begin,
+        before: None,
+        poll: false,
+    }
+    .await;
+}
+
+#[tokio::test(start_paused = true)]
+async fn interval() {
+    let start = Instant::now();
+
+    // TODO: Skip this
+    time::advance(ms(1)).await;
+
+    let mut i = task::spawn(time::interval_at(start, ms(300)));
+
+    assert_ready_eq!(poll_next(&mut i), start);
+    assert_pending!(poll_next(&mut i));
+
+    let before = Instant::now();
+    time::advance(ms(100)).await;
+    assert_elapsed!(before, 100);
+    assert_pending!(poll_next(&mut i));
+
+    let before = Instant::now();
+    time::advance(ms(200)).await;
+    assert_elapsed!(before, 200);
+    assert_ready_eq!(poll_next(&mut i), start + ms(300));
+    assert_pending!(poll_next(&mut i));
+
+    let before = Instant::now();
+    time::advance(ms(400)).await;
+    assert_elapsed!(before, 400);
+    assert_ready_eq!(poll_next(&mut i), start + ms(600));
+    assert_pending!(poll_next(&mut i));
+
+    let before = Instant::now();
+    time::advance(ms(500)).await;
+    assert_elapsed!(before, 500);
+    assert_ready_eq!(poll_next(&mut i), start + ms(900));
+    assert_ready_eq!(poll_next(&mut i), start + ms(1200));
+    assert_pending!(poll_next(&mut i));
+}
+
+fn poll_next(interval: &mut task::Spawn<time::Interval>) -> Poll<Instant> {
+    interval.enter(|cx, mut interval| interval.poll_tick(cx))
+}
+
+fn ms(n: u64) -> Duration {
+    Duration::from_millis(n)
+}