diff options
author | Douman <douman@gmx.se> | 2019-06-27 20:40:22 +0300 |
---|---|---|
committer | Carl Lerche <me@carllerche.com> | 2019-06-27 10:40:21 -0700 |
commit | 0af05e7408cef92b834718ed1bd3578fc0fbd40e (patch) | |
tree | b852f4ca8a8a22a6927060c0e498cc078ff0abeb /tokio-macros/src/lib.rs | |
parent | 6b9e7bdace71bf64e12f7ec461a351f7eb188f60 (diff) |
macros: allow configuring runtime used by main macro (#1185)
Diffstat (limited to 'tokio-macros/src/lib.rs')
-rw-r--r-- | tokio-macros/src/lib.rs | 71 |
1 files changed, 62 insertions, 9 deletions
diff --git a/tokio-macros/src/lib.rs b/tokio-macros/src/lib.rs index 3cf5d76a..1562b877 100644 --- a/tokio-macros/src/lib.rs +++ b/tokio-macros/src/lib.rs @@ -10,24 +10,50 @@ use proc_macro::TokenStream; use quote::{quote, quote_spanned}; use syn::spanned::Spanned; -/// Define the program entry point +/// Marks async function to be executed by selected runtime. /// -/// # Examples +/// ## Options: +/// +/// - `single_thread` - Uses `current_thread`. +/// - `multi_thread` - Uses multi-threaded runtime. Used by default. +/// +/// ## Usage /// +/// ### Select runtime +/// +/// ```rust +///#![feature(async_await)] +/// +/// #[tokio::main(single_thread)] +/// async fn main() { +/// println!("Hello world"); +/// } /// ``` +/// ### Using default +/// +/// ```rust +///#![feature(async_await)] +/// /// #[tokio::main] /// async fn main() { -/// println!("Hello from Tokio!"); +/// println!("Hello world"); /// } /// ``` #[proc_macro_attribute] #[cfg(not(test))] // Work around for rust-lang/rust#62127 -pub fn main(_attr: TokenStream, item: TokenStream) -> TokenStream { +pub fn main(args: TokenStream, item: TokenStream) -> TokenStream { + enum RuntimeType { + Single, + Multi, + } + let input = syn::parse_macro_input!(item as syn::ItemFn); + let args = syn::parse_macro_input!(args as syn::AttributeArgs); let ret = &input.decl.output; let name = &input.ident; let body = &input.block; + let attrs = &input.attrs; if input.asyncness.is_none() { let tokens = quote_spanned! { input.span() => @@ -37,21 +63,48 @@ pub fn main(_attr: TokenStream, item: TokenStream) -> TokenStream { return TokenStream::from(tokens); } - let result = quote! { - fn #name() #ret { - let mut rt = tokio::runtime::Runtime::new().unwrap(); - rt.block_on(async { #body }) + let mut runtime = RuntimeType::Multi; + + for arg in args { + match arg { + syn::NestedMeta::Meta(syn::Meta::Word(ident)) => match ident.to_string().to_lowercase().as_str() { + "multi_thread" => runtime = RuntimeType::Multi, + "single_thread" => runtime = RuntimeType::Single, + name => panic!("Unknown attribute {} is specified", name), + }, + _ => () + } + } + + let result = match runtime { + RuntimeType::Multi => quote! { + #(#attrs)* + fn #name() #ret { + let mut rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async { #body }) + } + }, + RuntimeType::Single => quote! { + #(#attrs)* + fn #name() #ret { + let mut rt = tokio::runtime::current_thread::Runtime::new().unwrap(); + rt.block_on(async { #body }) + } } }; result.into() } -/// Define a Tokio aware unit test +/// Marks async function to be executed by runtime, suitable to test enviornment +/// +/// Uses `current_thread` runtime. /// /// # Examples /// /// ```ignore +/// #![feature(async_await)] +/// /// #[tokio::test] /// async fn my_test() { /// assert!(true); |