From abd4c0025539f142ec48a09e01430f7ee3b83214 Mon Sep 17 00:00:00 2001 From: Carl Lerche Date: Thu, 17 Dec 2020 15:37:08 -0800 Subject: [PATCH] time: enforce `current_thread` rt for time::pause (#3289) Pausing time is a capability added to assist with testing Tokio code dependent on time. Currently, the capability implicitly requires the current_thread runtime. This change enforces the requirement by panicking if called from a multi-threaded runtime. --- tokio/src/runtime/builder.rs | 5 +++++ tokio/src/runtime/driver.rs | 9 ++++---- tokio/src/time/clock.rs | 28 ++++++++++++++++++------- tokio/src/time/driver/tests/mod.rs | 12 +++++------ tokio/tests/time_pause.rs | 33 ++++++++++++++++++++++++++++++ 5 files changed, 70 insertions(+), 17 deletions(-) create mode 100644 tokio/tests/time_pause.rs diff --git a/tokio/src/runtime/builder.rs b/tokio/src/runtime/builder.rs index e792c7dd1..b445c20b4 100644 --- a/tokio/src/runtime/builder.rs +++ b/tokio/src/runtime/builder.rs @@ -379,6 +379,11 @@ impl Builder { fn get_cfg(&self) -> driver::Cfg { driver::Cfg { + enable_pause_time: match self.kind { + Kind::CurrentThread => true, + #[cfg(feature = "rt-multi-thread")] + Kind::MultiThread => false, + }, enable_io: self.enable_io, enable_time: self.enable_time, } diff --git a/tokio/src/runtime/driver.rs b/tokio/src/runtime/driver.rs index e89de9dc6..b89fa4fca 100644 --- a/tokio/src/runtime/driver.rs +++ b/tokio/src/runtime/driver.rs @@ -103,8 +103,8 @@ cfg_time! { pub(crate) type Clock = crate::time::Clock; pub(crate) type TimeHandle = Option; - fn create_clock() -> Clock { - crate::time::Clock::new() + fn create_clock(enable_pausing: bool) -> Clock { + crate::time::Clock::new(enable_pausing) } fn create_time_driver( @@ -131,7 +131,7 @@ cfg_not_time! { pub(crate) type Clock = (); pub(crate) type TimeHandle = (); - fn create_clock() -> Clock { + fn create_clock(_enable_pausing: bool) -> Clock { () } @@ -161,13 +161,14 @@ pub(crate) struct Resources { pub(crate) struct Cfg { pub(crate) enable_io: bool, pub(crate) enable_time: bool, + pub(crate) enable_pause_time: bool, } impl Driver { pub(crate) fn new(cfg: Cfg) -> io::Result<(Self, Resources)> { let (io_stack, io_handle, signal_handle) = create_io_stack(cfg.enable_io)?; - let clock = create_clock(); + let clock = create_clock(cfg.enable_pause_time); let (time_driver, time_handle) = create_time_driver(cfg.enable_time, io_stack, clock.clone()); diff --git a/tokio/src/time/clock.rs b/tokio/src/time/clock.rs index fab7ecaf9..a62fbe390 100644 --- a/tokio/src/time/clock.rs +++ b/tokio/src/time/clock.rs @@ -17,7 +17,7 @@ cfg_not_test_util! { } impl Clock { - pub(crate) fn new() -> Clock { + pub(crate) fn new(_enable_pausing: bool) -> Clock { Clock {} } @@ -59,6 +59,9 @@ cfg_test_util! { #[derive(Debug)] struct Inner { + /// True if the ability to pause time is enabled. + enable_pausing: bool, + /// Instant to use as the clock's base instant. base: std::time::Instant, @@ -69,14 +72,18 @@ cfg_test_util! { /// Pause time /// /// The current value of `Instant::now()` is saved and all subsequent calls - /// to `Instant::now()` until the timer wheel is checked again will return the saved value. - /// Once the timer wheel is checked, time will immediately advance to the next registered - /// `Sleep`. This is useful for running tests that depend on time. + /// to `Instant::now()` until the timer wheel is checked again will return + /// the saved value. Once the timer wheel is checked, time will immediately + /// advance to the next registered `Sleep`. This is useful for running tests + /// that depend on time. + /// + /// Pausing time requires the `current_thread` Tokio runtime. This is the + /// default runtime used by `#[tokio::test]` /// /// # Panics /// - /// Panics if time is already frozen or if called from outside of the Tokio - /// runtime. + /// Panics if time is already frozen or if called from outside of a + /// `current_thread` Tokio runtime. pub fn pause() { let clock = clock().expect("time cannot be frozen from outside the Tokio runtime"); clock.pause(); @@ -142,11 +149,12 @@ cfg_test_util! { impl Clock { /// Return a new `Clock` instance that uses the current execution context's /// source of time. - pub(crate) fn new() -> Clock { + pub(crate) fn new(enable_pausing: bool) -> Clock { let now = std::time::Instant::now(); Clock { inner: Arc::new(Mutex::new(Inner { + enable_pausing, base: now, unfrozen: Some(now), })), @@ -156,6 +164,12 @@ cfg_test_util! { pub(crate) fn pause(&self) { let mut inner = self.inner.lock().unwrap(); + if !inner.enable_pausing { + drop(inner); // avoid poisoning the lock + panic!("`time::pause()` requires the `current_thread` Tokio runtime. \ + This is the default Runtime used by `#[tokio::test]."); + } + let elapsed = inner.unfrozen.as_ref().expect("time is already frozen").elapsed(); inner.base += elapsed; inner.unfrozen = None; diff --git a/tokio/src/time/driver/tests/mod.rs b/tokio/src/time/driver/tests/mod.rs index e6af798f0..cfefed32f 100644 --- a/tokio/src/time/driver/tests/mod.rs +++ b/tokio/src/time/driver/tests/mod.rs @@ -41,7 +41,7 @@ fn model(f: impl Fn() + Send + Sync + 'static) { #[test] fn single_timer() { model(|| { - let clock = crate::time::clock::Clock::new(); + let clock = crate::time::clock::Clock::new(true); let time_source = super::ClockTime::new(clock.clone()); let inner = super::Inner::new(time_source.clone(), MockUnpark::mock()); @@ -72,7 +72,7 @@ fn single_timer() { #[test] fn drop_timer() { model(|| { - let clock = crate::time::clock::Clock::new(); + let clock = crate::time::clock::Clock::new(true); let time_source = super::ClockTime::new(clock.clone()); let inner = super::Inner::new(time_source.clone(), MockUnpark::mock()); @@ -103,7 +103,7 @@ fn drop_timer() { #[test] fn change_waker() { model(|| { - let clock = crate::time::clock::Clock::new(); + let clock = crate::time::clock::Clock::new(true); let time_source = super::ClockTime::new(clock.clone()); let inner = super::Inner::new(time_source.clone(), MockUnpark::mock()); @@ -138,7 +138,7 @@ fn reset_future() { model(|| { let finished_early = Arc::new(AtomicBool::new(false)); - let clock = crate::time::clock::Clock::new(); + let clock = crate::time::clock::Clock::new(true); let time_source = super::ClockTime::new(clock.clone()); let inner = super::Inner::new(time_source.clone(), MockUnpark::mock()); @@ -185,7 +185,7 @@ fn reset_future() { #[test] #[cfg(not(loom))] fn poll_process_levels() { - let clock = crate::time::clock::Clock::new(); + let clock = crate::time::clock::Clock::new(true); clock.pause(); let time_source = super::ClockTime::new(clock.clone()); @@ -226,7 +226,7 @@ fn poll_process_levels() { fn poll_process_levels_targeted() { let mut context = Context::from_waker(noop_waker_ref()); - let clock = crate::time::clock::Clock::new(); + let clock = crate::time::clock::Clock::new(true); clock.pause(); let time_source = super::ClockTime::new(clock.clone()); diff --git a/tokio/tests/time_pause.rs b/tokio/tests/time_pause.rs new file mode 100644 index 000000000..49a7677f5 --- /dev/null +++ b/tokio/tests/time_pause.rs @@ -0,0 +1,33 @@ +#![warn(rust_2018_idioms)] +#![cfg(feature = "full")] + +use tokio_test::assert_err; + +#[tokio::test] +async fn pause_time_in_main() { + tokio::time::pause(); +} + +#[tokio::test] +async fn pause_time_in_task() { + let t = tokio::spawn(async { + tokio::time::pause(); + }); + + t.await.unwrap(); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[should_panic] +async fn pause_time_in_main_threads() { + tokio::time::pause(); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn pause_time_in_spawn_threads() { + let t = tokio::spawn(async { + tokio::time::pause(); + }); + + assert_err!(t.await); +}