From 2b9b55810847b4c7855e3de82f432ca997600f30 Mon Sep 17 00:00:00 2001 From: Carl Lerche Date: Wed, 21 Apr 2021 15:23:35 -0700 Subject: [PATCH] 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 | 15 +--- tokio/tests/time_pause.rs | 184 +++++++++++++++++++++++++++++++++++++- 2 files changed, 184 insertions(+), 15 deletions(-) diff --git a/tokio/src/time/clock.rs b/tokio/src/time/clock.rs index 8957800cb..c5ef86be4 100644 --- a/tokio/src/time/clock.rs +++ b/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. diff --git a/tokio/tests/time_pause.rs b/tokio/tests/time_pause.rs index bc84ac578..835db12f0 100644 --- a/tokio/tests/time_pause.rs +++ b/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 { 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>>), + AfterAdvance, +} + +struct Tester { + sleep: Pin>, + state: State, + before: Option, + poll: bool, +} + +impl Future for Tester { + type Output = (); + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + 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) -> Poll { + interval.enter(|cx, mut interval| interval.poll_tick(cx)) +} + +fn ms(n: u64) -> Duration { + Duration::from_millis(n) +}