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

✨ multi threaded workers

Browse source
This commit is contained in:
JMARyA 2025-03-08 18:01:16 +01:00
parent 0c190df5d7
commit e827faaa3f
Signed by: jmarya
GPG key ID: 901B2ADDF27C2263
7 changed files with 208 additions and 45 deletions
Download patch file Download diff file Expand all files Collapse all files

117
Cargo.lock generated
View file

@ -103,6 +103,12 @@ version = "2.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5c8214115b7bf84099f1309324e63141d4c5d7cc26862f97a0a857dbefe165bd" checksum = "5c8214115b7bf84099f1309324e63141d4c5d7cc26862f97a0a857dbefe165bd"
[[package]]
name = "byteorder"
version = "1.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b"
[[package]] [[package]]
name = "cfg-if" name = "cfg-if"
version = "1.0.0" version = "1.0.0"
@ -125,6 +131,7 @@ dependencies = [
"env_logger", "env_logger",
"log", "log",
"once_cell", "once_cell",
"rand",
"serde_json", "serde_json",
"tokio", "tokio",
] ]
@ -232,6 +239,18 @@ dependencies = [
"log", "log",
] ]
[[package]]
name = "getrandom"
version = "0.3.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "43a49c392881ce6d5c3b8cb70f98717b7c07aabbdff06687b9030dbfbe2725f8"
dependencies = [
"cfg-if",
"libc",
"wasi",
"windows-targets",
]
[[package]] [[package]]
name = "gimli" name = "gimli"
version = "0.31.1" version = "0.31.1"
@ -333,6 +352,15 @@ version = "0.2.16"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3b3cff922bd51709b605d9ead9aa71031d81447142d828eb4a6eba76fe619f9b" checksum = "3b3cff922bd51709b605d9ead9aa71031d81447142d828eb4a6eba76fe619f9b"
[[package]]
name = "ppv-lite86"
version = "0.2.20"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "77957b295656769bb8ad2b6a6b09d897d94f05c41b069aede1fcdaa675eaea04"
dependencies = [
"zerocopy 0.7.35",
]
[[package]] [[package]]
name = "proc-macro2" name = "proc-macro2"
version = "1.0.94" version = "1.0.94"
@ -351,6 +379,36 @@ dependencies = [
"proc-macro2", "proc-macro2",
] ]
[[package]]
name = "rand"
version = "0.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3779b94aeb87e8bd4e834cee3650289ee9e0d5677f976ecdb6d219e5f4f6cd94"
dependencies = [
"rand_chacha",
"rand_core",
"zerocopy 0.8.23",
]
[[package]]
name = "rand_chacha"
version = "0.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb"
dependencies = [
"ppv-lite86",
"rand_core",
]
[[package]]
name = "rand_core"
version = "0.9.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "99d9a13982dcf210057a8a78572b2217b667c3beacbf3a0d8b454f6f82837d38"
dependencies = [
"getrandom",
]
[[package]] [[package]]
name = "redox_syscall" name = "redox_syscall"
version = "0.5.10" version = "0.5.10"
@ -478,6 +536,15 @@ version = "0.2.2"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821"
[[package]]
name = "wasi"
version = "0.13.3+wasi-0.2.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "26816d2e1a4a36a2940b96c5296ce403917633dff8f3440e9b236ed6f6bacad2"
dependencies = [
"wit-bindgen-rt",
]
[[package]] [[package]]
name = "windows-sys" name = "windows-sys"
version = "0.59.0" version = "0.59.0"
@ -550,3 +617,53 @@ name = "windows_x86_64_msvc"
version = "0.52.6" version = "0.52.6"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec"
[[package]]
name = "wit-bindgen-rt"
version = "0.33.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3268f3d866458b787f390cf61f4bbb563b922d091359f9608842999eaee3943c"
dependencies = [
"bitflags",
]
[[package]]
name = "zerocopy"
version = "0.7.35"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1b9b4fd18abc82b8136838da5d50bae7bdea537c574d8dc1a34ed098d6c166f0"
dependencies = [
"byteorder",
"zerocopy-derive 0.7.35",
]
[[package]]
name = "zerocopy"
version = "0.8.23"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fd97444d05a4328b90e75e503a34bad781f14e28a823ad3557f0750df1ebcbc6"
dependencies = [
"zerocopy-derive 0.8.23",
]
[[package]]
name = "zerocopy-derive"
version = "0.7.35"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fa4f8080344d4671fb4e831a13ad1e68092748387dfc4f55e356242fae12ce3e"
dependencies = [
"proc-macro2",
"quote",
"syn",
]
[[package]]
name = "zerocopy-derive"
version = "0.8.23"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6352c01d0edd5db859a63e2605f4ea3183ddbd15e2c4a9e7d32184df75e4f154"
dependencies = [
"proc-macro2",
"quote",
"syn",
]

