Skip to content

Commit

Permalink
Feature: featrue flag monoio enables monoio as runtime
Browse files Browse the repository at this point in the history
`async-entry = { features = ["monoio"] }` for using `monoio` as the
async runtime. It produces a runtime such as:

```
let mut rt = monoio::RuntimeBuilder::<monoio::FusionDriver>::new()
    .enable_all()
    .build()
    .expect("Failed building the Runtime");
rt.block_on(body);
```

To use the default tokio runtime, add async-entry with
`features = ["tokio"]` or `features = []`.

When using `monoio` runtime: `flavor`, `worker_threads` and `start_paused` are ignored.
drmingdrmer committed Feb 19, 2024
1 parent 9fbfa83 commit 7da5496
Showing 15 changed files with 1,490 additions and 109 deletions.
3 changes: 3 additions & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
@@ -1,4 +1,7 @@
[workspace]

resolver = "1"

members = [
"async-entry",
"test-async-entry",
14 changes: 14 additions & 0 deletions async-entry/Cargo.toml
Original file line number Diff line number Diff line change
@@ -12,9 +12,22 @@ extended Tokio's proc macros.
"""
categories = ["asynchronous"]

[features]

default = []

# Use tokio as runtime:
# https://tokio.rs/
tokio = []

# Use monoio as runtime:
# https://github.com/bytedance/monoio
monoio = []

[lib]
proc-macro = true


[dependencies]
syn = { version = "1.0.56", features = ["full"] }
quote = "1"
@@ -24,4 +37,5 @@ proc-macro2 = "1.0.7"
[dev-dependencies]
anyhow = "1.0.55"
tokio = { version = "1.7.1", features = ["macros", "rt","rt-multi-thread", "sync", "test-util"] }
monoio = { version = "0.1.0", features = [] }
tracing = "0.1.29"
132 changes: 102 additions & 30 deletions async-entry/src/lib.rs
Original file line number Diff line number Diff line change
@@ -5,9 +5,21 @@ use proc_macro2::Span;
use quote::quote;
use quote::quote_spanned;
use quote::ToTokens;
use syn::__private::TokenStream2;
use syn::parse::Parser;
use syn::ItemFn;

fn get_runtime_name() -> &'static str {
if cfg!(feature = "tokio") {
return "tokio";
}
if cfg!(feature = "monoio") {
return "monoio";
} else {
return "tokio";
}
}

#[derive(Debug, Clone, Copy, PartialEq)]
enum RuntimeFlavor {
CurrentThread,
@@ -143,9 +155,17 @@ impl Configuration {

fn macro_name(&self) -> &'static str {
if self.is_test {
"tokio::test"
match get_runtime_name() {
"tokio" => "tokio::test",
"monoio" => "monoio::test",
_ => unreachable!(),
}
} else {
"tokio::main"
match get_runtime_name() {
"tokio" => "tokio::main",
"monoio" => "monoio::main",
_ => unreachable!(),
}
}
}

@@ -333,9 +353,17 @@ fn build_config(args: AttributeArgs, rt_multi_thread: bool) -> Result<FinalConfi

type AttributeArgs = syn::punctuated::Punctuated<syn::NestedMeta, syn::Token![,]>;

/// Marks async function to be executed by runtime, suitable to test environment
/// Marks async function to be executed by async runtime, suitable to test environment
///
/// It supports:
/// - [tokio](https://tokio.rs/)
/// - [monoio](https://github.com/bytedance/monoio)
///
/// ## Usage
/// By default it uses `tokio` runtime. Switch runtime with feature flags:
/// - `tokio`: tokio runtime;
/// - `monoio`: monoio runtime;
///
/// ## Usage for tokio runtime
///
/// ### Multi-thread runtime
///
@@ -432,6 +460,29 @@ type AttributeArgs = syn::punctuated::Punctuated<syn::NestedMeta, syn::Token![,]
/// // }
/// ```
///
/// ## Usage for monoio runtime
///
/// **When using `monoio` runtime with feature flag `monoio` enabled:
/// `flavor`, `worker_threads` and `start_paused` are ignored**.
///
/// It is the same as using `tokio` runtime, except the runtime is `monoio`:
///
/// ```no_run
/// #[async_entry::test()]
/// async fn my_test() {
/// assert!(true);
/// }
/// // Will produce:
/// //
/// // fn my_test() {
/// // // ...
/// //
/// // let body = async { assert!(true); };
/// // let rt = monoio::RuntimeBuilder::<_>::new()...
/// // rt.block_on(body);
/// // }
/// ```
///
/// ### NOTE:
///
/// If you rename the async_entry crate in your dependencies this macro will not work.
@@ -479,31 +530,7 @@ fn build_test_fn(mut item_fn: ItemFn, config: FinalConfig) -> Result<TokenStream

let test_attr = quote! { #[::core::prelude::v1::test] };

let mut rt_builder = quote! { tokio::runtime::Builder };

rt_builder = match config.flavor {
RuntimeFlavor::CurrentThread => quote_spanned! {last_stmt_start_span=>
#rt_builder::new_current_thread()
},
RuntimeFlavor::Threaded => quote_spanned! {last_stmt_start_span=>
#rt_builder::new_multi_thread()
},
};

if let Some(v) = config.worker_threads {
rt_builder = quote! { #rt_builder.worker_threads(#v) };
}

if let Some(v) = config.start_paused {
rt_builder = quote! { #rt_builder.start_paused(#v) };
}

let rt = quote! {
#rt_builder
.enable_all()
.build()
.expect("Failed building the Runtime")
};
let rt = build_runtime(last_stmt_start_span, &config)?;

let init = if let Some(init) = config.init {
let init_str = format!("let _g = {};", init.0);
@@ -563,7 +590,8 @@ fn build_test_fn(mut item_fn: ItemFn, config: FinalConfig) -> Result<TokenStream

#body_tracing_span

let rt = #rt;
#[allow(unused_mut)]
let mut rt = #rt;

#[allow(clippy::expect_used)]
#tail_return rt.block_on(body) #tail_semicolon
@@ -583,6 +611,50 @@ fn build_test_fn(mut item_fn: ItemFn, config: FinalConfig) -> Result<TokenStream
Ok(x)
}

/// Build a statement that builds a async runtime,
/// e.g. `let rt = Builder::new_multi_thread().build().expect("");`
fn build_runtime(span: Span, config: &FinalConfig) -> Result<TokenStream2, syn::Error> {
let rt_builder = {
match get_runtime_name() {
"tokio" => {
let mut rt_builder = quote! { tokio::runtime::Builder };

rt_builder = match config.flavor {
RuntimeFlavor::CurrentThread => quote_spanned! {span=>
#rt_builder::new_current_thread()
},
RuntimeFlavor::Threaded => quote_spanned! {span=>
#rt_builder::new_multi_thread()
},
};

if let Some(v) = config.worker_threads {
rt_builder = quote! { #rt_builder.worker_threads(#v) };
}

if let Some(v) = config.start_paused {
rt_builder = quote! { #rt_builder.start_paused(#v) };
}
rt_builder
}
"monoio" => {
let rt_builder = quote! { monoio::RuntimeBuilder::<monoio::FusionDriver>::new() };
rt_builder
}
_ => unreachable!(),
}
};

let rt: TokenStream2 = quote! {
#rt_builder
.enable_all()
.build()
.expect("Failed building the Runtime")
};

Ok(rt)
}

/// Parse TokenStream of some fn
fn parse_item_fn(item: TokenStream) -> Result<ItemFn, syn::Error> {
let input = syn::parse::<ItemFn>(item.clone())?;
Loading

0 comments on commit 7da5496

Please sign in to comment.