✨ worker
This commit is contained in:
parent
5cb4facc48
commit
e3393f1e09
GPG key ID:
901B2ADDF27C2263
11 changed files with 841 additions and 48 deletions
98
comrade-macro/Cargo.lock
generated
Normal file
98
comrade-macro/Cargo.lock
generated
Normal file
|
|
@ -0,0 +1,98 @@
|
|||
# This file is automatically @generated by Cargo.
|
||||
# It is not intended for manual editing.
|
||||
version = 4
|
||||
|
||||
[[package]]
|
||||
name = "comrade-macro"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"serde_json",
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "itoa"
|
||||
version = "1.0.15"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4a5f13b858c8d314ee3e8f639011f7ccefe71f97f96e50151fb991f267928e2c"
|
||||
|
||||
[[package]]
|
||||
name = "memchr"
|
||||
version = "2.7.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3"
|
||||
|
||||
[[package]]
|
||||
name = "proc-macro2"
|
||||
version = "1.0.94"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a31971752e70b8b2686d7e46ec17fb38dad4051d94024c88df49b667caea9c84"
|
||||
dependencies = [
|
||||
"unicode-ident",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "quote"
|
||||
version = "1.0.39"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c1f1914ce909e1658d9907913b4b91947430c7d9be598b15a1912935b8c04801"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ryu"
|
||||
version = "1.0.20"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "28d3b2b1366ec20994f1fd18c3c594f05c5dd4bc44d8bb0c1c632c8d6829481f"
|
||||
|
||||
[[package]]
|
||||
name = "serde"
|
||||
version = "1.0.218"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e8dfc9d19bdbf6d17e22319da49161d5d0108e4188e8b680aef6299eed22df60"
|
||||
dependencies = [
|
||||
"serde_derive",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_derive"
|
||||
version = "1.0.218"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f09503e191f4e797cb8aac08e9a4a4695c5edf6a2e70e376d961ddd5c969f82b"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_json"
|
||||
version = "1.0.140"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "20068b6e96dc6c9bd23e01df8827e6c7e1f2fddd43c21810382803c136b99373"
|
||||
dependencies = [
|
||||
"itoa",
|
||||
"memchr",
|
||||
"ryu",
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "syn"
|
||||
version = "2.0.99"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e02e925281e18ffd9d640e234264753c43edc62d64b2d4cf898f1bc5e75f3fc2"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"unicode-ident",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "unicode-ident"
|
||||
version = "1.0.18"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5a5f39404a5da50712a4c1eecf25e90dd62b613502b7e925fd4e4d19b5c96512"
|
||||
13
comrade-macro/Cargo.toml
Normal file
13
comrade-macro/Cargo.toml
Normal file
|
|
@ -0,0 +1,13 @@
|
|||
[package]
|
||||
name = "comrade-macro"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
|
||||
[dependencies]
|
||||
proc-macro2 = "1.0.94"
|
||||
quote = "1.0.39"
|
||||
serde_json = "1.0.140"
|
||||
syn = { version = "2.0.99", features = ["full"] }
|
||||
|
||||
[lib]
|
||||
proc-macro = true
|
||||
200
comrade-macro/src/lib.rs
Normal file
200
comrade-macro/src/lib.rs
Normal file
|
|
@ -0,0 +1,200 @@
|
|||
use proc_macro::TokenStream;
|
||||
use quote::{ToTokens, format_ident, quote};
|
||||
use syn::{FnArg, Ident, ItemFn, Pat, ReturnType, Type, parse_macro_input};
|
||||
|
||||
/// This macro turns this function into a worker.
|
||||
///
|
||||
/// This will upgrade the function and generate a few ones (`fn` is a placeholder for the functions name):
|
||||
/// - `fn()` - This function will be exactly the same as the original but it will be computed by a worker.
|
||||
/// - `fn_init(&ServiceManager) -> ServiceManager` - This function registers a worker thread on a `ServiceManager`.
|
||||
/// - `fn_shutdown()` - This function issues a shutdown request.
|
||||
/// - `fn_init_scoped(&ServiceManager) -> (ServiceManager, fn_Scoped)` - This function registers a worker thread on a `ServiceManager` and returns a scoped struct. You can call the underlying function with `.call()` on the struct and it will automatically shutdown any workers if it gets out of scope.
|
||||
///
|
||||
/// # Examples
|
||||
/// ```ignore
|
||||
/// use comrade::worker;
|
||||
///
|
||||
/// // Declare worker
|
||||
/// #[worker]
|
||||
/// pub fn multiply(a: i32, b: i32) -> i32 {
|
||||
/// a * b
|
||||
/// }
|
||||
///
|
||||
/// fn main() {
|
||||
/// let s = ServiceManager::new().mode(comrade::service::ServiceMode::Decay);
|
||||
///
|
||||
/// // Init worker thread
|
||||
/// let s = multiply_init(s);
|
||||
/// let s = s.spawn();
|
||||
///
|
||||
/// // Usage
|
||||
/// let x = multiply(2, 8);
|
||||
/// println!("myfn {x}");
|
||||
///
|
||||
/// // Shutdown worker thread
|
||||
/// multiply_shutdown();
|
||||
///
|
||||
/// s.join().unwrap();
|
||||
/// }
|
||||
/// ```
|
||||
#[proc_macro_attribute]
|
||||
pub fn worker(_attr: TokenStream, item: TokenStream) -> TokenStream {
|
||||
let input: ItemFn = parse_macro_input!(item as ItemFn);
|
||||
let fn_name = &input.sig.ident;
|
||||
|
||||
// Extract parameter names and types separately
|
||||
let params: Vec<(Ident, Type)> = input
|
||||
.sig
|
||||
.inputs
|
||||
.iter()
|
||||
.filter_map(|arg| {
|
||||
if let FnArg::Typed(pat_type) = arg {
|
||||
let name = if let Pat::Ident(pat_ident) = *pat_type.pat.clone() {
|
||||
pat_ident.ident.clone()
|
||||
} else {
|
||||
return None;
|
||||
};
|
||||
let ty = *pat_type.ty.clone();
|
||||
Some((name, ty))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Extract parameter names and types into separate lists for quoting
|
||||
let param_names: Vec<Ident> = params.iter().map(|(name, _)| name.clone()).collect();
|
||||
let param_types: Vec<Type> = params.iter().map(|(_, ty)| ty.clone()).collect();
|
||||
|
||||
for t in ¶m_types {
|
||||
println!("param {}", t.to_token_stream().to_string());
|
||||
}
|
||||
|
||||
// Extract return type
|
||||
let return_type = match &input.sig.output {
|
||||
ReturnType::Type(_, ty) => quote!(#ty),
|
||||
ReturnType::Default => quote!(()),
|
||||
};
|
||||
|
||||
// Extract function body
|
||||
let body = &input.block;
|
||||
|
||||
let wrapper_fn = format_ident!("{}_wrapper", fn_name);
|
||||
let worker_fn = format_ident!("{}_worker", fn_name);
|
||||
let init_fn = format_ident!("{}_init", fn_name);
|
||||
let init_fn_scoped = format_ident!("{}_init_scoped", fn_name);
|
||||
let fn_scope_struct = format_ident!("{}_Scoped", fn_name);
|
||||
let shutdown_fn = format_ident!("{}_shutdown", fn_name);
|
||||
|
||||
let param_unpacking = param_names.iter().enumerate().map(|(i, name)| {
|
||||
if param_names.len() == 1 {
|
||||
return quote! {
|
||||
let #name = i;
|
||||
|
||||
};
|
||||
}
|
||||
|
||||
let param_type = ¶m_types[i];
|
||||
if let Type::Path(_) = param_type {
|
||||
quote! {
|
||||
let #name = i.#i;
|
||||
}
|
||||
} else {
|
||||
quote! {
|
||||
let #name = i;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let output = quote! {
|
||||
pub fn #fn_name(#(#param_names: #param_types),*) -> #return_type {
|
||||
let i: comrade::serde_json::Value = comrade::serde_json::to_value( (#(#param_names),*) ).unwrap();
|
||||
serde_json::from_value(comrade::UNION.get(stringify!(#fn_name)).unwrap().send(i)).unwrap()
|
||||
}
|
||||
|
||||
fn #wrapper_fn(task: JobOrder<comrade::serde_json::Value, comrade::serde_json::Value>) {
|
||||
let i = task.param.clone();
|
||||
|
||||
// Deserialize the parameter into the function's expected types
|
||||
let i: (#(#param_types),*) = comrade::serde_json::from_value(i).unwrap();
|
||||
|
||||
#(#param_unpacking)*
|
||||
|
||||
let res = #body;
|
||||
|
||||
task.done(comrade::serde_json::to_value(&res).unwrap());
|
||||
}
|
||||
|
||||
pub fn #worker_fn(recv: Receiver<JobOrder<comrade::serde_json::Value, comrade::serde_json::Value>>) {
|
||||
loop {
|
||||
let task = recv.recv();
|
||||
|
||||
match task {
|
||||
Ok(task) => {
|
||||
if let comrade::serde_json::Value::Object(obj) = &task.param {
|
||||
if obj.contains_key("task") {
|
||||
log::info!("Shutdown requested for task worker {}", stringify!(#fn_name));
|
||||
task.done(comrade::serde_json::json!({"ok": 1}));
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
#wrapper_fn(task)
|
||||
},
|
||||
Err(e) => {
|
||||
log::error!("Error receiving task: {e:?}");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[doc = "Shutdown the worker"]
|
||||
pub fn #shutdown_fn() {
|
||||
comrade::UNION.get(stringify!(#fn_name)).unwrap().send(comrade::serde_json::json!({"task": "shutdown"}));
|
||||
}
|
||||
|
||||
#[doc = "Initialize a worker thread on `ServiceManager`"]
|
||||
pub fn #init_fn(sm: ServiceManager) -> ServiceManager {
|
||||
let (dispatch, recv): (JobDispatcher<_, _>, Receiver<JobOrder<_, _>>) = JobDispatcher::new();
|
||||
|
||||
let sm = sm.register(stringify!(#worker_fn), move |_| {
|
||||
#worker_fn(recv.clone());
|
||||
});
|
||||
|
||||
comrade::UNION.insert(stringify!(#fn_name), dispatch);
|
||||
|
||||
sm
|
||||
}
|
||||
|
||||
#[allow(non_camel_case_types)]
|
||||
pub struct #fn_scope_struct {}
|
||||
|
||||
impl #fn_scope_struct {
|
||||
pub fn call(&self, #(#param_names: #param_types),*) -> #return_type {
|
||||
#fn_name(#(#param_names),*)
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for #fn_scope_struct {
|
||||
fn drop(&mut self) {
|
||||
log::info!("Scoped task worker got dropped.");
|
||||
#shutdown_fn();
|
||||
}
|
||||
}
|
||||
|
||||
#[doc = "Initialize a worker thread on `ServiceManager` on a scoped lifetime"]
|
||||
pub fn #init_fn_scoped(sm: ServiceManager) -> (ServiceManager, #fn_scope_struct) {
|
||||
let (dispatch, recv): (JobDispatcher<_, _>, Receiver<JobOrder<_, _>>) = JobDispatcher::new();
|
||||
|
||||
let sm = sm.register(stringify!(#worker_fn), move |_| {
|
||||
#worker_fn(recv.clone());
|
||||
});
|
||||
|
||||
comrade::UNION.insert(stringify!(#fn_name), dispatch);
|
||||
|
||||
(sm, #fn_scope_struct {})
|
||||
}
|
||||
};
|
||||
|
||||
output.into()
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue