From 38bff0adda393f8121225727d93cb342d8363979 Mon Sep 17 00:00:00 2001 From: Carl Lerche Date: Tue, 21 Jan 2020 10:46:32 -0800 Subject: macros: fix `#[tokio::main]` without rt-core (#2139) The Tokio runtime provides a "shell" runtime when `rt-core` is not available. This shell runtime is enough to support `#[tokio::main`] and `#[tokio::test]. A previous change disabled these two attr macros when `rt-core` was not selected. This patch fixes this by re-enabling the `main` and `test` attr macros without `rt-core` and adds some integration tests to prevent future regressions. --- tokio-macros/src/entry.rs | 365 +++++++++++++++++++++++++++++++++++++++++++++ tokio-macros/src/lib.rs | 370 +--------------------------------------------- 2 files changed, 373 insertions(+), 362 deletions(-) create mode 100644 tokio-macros/src/entry.rs (limited to 'tokio-macros') diff --git a/tokio-macros/src/entry.rs b/tokio-macros/src/entry.rs new file mode 100644 index 00000000..e0b9b024 --- /dev/null +++ b/tokio-macros/src/entry.rs @@ -0,0 +1,365 @@ +#![doc(html_root_url = "https://docs.rs/tokio-macros/0.2.3")] +#![allow(clippy::needless_doctest_main)] +#![warn( + missing_debug_implementations, + missing_docs, + rust_2018_idioms, + unreachable_pub +)] +use proc_macro::TokenStream; +use quote::quote; +use std::num::NonZeroUsize; + +#[derive(Clone, Copy, PartialEq)] +enum Runtime { + Basic, + Threaded, +} + +fn parse_knobs( + input: syn::ItemFn, + args: syn::AttributeArgs, + is_test: bool, + rt_threaded: bool, +) -> Result { + let ret = &input.sig.output; + let name = &input.sig.ident; + let inputs = &input.sig.inputs; + let body = &input.block; + let attrs = &input.attrs; + let vis = input.vis; + + if input.sig.asyncness.is_none() { + let msg = "the async keyword is missing from the function declaration"; + return Err(syn::Error::new_spanned(input.sig.fn_token, msg)); + } + + let mut runtime = None; + let mut core_threads = None; + let mut max_threads = None; + + for arg in args { + match arg { + syn::NestedMeta::Meta(syn::Meta::NameValue(namevalue)) => { + let ident = namevalue.path.get_ident(); + if ident.is_none() { + let msg = "Must have specified ident"; + return Err(syn::Error::new_spanned(namevalue, msg)); + } + match ident.unwrap().to_string().to_lowercase().as_str() { + "core_threads" => { + if rt_threaded { + match &namevalue.lit { + syn::Lit::Int(expr) => { + let num = expr.base10_parse::().unwrap(); + if num.get() > 1 { + runtime = Some(Runtime::Threaded); + } else { + runtime = Some(Runtime::Basic); + } + + if let Some(v) = max_threads { + if v < num { + return Err(syn::Error::new_spanned( + namevalue, + "max_threads cannot be less than core_threads", + )); + } + } + + core_threads = Some(num); + } + _ => { + return Err(syn::Error::new_spanned( + namevalue, + "core_threads argument must be an int", + )) + } + } + } else { + return Err(syn::Error::new_spanned( + namevalue, + "core_threads can only be set with rt-threaded feature flag enabled", + )); + } + } + "max_threads" => match &namevalue.lit { + syn::Lit::Int(expr) => { + let num = expr.base10_parse::().unwrap(); + + if let Some(v) = core_threads { + if num < v { + return Err(syn::Error::new_spanned( + namevalue, + "max_threads cannot be less than core_threads", + )); + } + } + max_threads = Some(num); + } + _ => { + return Err(syn::Error::new_spanned( + namevalue, + "max_threads argument must be an int", + )) + } + }, + name => { + let msg = format!("Unknown attribute pair {} is specified; expected one of: `core_threads`, `max_threads`", name); + return Err(syn::Error::new_spanned(namevalue, msg)); + } + } + } + syn::NestedMeta::Meta(syn::Meta::Path(path)) => { + let ident = path.get_ident(); + if ident.is_none() { + let msg = "Must have specified ident"; + return Err(syn::Error::new_spanned(path, msg)); + } + match ident.unwrap().to_string().to_lowercase().as_str() { + "threaded_scheduler" => { + runtime = Some(runtime.unwrap_or_else(|| Runtime::Threaded)) + } + "basic_scheduler" => runtime = Some(runtime.unwrap_or_else(|| Runtime::Basic)), + name => { + let msg = format!("Unknown attribute {} is specified; expected `basic_scheduler` or `threaded_scheduler`", name); + return Err(syn::Error::new_spanned(path, msg)); + } + } + } + other => { + return Err(syn::Error::new_spanned( + other, + "Unknown attribute inside the macro", + )); + } + } + } + + let mut rt = quote! { tokio::runtime::Builder::new().basic_scheduler() }; + if rt_threaded && (runtime == Some(Runtime::Threaded) || (runtime.is_none() && !is_test)) { + rt = quote! { #rt.threaded_scheduler() }; + } + if let Some(v) = core_threads.map(|v| v.get()) { + rt = quote! { #rt.core_threads(#v) }; + } + if let Some(v) = max_threads.map(|v| v.get()) { + rt = quote! { #rt.max_threads(#v) }; + } + + let header = { + if is_test { + quote! { + #[test] + } + } else { + quote! {} + } + }; + + let result = quote! { + #header + #(#attrs)* + #vis fn #name(#inputs) #ret { + #rt + .enable_all() + .build() + .unwrap() + .block_on(async { #body }) + } + }; + + Ok(result.into()) +} + +#[cfg(not(test))] // Work around for rust-lang/rust#62127 +pub(crate) fn main(args: TokenStream, item: TokenStream, rt_threaded: bool) -> TokenStream { + let input = syn::parse_macro_input!(item as syn::ItemFn); + let args = syn::parse_macro_input!(args as syn::AttributeArgs); + + if input.sig.ident == "main" && !input.sig.inputs.is_empty() { + let msg = "the main function cannot accept arguments"; + return syn::Error::new_spanned(&input.sig.inputs, msg) + .to_compile_error() + .into(); + } + + parse_knobs(input, args, false, rt_threaded).unwrap_or_else(|e| e.to_compile_error().into()) +} + +pub(crate) fn test(args: TokenStream, item: TokenStream, rt_threaded: bool) -> TokenStream { + let input = syn::parse_macro_input!(item as syn::ItemFn); + let args = syn::parse_macro_input!(args as syn::AttributeArgs); + + for attr in &input.attrs { + if attr.path.is_ident("test") { + let msg = "second test attribute is supplied"; + return syn::Error::new_spanned(&attr, msg) + .to_compile_error() + .into(); + } + } + + if !input.sig.inputs.is_empty() { + let msg = "the test function cannot accept arguments"; + return syn::Error::new_spanned(&input.sig.inputs, msg) + .to_compile_error() + .into(); + } + + parse_knobs(input, args, true, rt_threaded).unwrap_or_else(|e| e.to_compile_error().into()) +} + +pub(crate) mod old { + use proc_macro::TokenStream; + use quote::quote; + + enum Runtime { + Basic, + Threaded, + Auto, + } + + #[cfg(not(test))] // Work around for rust-lang/rust#62127 + pub(crate) fn main(args: TokenStream, item: TokenStream) -> TokenStream { + let input = syn::parse_macro_input!(item as syn::ItemFn); + let args = syn::parse_macro_input!(args as syn::AttributeArgs); + + let ret = &input.sig.output; + let name = &input.sig.ident; + let inputs = &input.sig.inputs; + let body = &input.block; + let attrs = &input.attrs; + let vis = input.vis; + + if input.sig.asyncness.is_none() { + let msg = "the async keyword is missing from the function declaration"; + return syn::Error::new_spanned(input.sig.fn_token, msg) + .to_compile_error() + .into(); + } else if name == "main" && !inputs.is_empty() { + let msg = "the main function cannot accept arguments"; + return syn::Error::new_spanned(&input.sig.inputs, msg) + .to_compile_error() + .into(); + } + + let mut runtime = Runtime::Auto; + + for arg in args { + if let syn::NestedMeta::Meta(syn::Meta::Path(path)) = arg { + let ident = path.get_ident(); + if ident.is_none() { + let msg = "Must have specified ident"; + return syn::Error::new_spanned(path, msg).to_compile_error().into(); + } + match ident.unwrap().to_string().to_lowercase().as_str() { + "threaded_scheduler" => runtime = Runtime::Threaded, + "basic_scheduler" => runtime = Runtime::Basic, + name => { + let msg = format!("Unknown attribute {} is specified; expected `basic_scheduler` or `threaded_scheduler`", name); + return syn::Error::new_spanned(path, msg).to_compile_error().into(); + } + } + } + } + + let result = match runtime { + Runtime::Threaded | Runtime::Auto => quote! { + #(#attrs)* + #vis fn #name(#inputs) #ret { + tokio::runtime::Runtime::new().unwrap().block_on(async { #body }) + } + }, + Runtime::Basic => quote! { + #(#attrs)* + #vis fn #name(#inputs) #ret { + tokio::runtime::Builder::new() + .basic_scheduler() + .enable_all() + .build() + .unwrap() + .block_on(async { #body }) + } + }, + }; + + result.into() + } + + pub(crate) fn test(args: TokenStream, item: TokenStream) -> TokenStream { + let input = syn::parse_macro_input!(item as syn::ItemFn); + let args = syn::parse_macro_input!(args as syn::AttributeArgs); + + let ret = &input.sig.output; + let name = &input.sig.ident; + let body = &input.block; + let attrs = &input.attrs; + let vis = input.vis; + + for attr in attrs { + if attr.path.is_ident("test") { + let msg = "second test attribute is supplied"; + return syn::Error::new_spanned(&attr, msg) + .to_compile_error() + .into(); + } + } + + if input.sig.asyncness.is_none() { + let msg = "the async keyword is missing from the function declaration"; + return syn::Error::new_spanned(&input.sig.fn_token, msg) + .to_compile_error() + .into(); + } else if !input.sig.inputs.is_empty() { + let msg = "the test function cannot accept arguments"; + return syn::Error::new_spanned(&input.sig.inputs, msg) + .to_compile_error() + .into(); + } + + let mut runtime = Runtime::Auto; + + for arg in args { + if let syn::NestedMeta::Meta(syn::Meta::Path(path)) = arg { + let ident = path.get_ident(); + if ident.is_none() { + let msg = "Must have specified ident"; + return syn::Error::new_spanned(path, msg).to_compile_error().into(); + } + match ident.unwrap().to_string().to_lowercase().as_str() { + "threaded_scheduler" => runtime = Runtime::Threaded, + "basic_scheduler" => runtime = Runtime::Basic, + name => { + let msg = format!("Unknown attribute {} is specified; expected `basic_scheduler` or `threaded_scheduler`", name); + return syn::Error::new_spanned(path, msg).to_compile_error().into(); + } + } + } + } + + let result = match runtime { + Runtime::Threaded => quote! { + #[test] + #(#attrs)* + #vis fn #name() #ret { + tokio::runtime::Runtime::new().unwrap().block_on(async { #body }) + } + }, + Runtime::Basic | Runtime::Auto => quote! { + #[test] + #(#attrs)* + #vis fn #name() #ret { + tokio::runtime::Builder::new() + .basic_scheduler() + .enable_all() + .build() + .unwrap() + .block_on(async { #body }) + } + }, + }; + + result.into() + } +} diff --git a/tokio-macros/src/lib.rs b/tokio-macros/src/lib.rs index 6eb57ea0..09b8b093 100644 --- a/tokio-macros/src/lib.rs +++ b/tokio-macros/src/lib.rs @@ -14,173 +14,11 @@ //! Macros for use with Tokio +mod entry; + extern crate proc_macro; use proc_macro::TokenStream; -use quote::quote; -use std::num::NonZeroUsize; - -#[derive(Clone, Copy, PartialEq)] -enum Runtime { - Basic, - Threaded, -} - -fn parse_knobs( - input: syn::ItemFn, - args: syn::AttributeArgs, - is_test: bool, - rt_threaded: bool, -) -> Result { - let ret = &input.sig.output; - let name = &input.sig.ident; - let inputs = &input.sig.inputs; - let body = &input.block; - let attrs = &input.attrs; - let vis = input.vis; - - if input.sig.asyncness.is_none() { - let msg = "the async keyword is missing from the function declaration"; - return Err(syn::Error::new_spanned(input.sig.fn_token, msg)); - } - - let mut runtime = None; - let mut core_threads = None; - let mut max_threads = None; - - for arg in args { - match arg { - syn::NestedMeta::Meta(syn::Meta::NameValue(namevalue)) => { - let ident = namevalue.path.get_ident(); - if ident.is_none() { - let msg = "Must have specified ident"; - return Err(syn::Error::new_spanned(namevalue, msg)); - } - match ident.unwrap().to_string().to_lowercase().as_str() { - "core_threads" => { - if rt_threaded { - match &namevalue.lit { - syn::Lit::Int(expr) => { - let num = expr.base10_parse::().unwrap(); - if num.get() > 1 { - runtime = Some(Runtime::Threaded); - } else { - runtime = Some(Runtime::Basic); - } - - if let Some(v) = max_threads { - if v < num { - return Err(syn::Error::new_spanned( - namevalue, - "max_threads cannot be less than core_threads", - )); - } - } - - core_threads = Some(num); - } - _ => { - return Err(syn::Error::new_spanned( - namevalue, - "core_threads argument must be an int", - )) - } - } - } else { - return Err(syn::Error::new_spanned( - namevalue, - "core_threads can only be set with rt-threaded feature flag enabled", - )); - } - } - "max_threads" => match &namevalue.lit { - syn::Lit::Int(expr) => { - let num = expr.base10_parse::().unwrap(); - - if let Some(v) = core_threads { - if num < v { - return Err(syn::Error::new_spanned( - namevalue, - "max_threads cannot be less than core_threads", - )); - } - } - max_threads = Some(num); - } - _ => { - return Err(syn::Error::new_spanned( - namevalue, - "max_threads argument must be an int", - )) - } - }, - name => { - let msg = format!("Unknown attribute pair {} is specified; expected one of: `core_threads`, `max_threads`", name); - return Err(syn::Error::new_spanned(namevalue, msg)); - } - } - } - syn::NestedMeta::Meta(syn::Meta::Path(path)) => { - let ident = path.get_ident(); - if ident.is_none() { - let msg = "Must have specified ident"; - return Err(syn::Error::new_spanned(path, msg)); - } - match ident.unwrap().to_string().to_lowercase().as_str() { - "threaded_scheduler" => { - runtime = Some(runtime.unwrap_or_else(|| Runtime::Threaded)) - } - "basic_scheduler" => runtime = Some(runtime.unwrap_or_else(|| Runtime::Basic)), - name => { - let msg = format!("Unknown attribute {} is specified; expected `basic_scheduler` or `threaded_scheduler`", name); - return Err(syn::Error::new_spanned(path, msg)); - } - } - } - other => { - return Err(syn::Error::new_spanned( - other, - "Unknown attribute inside the macro", - )); - } - } - } - - let mut rt = quote! { tokio::runtime::Builder::new().basic_scheduler() }; - if rt_threaded && (runtime == Some(Runtime::Threaded) || (runtime.is_none() && !is_test)) { - rt = quote! { #rt.threaded_scheduler() }; - } - if let Some(v) = core_threads.map(|v| v.get()) { - rt = quote! { #rt.core_threads(#v) }; - } - if let Some(v) = max_threads.map(|v| v.get()) { - rt = quote! { #rt.max_threads(#v) }; - } - - let header = { - if is_test { - quote! { - #[test] - } - } else { - quote! {} - } - }; - - let result = quote! { - #header - #(#attrs)* - #vis fn #name(#inputs) #ret { - #rt - .enable_all() - .build() - .unwrap() - .block_on(async { #body }) - } - }; - - Ok(result.into()) -} /// Marks async function to be executed by selected runtime. /// @@ -215,7 +53,7 @@ fn parse_knobs( #[proc_macro_attribute] #[cfg(not(test))] // Work around for rust-lang/rust#62127 pub fn main_threaded(args: TokenStream, item: TokenStream) -> TokenStream { - main_impl(args, item, true) + entry::main(args, item, true) } /// Marks async function to be executed by selected runtime. @@ -251,7 +89,7 @@ pub fn main_threaded(args: TokenStream, item: TokenStream) -> TokenStream { #[proc_macro_attribute] #[cfg(not(test))] // Work around for rust-lang/rust#62127 pub fn main(args: TokenStream, item: TokenStream) -> TokenStream { - old::main(args, item) + entry::old::main(args, item) } /// Marks async function to be executed by selected runtime. @@ -277,22 +115,7 @@ pub fn main(args: TokenStream, item: TokenStream) -> TokenStream { #[proc_macro_attribute] #[cfg(not(test))] // Work around for rust-lang/rust#62127 pub fn main_basic(args: TokenStream, item: TokenStream) -> TokenStream { - main_impl(args, item, false) -} - -#[cfg(not(test))] // Work around for rust-lang/rust#62127 -fn main_impl(args: TokenStream, item: TokenStream, rt_threaded: bool) -> TokenStream { - let input = syn::parse_macro_input!(item as syn::ItemFn); - let args = syn::parse_macro_input!(args as syn::AttributeArgs); - - if input.sig.ident == "main" && !input.sig.inputs.is_empty() { - let msg = "the main function cannot accept arguments"; - return syn::Error::new_spanned(&input.sig.inputs, msg) - .to_compile_error() - .into(); - } - - parse_knobs(input, args, false, rt_threaded).unwrap_or_else(|e| e.to_compile_error().into()) + entry::main(args, item, false) } /// Marks async function to be executed by runtime, suitable to test enviornment @@ -323,7 +146,7 @@ fn main_impl(args: TokenStream, item: TokenStream, rt_threaded: bool) -> TokenSt /// ``` #[proc_macro_attribute] pub fn test_threaded(args: TokenStream, item: TokenStream) -> TokenStream { - test_impl(args, item, true) + entry::test(args, item, true) } /// Marks async function to be executed by runtime, suitable to test enviornment @@ -354,7 +177,7 @@ pub fn test_threaded(args: TokenStream, item: TokenStream) -> TokenStream { /// ``` #[proc_macro_attribute] pub fn test(args: TokenStream, item: TokenStream) -> TokenStream { - old::test(args, item) + entry::old::test(args, item) } /// Marks async function to be executed by runtime, suitable to test enviornment @@ -373,182 +196,5 @@ pub fn test(args: TokenStream, item: TokenStream) -> TokenStream { /// ``` #[proc_macro_attribute] pub fn test_basic(args: TokenStream, item: TokenStream) -> TokenStream { - test_impl(args, item, false) -} - -fn test_impl(args: TokenStream, item: TokenStream, rt_threaded: bool) -> TokenStream { - let input = syn::parse_macro_input!(item as syn::ItemFn); - let args = syn::parse_macro_input!(args as syn::AttributeArgs); - - for attr in &input.attrs { - if attr.path.is_ident("test") { - let msg = "second test attribute is supplied"; - return syn::Error::new_spanned(&attr, msg) - .to_compile_error() - .into(); - } - } - - if !input.sig.inputs.is_empty() { - let msg = "the test function cannot accept arguments"; - return syn::Error::new_spanned(&input.sig.inputs, msg) - .to_compile_error() - .into(); - } - - parse_knobs(input, args, true, rt_threaded).unwrap_or_else(|e| e.to_compile_error().into()) -} - -mod old { - use proc_macro::TokenStream; - use quote::quote; - - enum Runtime { - Basic, - Threaded, - Auto, - } - - #[cfg(not(test))] // Work around for rust-lang/rust#62127 - pub(crate) fn main(args: TokenStream, item: TokenStream) -> TokenStream { - let input = syn::parse_macro_input!(item as syn::ItemFn); - let args = syn::parse_macro_input!(args as syn::AttributeArgs); - - let ret = &input.sig.output; - let name = &input.sig.ident; - let inputs = &input.sig.inputs; - let body = &input.block; - let attrs = &input.attrs; - let vis = input.vis; - - if input.sig.asyncness.is_none() { - let msg = "the async keyword is missing from the function declaration"; - return syn::Error::new_spanned(input.sig.fn_token, msg) - .to_compile_error() - .into(); - } else if name == "main" && !inputs.is_empty() { - let msg = "the main function cannot accept arguments"; - return syn::Error::new_spanned(&input.sig.inputs, msg) - .to_compile_error() - .into(); - } - - let mut runtime = Runtime::Auto; - - for arg in args { - if let syn::NestedMeta::Meta(syn::Meta::Path(path)) = arg { - let ident = path.get_ident(); - if ident.is_none() { - let msg = "Must have specified ident"; - return syn::Error::new_spanned(path, msg).to_compile_error().into(); - } - match ident.unwrap().to_string().to_lowercase().as_str() { - "threaded_scheduler" => runtime = Runtime::Threaded, - "basic_scheduler" => runtime = Runtime::Basic, - name => { - let msg = format!("Unknown attribute {} is specified; expected `basic_scheduler` or `threaded_scheduler`", name); - return syn::Error::new_spanned(path, msg).to_compile_error().into(); - } - } - } - } - - let result = match runtime { - Runtime::Threaded | Runtime::Auto => quote! { - #(#attrs)* - #vis fn #name(#inputs) #ret { - tokio::runtime::Runtime::new().unwrap().block_on(async { #body }) - } - }, - Runtime::Basic => quote! { - #(#attrs)* - #vis fn #name(#inputs) #ret { - tokio::runtime::Builder::new() - .basic_scheduler() - .enable_all() - .build() - .unwrap() - .block_on(async { #body }) - } - }, - }; - - result.into() - } - - pub(crate) fn test(args: TokenStream, item: TokenStream) -> TokenStream { - let input = syn::parse_macro_input!(item as syn::ItemFn); - let args = syn::parse_macro_input!(args as syn::AttributeArgs); - - let ret = &input.sig.output; - let name = &input.sig.ident; - let body = &input.block; - let attrs = &input.attrs; - let vis = input.vis; - - for attr in attrs { - if attr.path.is_ident("test") { - let msg = "second test attribute is supplied"; - return syn::Error::new_spanned(&attr, msg) - .to_compile_error() - .into(); - } - } - - if input.sig.asyncness.is_none() { - let msg = "the async keyword is missing from the function declaration"; - return syn::Error::new_spanned(&input.sig.fn_token, msg) - .to_compile_error() - .into(); - } else if !input.sig.inputs.is_empty() { - let msg = "the test function cannot accept arguments"; - return syn::Error::new_spanned(&input.sig.inputs, msg) - .to_compile_error() - .into(); - } - - let mut runtime = Runtime::Auto; - - for arg in args { - if let syn::NestedMeta::Meta(syn::Meta::Path(path)) = arg { - let ident = path.get_ident(); - if ident.is_none() { - let msg = "Must have specified ident"; - return syn::Error::new_spanned(path, msg).to_compile_error().into(); - } - match ident.unwrap().to_string().to_lowercase().as_str() { - "threaded_scheduler" => runtime = Runtime::Threaded, - "basic_scheduler" => runtime = Runtime::Basic, - name => { - let msg = format!("Unknown attribute {} is specified; expected `basic_scheduler` or `threaded_scheduler`", name); - return syn::Error::new_spanned(path, msg).to_compile_error().into(); - } - } - } - } - - let result = match runtime { - Runtime::Threaded => quote! { - #[test] - #(#attrs)* - #vis fn #name() #ret { - tokio::runtime::Runtime::new().unwrap().block_on(async { #body }) - } - }, - Runtime::Basic | Runtime::Auto => quote! { - #[test] - #(#attrs)* - #vis fn #name() #ret { - tokio::runtime::Builder::new() - .basic_scheduler() - .enable_all() - .build() - .unwrap() - .block_on(async { #body }) - } - }, - }; - - result.into() - } + entry::test(args, item, false) } -- cgit v1.2.3