diff --git a/tests-build/tests/fail/macros_spawn_no_send.rs b/tests-build/tests/fail/macros_spawn_no_send.rs new file mode 100644 index 00000000000..80f639ff495 --- /dev/null +++ b/tests-build/tests/fail/macros_spawn_no_send.rs @@ -0,0 +1,17 @@ +use std::rc::Rc; + +use tests_build::tokio; + +async fn send_f() {} + +#[tokio::test(spawn = true)] +async fn no_send_test() { + let _value = Rc::new(0); + send_f().await; +} + +#[tokio::main(spawn = true)] +async fn no_send_main() { + let _value = Rc::new(0); + send_f().await; +} diff --git a/tests-build/tests/macros.rs b/tests-build/tests/macros.rs index 0a180dfb74f..63ffa26867c 100644 --- a/tests-build/tests/macros.rs +++ b/tests-build/tests/macros.rs @@ -14,6 +14,9 @@ fn compile_fail_full() { #[cfg(feature = "full")] t.compile_fail("tests/fail/macros_invalid_input.rs"); + #[cfg(feature = "full")] + t.compile_fail("tests/fail/macros_spawn_no_send.rs"); + #[cfg(feature = "full")] t.compile_fail("tests/fail/macros_dead_code.rs"); diff --git a/tokio-macros/src/entry.rs b/tokio-macros/src/entry.rs index bcf64aef9c8..2842fb8f320 100644 --- a/tokio-macros/src/entry.rs +++ b/tokio-macros/src/entry.rs @@ -29,6 +29,7 @@ struct FinalConfig { flavor: RuntimeFlavor, worker_threads: Option, start_paused: Option, + spawn: Option, crate_name: Option, } @@ -37,6 +38,7 @@ const DEFAULT_ERROR_CONFIG: FinalConfig = FinalConfig { flavor: RuntimeFlavor::CurrentThread, worker_threads: None, start_paused: None, + spawn: None, crate_name: None, }; @@ -46,6 +48,7 @@ struct Configuration { flavor: Option, worker_threads: Option<(usize, Span)>, start_paused: Option<(bool, Span)>, + spawn: Option, is_test: bool, crate_name: Option, } @@ -61,6 +64,7 @@ impl Configuration { flavor: None, worker_threads: None, start_paused: None, + spawn: None, is_test, crate_name: None, } @@ -108,6 +112,16 @@ impl Configuration { Ok(()) } + fn set_spawn(&mut self, spawn: syn::Lit, span: Span) -> Result<(), syn::Error> { + if self.spawn.is_some() { + return Err(syn::Error::new(span, "`spawn` set multiple times.")); + } + + let spawn = parse_bool(spawn, span, "spawn")?; + self.spawn = Some(spawn); + Ok(()) + } + fn set_crate_name(&mut self, name: syn::Lit, span: Span) -> Result<(), syn::Error> { if self.crate_name.is_some() { return Err(syn::Error::new(span, "`crate` set multiple times.")); @@ -168,6 +182,7 @@ impl Configuration { flavor, worker_threads, start_paused, + spawn: self.spawn, }) } } @@ -268,6 +283,9 @@ fn build_config( "start_paused" => { config.set_start_paused(lit.clone(), syn::spanned::Spanned::span(lit))?; } + "spawn" => { + config.set_spawn(lit.clone(), syn::spanned::Spanned::span(lit))?; + } "core_threads" => { let msg = "Attribute `core_threads` is renamed to `worker_threads`"; return Err(syn::Error::new_spanned(namevalue, msg)); @@ -303,11 +321,11 @@ fn build_config( macro_name ) } - "flavor" | "worker_threads" | "start_paused" => { + "flavor" | "worker_threads" | "start_paused" | "spawn" => { format!("The `{}` attribute requires an argument.", name) } name => { - format!("Unknown attribute {} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`", name) + format!("Unknown attribute {} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `spawn`, `crate`", name) } }; return Err(syn::Error::new_spanned(path, msg)); @@ -369,14 +387,32 @@ fn parse_knobs(mut input: ItemFn, is_test: bool, config: FinalConfig) -> TokenSt }; let body_ident = quote! { body }; - let last_block = quote_spanned! {last_stmt_end_span=> - #[allow(clippy::expect_used, clippy::diverging_sub_expression)] - { - return #rt - .enable_all() - .build() - .expect("Failed building the Runtime") - .block_on(#body_ident); + let spawn = config.spawn.unwrap_or(false); + let last_block = if spawn { + quote_spanned! {last_stmt_end_span=> + #[allow(clippy::expect_used, clippy::diverging_sub_expression)] + { + let runtime = #rt + .enable_all() + .build() + .expect("Failed building the Runtime"); + let handle = runtime.spawn(#body_ident); + runtime.block_on(async move { + // propagate the panic + handle.await.unwrap() + }); + } + } + } else { + quote_spanned! {last_stmt_end_span=> + #[allow(clippy::expect_used, clippy::diverging_sub_expression)] + { + return #rt + .enable_all() + .build() + .expect("Failed building the Runtime") + .block_on(#body_ident); + } } }; @@ -391,7 +427,7 @@ fn parse_knobs(mut input: ItemFn, is_test: bool, config: FinalConfig) -> TokenSt // // We don't do this for the main function as it should only be used once so // there will be no benefit. - let body = if is_test { + let body = if is_test && !spawn { let output_type = match &input.sig.output { // For functions with no return value syn doesn't print anything, // but that doesn't work as `Output` for our boxed `Future`, so diff --git a/tokio-macros/src/lib.rs b/tokio-macros/src/lib.rs index 919c4ac0ba9..b7869c1b632 100644 --- a/tokio-macros/src/lib.rs +++ b/tokio-macros/src/lib.rs @@ -414,6 +414,34 @@ pub fn main_rt(args: TokenStream, item: TokenStream) -> TokenStream { /// /// Note that `start_paused` requires the `test-util` feature to be enabled. /// +/// ### Spawn and join the test future +/// +/// ```no_run +/// #[tokio::test(spawn = true)] +/// async fn my_test() { +/// assert!(true); +/// } +/// ``` +/// +/// Equivalent code not using `#[tokio::test]` +/// +/// ```no_run +/// #[test] +/// fn my_test() { +/// let runtime = tokio::runtime::Builder::new_current_thread() +/// .enable_all() +/// .start_paused(true) +/// .build() +/// .unwrap(); +/// let handle = runtime.spawn(async { +/// assert!(true); +/// }); +/// runtime.block_on(async move { +/// handle.await.unwrap(); +/// }); +/// } +/// ``` +/// /// ### Rename package /// /// ```rust