1
Cargo.toml
View file

@ -12,3 +12,4 @@ once_cell = "1.20.3"
tokio = "1.43.0" tokio = "1.43.0"
comrade-macro = { path = "./comrade-macro" } comrade-macro = { path = "./comrade-macro" }
serde_json = "1.0.140" serde_json = "1.0.140"
rand = "0.9.0"

4
README.md
View file

@ -71,12 +71,14 @@ You can annotate a function with `#[worker]` which gives them superpowers. These
```rust ```rust
use comrade::{worker}; use comrade::{worker};
// Single local worker
#[worker] #[worker]
pub fn myfn(i: i32) -> i32 { pub fn myfn(i: i32) -> i32 {
i * 2 i * 2
} }
#[worker] // 4 local worker threads
#[worker(4)]
pub fn multiply(a: i32, b: i32) -> i32 { pub fn multiply(a: i32, b: i32) -> i32 {
a * b a * b
} }

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

@ -1,5 +1,5 @@
use proc_macro::TokenStream; use proc_macro::TokenStream;
use quote::{ToTokens, format_ident, quote}; use quote::{format_ident, quote};
use syn::{FnArg, Ident, ItemFn, Pat, ReturnType, Type, parse_macro_input}; use syn::{FnArg, Ident, ItemFn, Pat, ReturnType, Type, parse_macro_input};
/// This macro turns this function into a worker. /// This macro turns this function into a worker.
@ -38,10 +38,13 @@ use syn::{FnArg, Ident, ItemFn, Pat, ReturnType, Type, parse_macro_input};
/// } /// }
/// ``` /// ```
#[proc_macro_attribute] #[proc_macro_attribute]
pub fn worker(_attr: TokenStream, item: TokenStream) -> TokenStream { pub fn worker(attr: TokenStream, item: TokenStream) -> TokenStream {
let input: ItemFn = parse_macro_input!(item as ItemFn); let input: ItemFn = parse_macro_input!(item as ItemFn);
let fn_name = &input.sig.ident; let fn_name = &input.sig.ident;
// Parse optional attribute argument (e.g., #[worker(4)])
let worker_count = attr.to_string().parse::<usize>().unwrap_or(1);
// Extract parameter names and types separately // Extract parameter names and types separately
let params: Vec<(Ident, Type)> = input let params: Vec<(Ident, Type)> = input
.sig .sig
@ -66,10 +69,6 @@ pub fn worker(_attr: TokenStream, item: TokenStream) -> TokenStream {
let param_names: Vec<Ident> = params.iter().map(|(name, _)| name.clone()).collect(); let param_names: Vec<Ident> = params.iter().map(|(name, _)| name.clone()).collect();
let param_types: Vec<Type> = params.iter().map(|(_, ty)| ty.clone()).collect(); let param_types: Vec<Type> = params.iter().map(|(_, ty)| ty.clone()).collect();
for t in &param_types {
println!("param {}", t.to_token_stream().to_string());
}
// Extract return type // Extract return type
let return_type = match &input.sig.output { let return_type = match &input.sig.output {
ReturnType::Type(_, ty) => quote!(#ty), ReturnType::Type(_, ty) => quote!(#ty),
@ -133,6 +132,7 @@ pub fn worker(_attr: TokenStream, item: TokenStream) -> TokenStream {
} }
pub fn #worker_fn(recv: Receiver<JobOrder<comrade::serde_json::Value, comrade::serde_json::Value>>) { pub fn #worker_fn(recv: Receiver<JobOrder<comrade::serde_json::Value, comrade::serde_json::Value>>) {
let mut metrics = (0, 0);
loop { loop {
let task = recv.recv(); let task = recv.recv();
@ -140,15 +140,17 @@ pub fn worker(_attr: TokenStream, item: TokenStream) -> TokenStream {
Ok(task) => { Ok(task) => {
if let comrade::serde_json::Value::Object(obj) = &task.param { if let comrade::serde_json::Value::Object(obj) = &task.param {
if obj.contains_key("task") { if obj.contains_key("task") {
log::info!("Shutdown requested for task worker {}", stringify!(#fn_name)); log::info!("Shutdown requested for task worker {}. Processed {} tasks since start with {} errors.", stringify!(#fn_name), metrics.0, metrics.1);
task.done(comrade::serde_json::json!({"ok": 1})); task.done(comrade::serde_json::json!({"ok": 1}));
break; break;
} }
} }
#wrapper_fn(task) #wrapper_fn(task);
metrics.0 += 1;
}, },
Err(e) => { Err(e) => {
metrics.1 += 1;
log::error!("Error receiving task: {e:?}"); log::error!("Error receiving task: {e:?}");
} }
} }
@ -157,20 +159,30 @@ pub fn worker(_attr: TokenStream, item: TokenStream) -> TokenStream {
#[doc = "Shutdown the worker"] #[doc = "Shutdown the worker"]
pub fn #shutdown_fn() { pub fn #shutdown_fn() {
comrade::UNION.get(stringify!(#fn_name)).unwrap().send(comrade::serde_json::json!({"task": "shutdown"})); comrade::UNION.get(stringify!(#fn_name)).unwrap().send_all(comrade::serde_json::json!({"task": "shutdown"}));
} }
#[doc = "Initialize a worker thread on `ServiceManager`"] #[doc = "Initialize a worker thread on `ServiceManager`"]
pub fn #init_fn(sm: ServiceManager) -> ServiceManager { pub fn #init_fn(sm: ServiceManager) -> ServiceManager {
let (dispatch, recv): (JobDispatcher<_, _>, Receiver<JobOrder<_, _>>) = JobDispatcher::new(); let mut dispatchers = Vec::new();
let mut s = sm;
let sm = sm.register(stringify!(#worker_fn), move |_| { for i in 0..#worker_count {
#worker_fn(recv.clone()); let (dispatch, recv): (JobDispatcher<_, _>, Receiver<JobOrder<_, _>>) = JobDispatcher::new();
});
comrade::UNION.insert(stringify!(#fn_name), dispatch); s = s.register(
&format!("{}_{i}", stringify!(#worker_fn)),
move |_| {
#worker_fn(recv.clone())
}
);
sm dispatchers.push(dispatch);
}
comrade::UNION.insert(stringify!(#fn_name), comrade::job::JobMultiplexer::from(dispatchers));
s
} }
#[allow(non_camel_case_types)] #[allow(non_camel_case_types)]
@ -191,14 +203,7 @@ pub fn worker(_attr: TokenStream, item: TokenStream) -> TokenStream {
#[doc = "Initialize a worker thread on `ServiceManager` on a scoped lifetime"] #[doc = "Initialize a worker thread on `ServiceManager` on a scoped lifetime"]
pub fn #init_fn_scoped(sm: ServiceManager) -> (ServiceManager, #fn_scope_struct) { pub fn #init_fn_scoped(sm: ServiceManager) -> (ServiceManager, #fn_scope_struct) {
let (dispatch, recv): (JobDispatcher<_, _>, Receiver<JobOrder<_, _>>) = JobDispatcher::new(); let sm = #init_fn(sm);
let sm = sm.register(stringify!(#worker_fn), move |_| {
#worker_fn(recv.clone());
});
comrade::UNION.insert(stringify!(#fn_name), dispatch);
(sm, #fn_scope_struct {}) (sm, #fn_scope_struct {})
} }
}; };

2
examples/work.rs
View file

@ -17,7 +17,7 @@ pub fn myfn(i: i32) -> i32 {
i * 2 i * 2
} }
#[worker] #[worker(4)]
pub fn multiply(a: i32, b: i32) -> i32 { pub fn multiply(a: i32, b: i32) -> i32 {
a * b a * b
} }

76
src/job.rs
View file

@ -1,22 +1,23 @@
use crossbeam::channel::{Receiver, Sender}; use crossbeam::channel::{Receiver, Sender};
use rand::Rng;
use std::sync::mpsc; use std::sync::mpsc;
#[derive(Clone)] #[derive(Clone)]
/// A generic job dispatcher struct that allows sending jobs of type `T` and receiving results of type `V` using message passing. /// A generic job dispatcher struct that allows sending jobs of type `I` and receiving results of type `O` using message passing.
pub struct JobDispatcher<T: Send + 'static, V: Send + 'static> { pub struct JobDispatcher<I: Send + 'static, O: Send + 'static> {
sender: Sender<JobOrder<T, V>>, sender: Sender<JobOrder<I, O>>,
} }
pub struct JobResult<V>(std::sync::mpsc::Receiver<V>); pub struct JobResult<O>(std::sync::mpsc::Receiver<O>);
impl<V> JobResult<V> { impl<O> JobResult<O> {
/// Wait for the Result of a Job. /// Wait for the Result of a Job.
pub fn wait(self) -> V { pub fn wait(self) -> O {
self.0.recv().unwrap() self.0.recv().unwrap()
} }
} }
impl<T: Send + 'static, V: Send + 'static> JobDispatcher<T, V> { impl<I: Send + 'static, O: Send + 'static> JobDispatcher<I, O> {
/// Creates a new instance of `JobDispatcher` and returns a tuple that contains it and a receiver end for `JobOrder`s. /// Creates a new instance of `JobDispatcher` and returns a tuple that contains it and a receiver end for `JobOrder`s.
/// # Example: /// # Example:
/// ``` /// ```
@ -37,7 +38,7 @@ impl<T: Send + 'static, V: Send + 'static> JobDispatcher<T, V> {
/// assert_eq!(result, 4); /// assert_eq!(result, 4);
/// ``` /// ```
#[must_use] #[must_use]
pub fn new() -> (Self, Receiver<JobOrder<T, V>>) { pub fn new() -> (Self, Receiver<JobOrder<I, O>>) {
let (sender, receiver) = crossbeam::channel::bounded(12); let (sender, receiver) = crossbeam::channel::bounded(12);
(Self { sender: sender }, receiver) (Self { sender: sender }, receiver)
@ -48,7 +49,7 @@ impl<T: Send + 'static, V: Send + 'static> JobDispatcher<T, V> {
/// # Panics /// # Panics
/// This function panics when the `JobOrder` struct gets out of scope without returning a finished result. /// This function panics when the `JobOrder` struct gets out of scope without returning a finished result.
/// Additionally if the internal `Mutex` is poisoned, this function will panic as well. /// Additionally if the internal `Mutex` is poisoned, this function will panic as well.
pub fn send(&self, param: T) -> V { pub fn send(&self, param: I) -> O {
let (tx, rx) = mpsc::channel(); let (tx, rx) = mpsc::channel();
let job_order = JobOrder::new(param, move |ret| { let job_order = JobOrder::new(param, move |ret| {
tx.send(ret).unwrap(); tx.send(ret).unwrap();
@ -57,7 +58,7 @@ impl<T: Send + 'static, V: Send + 'static> JobDispatcher<T, V> {
rx.recv().unwrap() rx.recv().unwrap()
} }
pub fn send_async(&self, param: T) -> JobResult<V> { pub fn send_async(&self, param: I) -> JobResult<O> {
let (tx, rx) = mpsc::channel(); let (tx, rx) = mpsc::channel();
let job_order = JobOrder::new(param, move |ret| { let job_order = JobOrder::new(param, move |ret| {
tx.send(ret).unwrap(); tx.send(ret).unwrap();
@ -68,7 +69,7 @@ impl<T: Send + 'static, V: Send + 'static> JobDispatcher<T, V> {
/// Sends a job of type `T` to the job dispatcher and waits for its result of type `V`. /// Sends a job of type `T` to the job dispatcher and waits for its result of type `V`.
/// Returns `Some(V)` when the job returns an result, `None` if somehow nothing was returned or the internal `Mutex` is poisoned. /// Returns `Some(V)` when the job returns an result, `None` if somehow nothing was returned or the internal `Mutex` is poisoned.
pub fn try_send(&self, param: T) -> Option<V> { pub fn try_send(&self, param: I) -> Option<O> {
let (tx, rx) = mpsc::channel(); let (tx, rx) = mpsc::channel();
let job_order = JobOrder::new(param, move |ret| { let job_order = JobOrder::new(param, move |ret| {
tx.send(ret).unwrap(); tx.send(ret).unwrap();
@ -78,17 +79,17 @@ impl<T: Send + 'static, V: Send + 'static> JobDispatcher<T, V> {
} }
} }
/// A struct that represents a job order that encapsulates a job of type `T` and its result of type `V`, along with a callback function that will send the result back to the job origin. /// A struct that represents a job order that encapsulates a job of type `I` and its result of type `O`, along with a callback function that will send the result back to the job origin.
pub struct JobOrder<T, V> { pub struct JobOrder<I, O> {
/// The job parameter of type `T`. /// The job parameter of type `T`.
pub param: T, pub param: I,
callback: Box<dyn FnOnce(V) + Send>, callback: Box<dyn FnOnce(O) + Send>,
} }
impl<T, V> JobOrder<T, V> { impl<I, O> JobOrder<I, O> {
/// Creates a new `JobOrder` instance with the specified job parameter `param` of type `T` and a callback function that takes the job result of type `V` as an argument. /// Creates a new `JobOrder` instance with the specified job parameter `param` of type `I` and a callback function that takes the job result of type `O` as an argument.
#[must_use] #[must_use]
fn new(param: T, callback: impl FnOnce(V) + Send + 'static) -> Self { fn new(param: I, callback: impl FnOnce(O) + Send + 'static) -> Self {
Self { Self {
param, param,
callback: Box::new(callback), callback: Box::new(callback),
@ -96,7 +97,44 @@ impl<T, V> JobOrder<T, V> {
} }
/// Send the result of the `JobOrder` back to it's origin /// Send the result of the `JobOrder` back to it's origin
pub fn done(self, val: V) { pub fn done(self, val: O) {
(self.callback)(val); (self.callback)(val);
} }
} }
pub struct JobMultiplexer<I: Send + 'static, O: Send + 'static> {
dispatchers: Vec<JobDispatcher<I, O>>,
}
fn get_random_item<T>(list: &[T]) -> Option<&T> {
if list.is_empty() {
return None;
}
let mut rng = rand::rng();
let index = rng.random_range(0..list.len());
list.get(index)
}
impl<I: Send + 'static, O: Send + 'static> JobMultiplexer<I, O> {
pub fn from(dispatchers: Vec<JobDispatcher<I, O>>) -> Self {
Self { dispatchers }
}
pub fn send(&self, param: I) -> O {
let d = get_random_item(&self.dispatchers).unwrap();
d.send(param)
}
pub fn send_async(&self, param: I) -> JobResult<O> {
let d = get_random_item(&self.dispatchers).unwrap();
d.send_async(param)
}
}
impl<I: Clone + Send + 'static, O: Send + 'static> JobMultiplexer<I, O> {
pub fn send_all(&self, param: I) {
for d in &self.dispatchers {
let _ = d.send(param.clone());
}
}
}

2
src/lib.rs
View file

@ -12,7 +12,7 @@ pub use serde_json;
// TODO : refactor dispatcher backends (memory, valkey) // TODO : refactor dispatcher backends (memory, valkey)
pub static UNION: Lazy< pub static UNION: Lazy<
DashMap<&'static str, job::JobDispatcher<serde_json::Value, serde_json::Value>>, DashMap<&'static str, job::JobMultiplexer<serde_json::Value, serde_json::Value>>,
> = Lazy::new(DashMap::new); > = Lazy::new(DashMap::new);
/// Rally Function /// Rally Function