jmarya/comrade
jmarya/comrade
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

✨ valkey task distribution

Browse source
This commit is contained in:
JMARyA 2025-03-08 21:25:02 +01:00
parent e827faaa3f
commit f381f30e27
Signed by: jmarya
GPG key ID: 901B2ADDF27C2263
7 changed files with 866 additions and 26 deletions
Download patch file Download diff file Expand all files Collapse all files

76
comrade-macro/src/lib.rs
View file

@ -5,10 +5,13 @@ 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.
/// - `fn()` - Behaves exactly the same as the original but it will be computed by a worker.
/// - `fn_init(ServiceManager) -> ServiceManager` - Registers worker threads on a `ServiceManager`.
/// - `fn_shutdown()` - Issue a shutdown request.
/// - `fn_init_scoped(ServiceManager) -> (ServiceManager, fn_Scoped)` - Registers worker threads on a `ServiceManager` and return 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.
/// - `fn_init_union(ServiceManager) -> ServiceManager` - Register worker threads and work in a union using Valkey.
/// - `fn_init_union_scoped(ServiceManager) -> ServiceManager` - Register worker threads and work in a union using Valkey and return a scoped struct.
/// - `fn_register_union()` - Register the worker in a union setup without starting local worker threads.
///
/// # Examples
/// ```ignore
@ -81,9 +84,12 @@ pub fn worker(attr: TokenStream, item: TokenStream) -> TokenStream {
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_union = format_ident!("{}_init_union", fn_name);
let init_fn_scoped = format_ident!("{}_init_scoped", fn_name);
let init_fn_scoped_union = format_ident!("{}_init_union_scoped", fn_name);
let fn_scope_struct = format_ident!("{}_Scoped", fn_name);
let fn_name_async = format_ident!("{}_async", fn_name);
let init_fn_register = format_ident!("{}_register_union", fn_name);
let shutdown_fn = format_ident!("{}_shutdown", fn_name);
let param_unpacking = param_names.iter().enumerate().map(|(i, name)| {
@ -109,13 +115,17 @@ pub fn worker(attr: TokenStream, item: TokenStream) -> TokenStream {
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()
serde_json::from_value(comrade::UNION.get(stringify!(#fn_name)).expect("Function call is not registered in UNION").send(i)).unwrap()
}
#[doc = "Will run the function non blocking returning a `JobResult<_>` for fetching a result later."]
pub fn #fn_name_async(#(#param_names: #param_types),*) -> comrade::job::JobResult<comrade::serde_json::Value> {
let i: comrade::serde_json::Value = comrade::serde_json::to_value( (#(#param_names),*) ).unwrap();
comrade::UNION.get(stringify!(#fn_name)).unwrap().send_async(i)
if let Some(dispatch) = comrade::UNION.get(stringify!(#fn_name)) {
dispatch.send_async(i)
} else {
panic!("Function call {} is not registered in UNION", stringify!(#fn_name))
}
}
fn #wrapper_fn(task: JobOrder<comrade::serde_json::Value, comrade::serde_json::Value>) {
@ -131,7 +141,7 @@ pub fn worker(attr: TokenStream, item: TokenStream) -> TokenStream {
task.done(comrade::serde_json::to_value(&res).unwrap());
}
pub fn #worker_fn(recv: Receiver<JobOrder<comrade::serde_json::Value, comrade::serde_json::Value>>) {
pub fn #worker_fn(recv: comrade::job::TaskReceiverBackend<comrade::serde_json::Value, comrade::serde_json::Value>) {
let mut metrics = (0, 0);
loop {
let task = recv.recv();
@ -167,17 +177,59 @@ pub fn worker(attr: TokenStream, item: TokenStream) -> TokenStream {
let mut dispatchers = Vec::new();
let mut s = sm;
log::info!("Initializing worker {} with {} threads", stringify!(#worker_fn), #worker_count);
for i in 0..#worker_count {
let (dispatch, recv): (JobDispatcher<_, _>, Receiver<JobOrder<_, _>>) = JobDispatcher::new();
s = s.register(
&format!("{}_{i}", stringify!(#worker_fn)),
move |_| {
#worker_fn(recv.clone())
#worker_fn(comrade::job::TaskReceiverBackend::Local(recv.clone()))
}
);
dispatchers.push(dispatch);
dispatchers.push(comrade::job::Dispatcher::Local(dispatch));
}
comrade::UNION.insert(stringify!(#fn_name), comrade::job::JobMultiplexer::from(dispatchers));
s
}
#[doc = "Register worker union on Valkey backend without starting local workers."]
pub fn #init_fn_register() {
let mut dispatchers = Vec::new();
log::info!("Registering worker union {}", stringify!(#worker_fn));
let dispatch = comrade::job::ValkeyJobDispatcher::<_, _>::new_topic(stringify!(#worker_fn), false);
dispatchers.push(comrade::job::Dispatcher::Union(dispatch));
comrade::UNION.insert(stringify!(#fn_name), comrade::job::JobMultiplexer::from(dispatchers));
}
#[doc = "Initialize worker threads on `ServiceManager` with Valkey backend"]
pub fn #init_fn_union(sm: ServiceManager) -> ServiceManager {
let mut dispatchers = Vec::new();
let mut s = sm;
log::info!("Initializing worker union {} with {} threads", stringify!(#worker_fn), #worker_count);
for i in 0..#worker_count {
let dispatch = comrade::job::ValkeyJobDispatcher::<_, _>::new_topic(stringify!(#worker_fn), true);
let recv = dispatch.clone();
s = s.register(
&format!("{}_union_{i}", stringify!(#worker_fn)),
move |_| {
#worker_fn(comrade::job::TaskReceiverBackend::Union(recv.clone()))
}
);
dispatchers.push(comrade::job::Dispatcher::Union(dispatch));
}
comrade::UNION.insert(stringify!(#fn_name), comrade::job::JobMultiplexer::from(dispatchers));
@ -206,6 +258,12 @@ pub fn worker(attr: TokenStream, item: TokenStream) -> TokenStream {
let sm = #init_fn(sm);
(sm, #fn_scope_struct {})
}
#[doc = "Initialize a worker union on `ServiceManager` on a scoped lifetime"]
pub fn #init_fn_scoped_union(sm: ServiceManager) -> (ServiceManager, #fn_scope_struct) {
let sm = #init_fn_union(sm);
(sm, #fn_scope_struct {})
}
};
output.into()