summaryrefslogtreecommitdiffstats
path: root/tokio-macros/src/lib.rs
diff options
context:
space:
mode:
Diffstat (limited to 'tokio-macros/src/lib.rs')
-rw-r--r--tokio-macros/src/lib.rs370
1 files changed, 8 insertions, 362 deletions
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<TokenStream, syn::Error> {
- 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::<NonZeroUsize>().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::<NonZeroUsize>().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)
}
generated by cgit v1.2.3 (git 2.25.1) at 2024-05-12 10:17:44 +0000