Skip to content

Commit

Permalink
time: Allow disabling time auto-advance behavior in tests
Browse files Browse the repository at this point in the history
  • Loading branch information
Dentosal committed Nov 16, 2022
1 parent a668020 commit 2877174
Show file tree
Hide file tree
Showing 7 changed files with 258 additions and 30 deletions.
36 changes: 36 additions & 0 deletions tokio-macros/src/entry.rs
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ impl RuntimeFlavor {
struct FinalConfig {
flavor: RuntimeFlavor,
worker_threads: Option<usize>,
auto_advance: Option<bool>,
start_paused: Option<bool>,
crate_name: Option<String>,
}
Expand All @@ -36,6 +37,7 @@ struct FinalConfig {
const DEFAULT_ERROR_CONFIG: FinalConfig = FinalConfig {
flavor: RuntimeFlavor::CurrentThread,
worker_threads: None,
auto_advance: None,
start_paused: None,
crate_name: None,
};
Expand All @@ -45,6 +47,7 @@ struct Configuration {
default_flavor: RuntimeFlavor,
flavor: Option<RuntimeFlavor>,
worker_threads: Option<(usize, Span)>,
auto_advance: Option<(bool, Span)>,
start_paused: Option<(bool, Span)>,
is_test: bool,
crate_name: Option<String>,
Expand All @@ -60,6 +63,7 @@ impl Configuration {
},
flavor: None,
worker_threads: None,
auto_advance: None,
start_paused: None,
is_test,
crate_name: None,
Expand Down Expand Up @@ -98,6 +102,16 @@ impl Configuration {
Ok(())
}

fn set_auto_advance(&mut self, auto_advance: syn::Lit, span: Span) -> Result<(), syn::Error> {
if self.auto_advance.is_some() {
return Err(syn::Error::new(span, "`auto_advance` set multiple times."));
}

let auto_advance = parse_bool(auto_advance, span, "auto_advance")?;
self.auto_advance = Some((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."));
Expand Down Expand Up @@ -151,6 +165,18 @@ impl Configuration {
}
};

let auto_advance = match (flavor, self.auto_advance) {
(Threaded, Some((_, auto_advance_span))) => {
let msg = format!(
"The `auto_advance` option requires the `current_thread` runtime flavor. Use `#[{}(flavor = \"current_thread\")]`",
self.macro_name(),
);
return Err(syn::Error::new(auto_advance_span, msg));
}
(CurrentThread, Some((auto_advance, _))) => Some(auto_advance),
(_, None) => None,
};

let start_paused = match (flavor, self.start_paused) {
(Threaded, Some((_, start_paused_span))) => {
let msg = format!(
Expand All @@ -167,6 +193,7 @@ impl Configuration {
crate_name: self.crate_name.clone(),
flavor,
worker_threads,
auto_advance,
start_paused,
})
}
Expand Down Expand Up @@ -268,6 +295,12 @@ fn build_config(
syn::spanned::Spanned::span(&namevalue.lit),
)?;
}
"auto_advance" => {
config.set_auto_advance(
namevalue.lit.clone(),
syn::spanned::Spanned::span(&namevalue.lit),
)?;
}
"start_paused" => {
config.set_start_paused(
namevalue.lit.clone(),
Expand Down Expand Up @@ -369,6 +402,9 @@ fn parse_knobs(mut input: syn::ItemFn, is_test: bool, config: FinalConfig) -> To
if let Some(v) = config.worker_threads {
rt = quote! { #rt.worker_threads(#v) };
}
if let Some(v) = config.auto_advance {
rt = quote! { #rt.auto_advance(#v) };
}
if let Some(v) = config.start_paused {
rt = quote! { #rt.start_paused(#v) };
}
Expand Down
50 changes: 44 additions & 6 deletions tokio/src/runtime/builder.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
use crate::runtime::handle::Handle;
use crate::runtime::{blocking, driver, Callback, Runtime};
use crate::util::rand::{RngSeed, RngSeedGenerator};
use crate::time::PauseSettings;

use std::fmt;
use std::io;
Expand Down Expand Up @@ -48,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.
auto_advance: bool,

/// Whether or not the clock should start paused.
start_paused: bool,

Expand Down Expand Up @@ -181,6 +185,7 @@ cfg_unstable! {

pub(crate) type ThreadNameFn = std::sync::Arc<dyn Fn() -> String + Send + Sync + 'static>;

#[derive(Clone, Copy)]
pub(crate) enum Kind {
CurrentThread,
#[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))]
Expand Down Expand Up @@ -232,6 +237,14 @@ impl Builder {
// Time defaults to "off"
enable_time: false,

// By default time can be paused and will auto-advance,
// but only on the `CurrentThread` runtime
auto_advance: match &kind {
Kind::CurrentThread => true,
#[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))]
Kind::MultiThread => false,
},

// The clock starts not-paused
start_paused: false,

Expand Down Expand Up @@ -639,14 +652,18 @@ impl Builder {

fn get_cfg(&self) -> driver::Cfg {
driver::Cfg {
enable_pause_time: match self.kind {
Kind::CurrentThread => true,
#[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))]
Kind::MultiThread => false,
},
enable_io: self.enable_io,
enable_time: self.enable_time,
start_paused: self.start_paused,
time_pausing: PauseSettings {
enabled: self.enable_time
&& match &self.kind {
Kind::CurrentThread => true,
#[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))]
Kind::MultiThread => false,
},
auto_advance: self.auto_advance,
start_paused: self.start_paused,
},
}
}

Expand Down Expand Up @@ -966,6 +983,27 @@ cfg_time! {

cfg_test_util! {
impl Builder {
/// Controls if the runtime's clock auto-advance behavior when paused.
///
/// Pausing time requires the current-thread runtime; construction of
/// the runtime will panic otherwise.
///
/// # Examples
///
/// ```
/// use tokio::runtime;
///
/// let rt = runtime::Builder::new_current_thread()
/// .enable_time()
/// .auto_advance(false)
/// .build()
/// .unwrap();
/// ```
pub fn auto_advance(&mut self, auto_advance: bool) -> &mut Self {
self.auto_advance = auto_advance;
self
}

/// Controls if the runtime's clock starts paused or advancing.
///
/// Pausing time requires the current-thread runtime; construction of
Expand Down
22 changes: 10 additions & 12 deletions tokio/src/runtime/driver.rs
Original file line number Diff line number Diff line change
Expand Up @@ -34,15 +34,19 @@ pub(crate) struct Handle {
pub(crate) struct Cfg {
pub(crate) enable_io: bool,
pub(crate) enable_time: bool,
pub(crate) enable_pause_time: bool,
pub(crate) start_paused: bool,
pub(crate) time_pausing: PauseSettings,
}

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)?;

let clock = create_clock(cfg.enable_pause_time, cfg.start_paused);
#[cfg(feature = "time")]
#[cfg_attr(docsrs, doc(cfg(feature = "time")))]
let clock = crate::time::Clock::new(cfg.time_pausing);

#[cfg(not(feature = "time"))]
let clock = ();

let (time_driver, time_handle) =
create_time_driver(cfg.enable_time, io_stack, clock.clone());
Expand Down Expand Up @@ -280,11 +284,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)
}

pub(crate) type PauseSettings = crate::time::PauseSettings;

fn create_time_driver(
enable: bool,
io_stack: IoStack,
Expand Down Expand Up @@ -327,12 +328,9 @@ cfg_not_time! {
type TimeDriver = IoStack;

pub(crate) type Clock = ();
pub(crate) type PauseSettings = ();
pub(crate) type TimeHandle = ();

fn create_clock(_enable_pausing: bool, _start_paused: bool) -> Clock {
()
}

fn create_time_driver(
_enable: bool,
io_stack: IoStack,
Expand Down
6 changes: 4 additions & 2 deletions tokio/src/runtime/time/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -230,8 +230,10 @@ impl Driver {
// yield in `Runtime::block_on`). In this case, we don't
// advance the clock.
if !handle.did_wake() {
// Simulate advancing time
clock.advance(duration);
// Simulate advancing time if enabled
if clock.auto_advance() {
clock.advance(duration);
}
}
} else {
self.park.park_timeout(rt_handle, duration);
Expand Down
Loading

0 comments on commit 2877174

Please sign in to comment.