time: Allow disabling time auto-advance behavior in tests

Co-authored-by: Hannes Karppila <hannes.karppila@gmail.com>

tokio-macros/src/entry.rs

@@ -28,6 +28,7 @@ impl RuntimeFlavor {
 struct FinalConfig {
     flavor: RuntimeFlavor,
     worker_threads: Option<usize>,
+    time_auto_advance: Option<bool>,
     start_paused: Option<bool>,
     crate_name: Option<Path>,
 }
@@ -36,6 +37,7 @@ struct FinalConfig {
 const DEFAULT_ERROR_CONFIG: FinalConfig = FinalConfig {
     flavor: RuntimeFlavor::CurrentThread,
     worker_threads: None,
+    time_auto_advance: None,
     start_paused: None,
     crate_name: None,
 };
@@ -45,6 +47,7 @@ struct Configuration {
     default_flavor: RuntimeFlavor,
     flavor: Option<RuntimeFlavor>,
     worker_threads: Option<(usize, Span)>,
+    time_auto_advance: Option<(bool, Span)>,
     start_paused: Option<(bool, Span)>,
     is_test: bool,
     crate_name: Option<Path>,
@@ -60,6 +63,7 @@ impl Configuration {
             },
             flavor: None,
             worker_threads: None,
+            time_auto_advance: None,
             start_paused: None,
             is_test,
             crate_name: None,
@@ -98,6 +102,23 @@ impl Configuration {
         Ok(())
     }
 
+    fn set_time_auto_advance(
+        &mut self,
+        time_auto_advance: syn::Lit,
+        span: Span,
+    ) -> Result<(), syn::Error> {
+        if self.time_auto_advance.is_some() {
+            return Err(syn::Error::new(
+                span,
+                "`time_auto_advance` set multiple times.",
+            ));
+        }
+
+        let time_auto_advance = parse_bool(time_auto_advance, span, "time_auto_advance")?;
+        self.time_auto_advance = Some((time_auto_advance, span));
+        Ok(())
+    }
+
     fn set_start_paused(&mut self, start_paused: syn::Lit, span: Span) -> Result<(), syn::Error> {
         if self.start_paused.is_some() {
             return Err(syn::Error::new(span, "`start_paused` set multiple times."));
@@ -151,6 +172,18 @@ impl Configuration {
             }
         };
 
+        let time_auto_advance = match (flavor, self.time_auto_advance) {
+            (F::Threaded, Some((_, time_auto_advance_span))) => {
+                let msg = format!(
+                    "The `time_auto_advance` option requires the `current_thread` runtime flavor. Use `#[{}(flavor = \"current_thread\")]`",
+                    self.macro_name(),
+                );
+                return Err(syn::Error::new(time_auto_advance_span, msg));
+            }
+            (F::CurrentThread, Some((time_auto_advance, _))) => Some(time_auto_advance),
+            (_, None) => None,
+        };
+
         let start_paused = match (flavor, self.start_paused) {
             (F::Threaded, Some((_, start_paused_span))) => {
                 let msg = format!(
@@ -167,6 +200,7 @@ impl Configuration {
             crate_name: self.crate_name.clone(),
             flavor,
             worker_threads,
+            time_auto_advance,
             start_paused,
         })
     }
@@ -265,6 +299,10 @@ fn build_config(
                     "flavor" => {
                         config.set_flavor(lit.clone(), syn::spanned::Spanned::span(lit))?;
                     }
+                    "time_auto_advance" => {
+                        config
+                            .set_time_auto_advance(lit.clone(), syn::spanned::Spanned::span(lit))?;
+                    }
                     "start_paused" => {
                         config.set_start_paused(lit.clone(), syn::spanned::Spanned::span(lit))?;
                     }
@@ -307,7 +345,11 @@ fn build_config(
                         format!("The `{}` attribute requires an argument.", name)
                     }
                     name => {
-                        format!("Unknown attribute {} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`", name)
+                        format!(
+                            concat!("Unknown attribute {} is specified; expected one of:",
+                                    " `flavor`, `worker_threads`, `time_auto_advance`, `start_paused`, `crate`"),
+                            name
+                        )
                     }
                 };
                 return Err(syn::Error::new_spanned(path, msg));
@@ -356,6 +398,9 @@ fn parse_knobs(mut input: ItemFn, is_test: bool, config: FinalConfig) -> TokenSt
     if let Some(v) = config.worker_threads {
         rt = quote_spanned! {last_stmt_start_span=> #rt.worker_threads(#v) };
     }
+    if let Some(v) = config.time_auto_advance {
+        rt = quote_spanned! {last_stmt_start_span=> #rt.time_auto_advance(#v) };
+    }
     if let Some(v) = config.start_paused {
         rt = quote_spanned! {last_stmt_start_span=> #rt.start_paused(#v) };
     }

tokio/src/runtime/builder.rs

@@ -49,6 +49,9 @@ pub struct Builder {
     /// Whether or not to enable the time driver
     enable_time: bool,
 
+    /// Whether or not clock should auto-advance when sleeping while time is paused.
+    enable_time_auto_advance: bool,
+
     /// Whether or not the clock should start paused.
     start_paused: bool,
 
@@ -271,6 +274,9 @@ impl Builder {
             // Time defaults to "off"
             enable_time: false,
 
+            // Auto-advance defaults to "on"
+            enable_time_auto_advance: true,
+
             // The clock starts not-paused
             start_paused: false,
 
@@ -715,6 +721,7 @@ impl Builder {
             },
             enable_io: self.enable_io,
             enable_time: self.enable_time,
+            enable_time_auto_advance: self.enable_time_auto_advance,
             start_paused: self.start_paused,
             nevents: self.nevents,
         }
@@ -1195,6 +1202,26 @@ cfg_time! {
             self.enable_time = true;
             self
         }
+
+        /// Sets auto-advance feature for paused timer.
+        ///
+        /// Read more at [crate::time::pause] doc.
+        ///
+        /// # Examples
+        ///
+        /// ```
+        /// use tokio::runtime;
+        ///
+        /// let rt = runtime::Builder::new_current_thread()
+        ///     .enable_time()
+        ///     .time_auto_advance(false)
+        ///     .build()
+        ///     .unwrap();
+        /// ```
+        pub fn time_auto_advance(&mut self, time_auto_advance: bool) -> &mut Self {
+            self.enable_time_auto_advance = time_auto_advance;
+            self
+        }
     }
 }
 

tokio/src/runtime/driver.rs

@@ -38,6 +38,7 @@ pub(crate) struct Cfg {
     pub(crate) enable_io: bool,
     pub(crate) enable_time: bool,
     pub(crate) enable_pause_time: bool,
+    pub(crate) enable_time_auto_advance: bool,
     pub(crate) start_paused: bool,
     pub(crate) nevents: usize,
 }
@@ -46,7 +47,11 @@ impl Driver {
     pub(crate) fn new(cfg: Cfg) -> io::Result<(Self, Handle)> {
         let (io_stack, io_handle, signal_handle) = create_io_stack(cfg.enable_io, cfg.nevents)?;
 
-        let clock = create_clock(cfg.enable_pause_time, cfg.start_paused);
+        let clock = create_clock(
+            cfg.enable_pause_time,
+            cfg.enable_time_auto_advance,
+            cfg.start_paused,
+        );
 
         let (time_driver, time_handle) = create_time_driver(cfg.enable_time, io_stack, &clock);
 
@@ -298,8 +303,8 @@ cfg_time! {
     pub(crate) type Clock = crate::time::Clock;
     pub(crate) type TimeHandle = Option<crate::runtime::time::Handle>;
 
-    fn create_clock(enable_pausing: bool, start_paused: bool) -> Clock {
-        crate::time::Clock::new(enable_pausing, start_paused)
+    fn create_clock(enable_pausing: bool, enable_auto_advance: bool, start_paused: bool) -> Clock {
+        crate::time::Clock::new(enable_pausing, enable_auto_advance, start_paused)
     }
 
     fn create_time_driver(

tokio/src/time/clock.rs

@@ -81,6 +81,9 @@ cfg_test_util! {
         /// True if the ability to pause time is enabled.
         enable_pausing: bool,
 
+        /// True if auto-advance features is enabled.
+        enable_auto_advance: bool,
+
         /// Instant to use as the clock's base instant.
         base: std::time::Instant,
 
@@ -119,13 +122,16 @@ cfg_test_util! {
     ///
     /// # Auto-advance
     ///
-    /// If time is paused and the runtime has no work to do, the clock is
+    /// By default if time is paused and the runtime has no work to do, the clock is
     /// auto-advanced to the next pending timer. This means that [`Sleep`] or
     /// other timer-backed primitives can cause the runtime to advance the
     /// current time when awaited.
     ///
+    /// This behavior could be disabled with [`set_time_auto_advance`].
+    ///
     /// [`Sleep`]: crate::time::Sleep
     /// [`advance`]: crate::time::advance
+    /// [`set_time_auto_advance`]: crate::time::set_time_auto_advance
     #[track_caller]
     pub fn pause() {
         with_clock(|maybe_clock| {
@@ -190,7 +196,7 @@ cfg_test_util! {
     ///
     /// # Auto-advance
     ///
-    /// If the time is paused and there is no work to do, the runtime advances
+    /// By default if the time is paused and there is no work to do, the runtime advances
     /// time to the next timer. See [`pause`](pause#auto-advance) for more
     /// details.
     ///
@@ -208,6 +214,35 @@ cfg_test_util! {
         crate::task::yield_now().await;
     }
 
+    /// Controls time auto-advance behavior.
+    ///
+    /// When time is paused, by default the runtime automatically advances
+    /// time to the next timer. See [`pause`](pause#auto-advance) for more details.
+    ///
+    /// This function allows enabling and disabling the auto-advance behavior.
+    /// When auto-advance feature is disabled, sleeping will wait indefinitely
+    /// until it's re-enabled, or the time is manually advanced using [`advance`].
+    /// This means that if all tasks call `sleep` simultaneously, the program will
+    /// deadlock.
+    ///
+    /// # Panics
+    ///
+    /// If called from outside of the Tokio runtime.
+    ///
+    /// [`sleep`]: fn@crate::time::sleep
+    /// [`advance`]: crate::time::advance
+    pub fn set_time_auto_advance(enable: bool) {
+        with_clock(|maybe_clock| {
+            let clock = match maybe_clock {
+                Some(clock) => clock,
+                None => return Err("time cannot be frozen from outside the Tokio runtime"),
+            };
+
+            clock.set_time_auto_advance(enable);
+            Ok(())
+        })
+    }
+
     /// Returns the current instant, factoring in frozen time.
     pub(crate) fn now() -> Instant {
         if !DID_PAUSE_CLOCK.load(Ordering::Acquire) {
@@ -226,12 +261,13 @@ cfg_test_util! {
     impl Clock {
         /// Returns a new `Clock` instance that uses the current execution context's
         /// source of time.
-        pub(crate) fn new(enable_pausing: bool, start_paused: bool) -> Clock {
+        pub(crate) fn new(enable_pausing: bool, enable_auto_advance: bool, start_paused: bool) -> Clock {
             let now = std::time::Instant::now();
 
             let clock = Clock {
                 inner: Mutex::new(Inner {
                     enable_pausing,
+                    enable_auto_advance,
                     base: now,
                     unfrozen: Some(now),
                     auto_advance_inhibit_count: 0,
@@ -269,6 +305,18 @@ cfg_test_util! {
             Ok(())
         }
 
+        pub(crate) fn set_time_auto_advance(&self, enable: bool) {
+            let mut inner = self.inner.lock();
+
+            if !inner.enable_pausing {
+                drop(inner); // avoid poisoning the lock
+                panic!("`time::set_auto_advance()` requires the `current_thread` Tokio runtime. \
+                        This is the default Runtime used by `#[tokio::test].");
+            }
+
+            inner.enable_auto_advance = enable;
+        }
+
         /// Temporarily stop auto-advancing the clock (see `tokio::time::pause`).
         pub(crate) fn inhibit_auto_advance(&self) {
             let mut inner = self.inner.lock();
@@ -282,7 +330,7 @@ cfg_test_util! {
 
         pub(crate) fn can_auto_advance(&self) -> bool {
             let inner = self.inner.lock();
-            inner.unfrozen.is_none() && inner.auto_advance_inhibit_count == 0
+            inner.enable_auto_advance && inner.unfrozen.is_none() && inner.auto_advance_inhibit_count == 0
         }
 
         pub(crate) fn advance(&self, duration: Duration) -> Result<(), &'static str> {

tokio/src/time/mod.rs

@@ -87,7 +87,7 @@
 mod clock;
 pub(crate) use self::clock::Clock;
 #[cfg(feature = "test-util")]
-pub use clock::{advance, pause, resume};
+pub use clock::{advance, pause, resume, set_time_auto_advance};
 
 pub mod error;