diff --git a/tokio-macros/src/entry.rs b/tokio-macros/src/entry.rs index 5cb4a49b430..059fac415bf 100644 --- a/tokio-macros/src/entry.rs +++ b/tokio-macros/src/entry.rs @@ -1,5 +1,5 @@ use proc_macro::TokenStream; -use proc_macro2::Span; +use proc_macro2::{Ident, Span}; use quote::{quote, quote_spanned, ToTokens}; use syn::parse::Parser; @@ -29,6 +29,7 @@ struct FinalConfig { flavor: RuntimeFlavor, worker_threads: Option, start_paused: Option, + package_name: Option, } /// Config used in case of the attribute not being able to build a valid config @@ -36,6 +37,7 @@ const DEFAULT_ERROR_CONFIG: FinalConfig = FinalConfig { flavor: RuntimeFlavor::CurrentThread, worker_threads: None, start_paused: None, + package_name: None, }; struct Configuration { @@ -45,6 +47,7 @@ struct Configuration { worker_threads: Option<(usize, Span)>, start_paused: Option<(bool, Span)>, is_test: bool, + package_name: Option, } impl Configuration { @@ -59,6 +62,7 @@ impl Configuration { worker_threads: None, start_paused: None, is_test, + package_name: None, } } @@ -104,6 +108,15 @@ impl Configuration { Ok(()) } + fn set_package_name(&mut self, name: syn::Lit, span: Span) -> Result<(), syn::Error> { + if self.package_name.is_some() { + return Err(syn::Error::new(span, "`package` set multiple times.")); + } + let name_ident = parse_ident(name, span, "package")?; + self.package_name = Some(name_ident); + Ok(()) + } + fn macro_name(&self) -> &'static str { if self.is_test { "tokio::test" @@ -151,6 +164,7 @@ impl Configuration { }; Ok(FinalConfig { + package_name: self.package_name.clone(), flavor, worker_threads, start_paused, @@ -185,6 +199,16 @@ fn parse_string(int: syn::Lit, span: Span, field: &str) -> Result Result { + match lit { + syn::Lit::Str(s) => s.parse::(), + _ => Err(syn::Error::new( + span, + format!("Failed to parse value of `{}` as ident.", field), + )), + } +} + fn parse_bool(bool: syn::Lit, span: Span, field: &str) -> Result { match bool { syn::Lit::Bool(b) => Ok(b.value), @@ -243,6 +267,12 @@ fn build_config( let msg = "Attribute `core_threads` is renamed to `worker_threads`"; return Err(syn::Error::new_spanned(namevalue, msg)); } + "package" => { + config.set_package_name( + namevalue.lit.clone(), + syn::spanned::Spanned::span(&namevalue.lit), + )?; + } name => { let msg = format!( "Unknown attribute {} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`", @@ -313,12 +343,16 @@ fn parse_knobs(mut input: syn::ItemFn, is_test: bool, config: FinalConfig) -> To (start, end) }; + let package_name = config + .package_name + .unwrap_or_else(|| Ident::new("tokio", Span::call_site())); + let mut rt = match config.flavor { RuntimeFlavor::CurrentThread => quote_spanned! {last_stmt_start_span=> - tokio::runtime::Builder::new_current_thread() + #package_name::runtime::Builder::new_current_thread() }, RuntimeFlavor::Threaded => quote_spanned! {last_stmt_start_span=> - tokio::runtime::Builder::new_multi_thread() + #package_name::runtime::Builder::new_multi_thread() }, }; if let Some(v) = config.worker_threads { diff --git a/tokio-macros/src/lib.rs b/tokio-macros/src/lib.rs index 38638a1df8a..ab491876b4f 100644 --- a/tokio-macros/src/lib.rs +++ b/tokio-macros/src/lib.rs @@ -168,12 +168,28 @@ use proc_macro::TokenStream; /// /// Note that `start_paused` requires the `test-util` feature to be enabled. /// -/// ### NOTE: +/// ### Rename package /// -/// If you rename the Tokio crate in your dependencies this macro will not work. -/// If you must rename the current version of Tokio because you're also using an -/// older version of Tokio, you _must_ make the current version of Tokio -/// available as `tokio` in the module where this macro is expanded. +/// ```rust +/// #[tokio1::main(package = "tokio1")] +/// async fn main() { +/// println!("Hello world"); +/// } +/// ``` +/// +/// Equivalent code not using `#[tokio::main]` +/// +/// ```rust +/// fn main() { +/// tokio1::runtime::Builder::new_multi_thread() +/// .enable_all() +/// .build() +/// .unwrap() +/// .block_on(async { +/// println!("Hello world"); +/// }) +/// } +/// ``` #[proc_macro_attribute] #[cfg(not(test))] // Work around for rust-lang/rust#62127 pub fn main(args: TokenStream, item: TokenStream) -> TokenStream { @@ -213,12 +229,28 @@ pub fn main(args: TokenStream, item: TokenStream) -> TokenStream { /// } /// ``` /// -/// ### NOTE: +/// ### Rename package /// -/// If you rename the Tokio crate in your dependencies this macro will not work. -/// If you must rename the current version of Tokio because you're also using an -/// older version of Tokio, you _must_ make the current version of Tokio -/// available as `tokio` in the module where this macro is expanded. +/// ```rust +/// #[tokio1::main(package = "tokio1")] +/// async fn main() { +/// println!("Hello world"); +/// } +/// ``` +/// +/// Equivalent code not using `#[tokio::main]` +/// +/// ```rust +/// fn main() { +/// tokio1::runtime::Builder::new_multi_thread() +/// .enable_all() +/// .build() +/// .unwrap() +/// .block_on(async { +/// println!("Hello world"); +/// }) +/// } +/// ``` #[proc_macro_attribute] #[cfg(not(test))] // Work around for rust-lang/rust#62127 pub fn main_rt(args: TokenStream, item: TokenStream) -> TokenStream { @@ -260,12 +292,14 @@ pub fn main_rt(args: TokenStream, item: TokenStream) -> TokenStream { /// /// Note that `start_paused` requires the `test-util` feature to be enabled. /// -/// ### NOTE: +/// ### Rename package /// -/// If you rename the Tokio crate in your dependencies this macro will not work. -/// If you must rename the current version of Tokio because you're also using an -/// older version of Tokio, you _must_ make the current version of Tokio -/// available as `tokio` in the module where this macro is expanded. +/// ```rust +/// #[tokio1::test(package = "tokio1")] +/// async fn my_test() { +/// println!("Hello world"); +/// } +/// ``` #[proc_macro_attribute] pub fn test(args: TokenStream, item: TokenStream) -> TokenStream { entry::test(args, item, true) @@ -281,13 +315,6 @@ pub fn test(args: TokenStream, item: TokenStream) -> TokenStream { /// assert!(true); /// } /// ``` -/// -/// ### NOTE: -/// -/// If you rename the Tokio crate in your dependencies this macro will not work. -/// If you must rename the current version of Tokio because you're also using an -/// older version of Tokio, you _must_ make the current version of Tokio -/// available as `tokio` in the module where this macro is expanded. #[proc_macro_attribute] pub fn test_rt(args: TokenStream, item: TokenStream) -> TokenStream { entry::test(args, item, false) diff --git a/tokio/src/sync/tests/notify.rs b/tokio/src/sync/tests/notify.rs index 20153b7a5a8..6ffbb427465 100644 --- a/tokio/src/sync/tests/notify.rs +++ b/tokio/src/sync/tests/notify.rs @@ -62,20 +62,14 @@ fn notify_simple() { assert!(fut2.poll().is_ready()); } -#[test] +#[crate::test(package = "crate", flavor = "current_thread")] #[cfg(not(target_arch = "wasm32"))] -fn watch_test() { - let rt = crate::runtime::Builder::new_current_thread() - .build() - .unwrap(); - - rt.block_on(async { - let (tx, mut rx) = crate::sync::watch::channel(()); +async fn watch_test() { + let (tx, mut rx) = crate::sync::watch::channel(()); - crate::spawn(async move { - let _ = tx.send(()); - }); - - let _ = rx.changed().await; + crate::spawn(async move { + let _ = tx.send(()); }); + + let _ = rx.changed().await; }