Skip to content

Commit

Permalink
macros: custom package name for #[tokio::main] and #[tokio::test]
Browse files Browse the repository at this point in the history 
This also enables `#[crate::test(package = "crate")]` in unit tests.

Sees rust-lang/cargo#5653.
Fixes #2312.
  • Loading branch information
@kezhuw
kezhuw committed Apr 11, 2022
1 parent 252b0fa commit 6610eab
Show file tree
Hide file tree
Showing 3 changed files with 93 additions and 38 deletions.
40 changes: 37 additions & 3 deletions tokio-macros/src/entry.rs
View file
@@ -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;

Expand Down Expand Up @@ -29,13 +29,15 @@ struct FinalConfig {
flavor: RuntimeFlavor,
worker_threads: Option<usize>,
start_paused: Option<bool>,
package_name: Option<Ident>,
}

/// Config used in case of the attribute not being able to build a valid config
const DEFAULT_ERROR_CONFIG: FinalConfig = FinalConfig {
flavor: RuntimeFlavor::CurrentThread,
worker_threads: None,
start_paused: None,
package_name: None,
};

struct Configuration {
Expand All @@ -45,6 +47,7 @@ struct Configuration {
worker_threads: Option<(usize, Span)>,
start_paused: Option<(bool, Span)>,
is_test: bool,
package_name: Option<Ident>,
}

impl Configuration {
Expand All @@ -59,6 +62,7 @@ impl Configuration {
worker_threads: None,
start_paused: None,
is_test,
package_name: None,
}
}

Expand Down Expand Up @@ -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"
Expand Down Expand Up @@ -151,6 +164,7 @@ impl Configuration {
};

Ok(FinalConfig {
package_name: self.package_name.clone(),
flavor,
worker_threads,
start_paused,
Expand Down Expand Up @@ -185,6 +199,16 @@ fn parse_string(int: syn::Lit, span: Span, field: &str) -> Result<String, syn::E
}
}

fn parse_ident(lit: syn::Lit, span: Span, field: &str) -> Result<Ident, syn::Error> {
match lit {
syn::Lit::Str(s) => s.parse::<syn::Ident>(),
_ => 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<bool, syn::Error> {
match bool {
syn::Lit::Bool(b) => Ok(b.value),
Expand Down Expand Up @@ -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`",
Expand Down Expand Up @@ -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 {
Expand Down
71 changes: 49 additions & 22 deletions tokio-macros/src/lib.rs
View file
Expand Up @@ -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 {
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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)
Expand All @@ -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)
Expand Down
20 changes: 7 additions & 13 deletions tokio/src/sync/tests/notify.rs
View file
Expand Up @@ -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;
}

0 comments on commit 6610eab

Please sign in to comment.