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) +}