iterated fn + defer + retry

This commit is contained in:
JMARyA 2025-03-10 11:27:20 +01:00
parent c010cc13e8
commit 18c663fcdb
Signed by: jmarya
GPG key ID: 901B2ADDF27C2263
11 changed files with 678 additions and 65 deletions

1
Cargo.lock generated
View file

@ -167,6 +167,7 @@ version = "0.1.0"
dependencies = [
"proc-macro2",
"quote",
"rand",
"serde_json",
"syn",
]

120
README.md
View file

@ -11,37 +11,47 @@
## Core Concepts
### Parallel Execution
### Higher Level Functions
`comrade` provides a simple interface for running tasks in parallel, perfect for independent tasks that can be processed concurrently.
`comrade` provides various convenient functions.
```rust
// process every item in parallel
let results: Vec<i32> = parallel(items, |item: &i32| {
// ...
});
```
### Rally Execution
The `rally` function allows you to run multiple tasks in parallel and return the result of the **first task to finish**. This is useful when you want to prioritize the first available result from several tasks (example: download from multiple HTTP mirrors).
```rust
// rally (return fastest computed result out of items)
// example: run multiple downloads and return the first finished one
let res: (i32, i32) = rally(items, |item: &i32| {
// ...
});
```
### Background Tasks
// Run background tasks without blocking the main thread
background(|| {
// Background task logic
println!("This is a background task!");
});
Easily run tasks in the background without blocking the main thread. This is useful for code that needs to be run without waiting for a result.
fn some_fn() {
println!("Hello World!");
```rust
fn handle() {
background(|| {
// Background task logic
println!("This is a background task!");
defer!(|| {
// this will run at the end of the scope
println!("Bye World!");
});
println!("doing something");
}
// Retry a `Fn() -> Option<X>` until it returns `Some(_)`.
let value: &str = retry(|| {
if rand::rng().random_range(0..10) > 5 {
Some("hello")
} else {
None
}
})
```
### Service Management
@ -113,6 +123,21 @@ fn main() {
}
```
You could effeciently run batch work:
```rust
fn batch_work() {
let mut work = Vec::new();
for i in 0..10 {
work.push((i.to_string(), multiply_async(i, i)));
}
for (label, res) in LabelPendingTaskIterator(work) {
println!("Finished task {label} -> {res}");
}
}
```
These tasks can now be distributed with Valkey.
Make sure you have a Valkey server running and the `$VALKEY_URL` environment variable is set for your application:
@ -155,3 +180,66 @@ fn main() {
println!("x is {x}");
}
```
### Stateful Functions
If you have a workload which can be iteratively computed by modifying state it can be modeled as a `IteratedFunction`.
These functions can be paused, stopped, saved to disk and revived later.
First define a iterative function:
```rust
#[derive(Debug, Default, Clone)]
pub struct AccFnContext {
pub iter: u64,
pub acc: u64,
}
pub fn multiply_iterated(
mut ctx: FunctionContext<AccFnContext, (u64, u64), u64>,
) -> FunctionContext<AccFnContext, (u64, u64), u64> {
// init
let (a, b) = ctx.initial;
// end condition (return)
if b == ctx.context.iter {
ctx.done(ctx.context.acc);
return ctx;
}
// computation
println!("doing calc {}", ctx.context.acc);
std::thread::sleep(Duration::from_millis(50));
let val = ctx.context.acc + a;
// saving state
ctx.state(|x| {
x.iter += 1;
x.acc = val;
});
return ctx;
}
```
Then you can use it like:
```rust
fn main() {
// run normally
let f = IteratedFunction::new_threaded(multiply_iterated, (5, 5));
println!("Result is {}", f.output());
// function starts running
let f = IteratedFunction::new_threaded(multiply_iterated, (5, 50));
// pause the function
f.pause();
// stop the function and get state
let state = f.stop();
// revive and start running again from state
let f = IteratedFunction::new_threaded_from_state(multiply_iterated, state);
// get output
println!("Result is {}", f.output());
}
```

172
comrade-macro/Cargo.lock generated
View file

@ -2,28 +2,68 @@
# It is not intended for manual editing.
version = 4
[[package]]
name = "bitflags"
version = "2.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5c8214115b7bf84099f1309324e63141d4c5d7cc26862f97a0a857dbefe165bd"
[[package]]
name = "cfg-if"
version = "1.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
[[package]]
name = "comrade-macro"
version = "0.1.0"
dependencies = [
"proc-macro2",
"quote",
"rand",
"serde_json",
"syn",
]
[[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]]
name = "itoa"
version = "1.0.15"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4a5f13b858c8d314ee3e8f639011f7ccefe71f97f96e50151fb991f267928e2c"
[[package]]
name = "libc"
version = "0.2.170"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "875b3680cb2f8f71bdcf9a30f38d48282f5d3c95cbf9b3fa57269bb5d5c06828"
[[package]]
name = "memchr"
version = "2.7.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3"
[[package]]
name = "ppv-lite86"
version = "0.2.21"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9"
dependencies = [
"zerocopy",
]
[[package]]
name = "proc-macro2"
version = "1.0.94"
@ -42,6 +82,36 @@ dependencies = [
"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",
]
[[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]]
name = "ryu"
version = "1.0.20"
@ -96,3 +166,105 @@ name = "unicode-ident"
version = "1.0.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5a5f39404a5da50712a4c1eecf25e90dd62b613502b7e925fd4e4d19b5c96512"
[[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]]
name = "windows-targets"
version = "0.52.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973"
dependencies = [
"windows_aarch64_gnullvm",
"windows_aarch64_msvc",
"windows_i686_gnu",
"windows_i686_gnullvm",
"windows_i686_msvc",
"windows_x86_64_gnu",
"windows_x86_64_gnullvm",
"windows_x86_64_msvc",
]
[[package]]
name = "windows_aarch64_gnullvm"
version = "0.52.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3"
[[package]]
name = "windows_aarch64_msvc"
version = "0.52.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469"
[[package]]
name = "windows_i686_gnu"
version = "0.52.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b"
[[package]]
name = "windows_i686_gnullvm"
version = "0.52.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66"
[[package]]
name = "windows_i686_msvc"
version = "0.52.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66"
[[package]]
name = "windows_x86_64_gnu"
version = "0.52.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78"
[[package]]
name = "windows_x86_64_gnullvm"
version = "0.52.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d"
[[package]]
name = "windows_x86_64_msvc"
version = "0.52.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
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.8.23"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fd97444d05a4328b90e75e503a34bad781f14e28a823ad3557f0750df1ebcbc6"
dependencies = [
"zerocopy-derive",
]
[[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",
]

View file

@ -6,6 +6,7 @@ edition = "2024"
[dependencies]
proc-macro2 = "1.0.94"
quote = "1.0.39"
rand = "0.9.0"
serde_json = "1.0.140"
syn = { version = "2.0.99", features = ["full"] }

View file

@ -1,6 +1,7 @@
use proc_macro::TokenStream;
use quote::{format_ident, quote};
use syn::{FnArg, Ident, ItemFn, Pat, ReturnType, Type, parse_macro_input};
use rand::Rng;
use syn::{ExprClosure, FnArg, Ident, ItemFn, Pat, ReturnType, Type, parse_macro_input};
/// This macro turns this function into a worker.
///
@ -271,3 +272,35 @@ pub fn worker(attr: TokenStream, item: TokenStream) -> TokenStream {
output.into()
}
/// A procedural macro that defers execution of a closure until the surrounding scope ends similiarly to Go's defer statement.
///
/// # Example
/// ```
/// use comrade::defer;
///
/// fn main() {
/// defer!(|| {
/// println!("This will be executed at the end of the scope.");
/// });
/// println!("This runs first.");
/// }
/// ```
///
/// The `defer!` macro ensures that the provided closure is wrapped in a `Defer` instance, which will execute the closure when the variable goes out of scope.
#[proc_macro]
pub fn defer(input: TokenStream) -> TokenStream {
// Parse the input as a closure expression (|| { ... })
let closure = parse_macro_input!(input as ExprClosure);
// Generate a random number for a unique variable name
let random_number = rand::rng().random_range(1000..9999);
let rand_ident = format_ident!("{}", format!("defer_{}", random_number));
// Expand into a let statement that wraps the closure in Defer::new
let expanded = quote! {
let #rand_ident = comrade::Defer::new(#closure);
};
TokenStream::from(expanded)
}

12
examples/defer.rs Normal file
View file

@ -0,0 +1,12 @@
use comrade::defer;
fn main() {
println!("Hello World!");
defer!(|| {
// this will run at the end of the scope
println!("Bye World!");
});
println!("doing something");
}

67
examples/iterated.rs Normal file
View file

@ -0,0 +1,67 @@
use std::time::Duration;
use comrade::iterated::{FunctionContext, IteratedFunction};
#[derive(Debug, Default, Clone)]
pub struct AccFnContext {
pub iter: u64,
pub acc: u64,
}
pub fn multiply_iterated(
mut ctx: FunctionContext<AccFnContext, (u64, u64), u64>,
) -> FunctionContext<AccFnContext, (u64, u64), u64> {
// init
let (a, b) = ctx.initial;
// end condition (return)
if b == ctx.context.iter {
ctx.done(ctx.context.acc);
return ctx;
}
// computation
println!("doing calc {}", ctx.context.acc);
std::thread::sleep(Duration::from_millis(50));
let val = ctx.context.acc + a;
// saving state
ctx.state(|x| {
x.iter += 1;
x.acc = val;
});
return ctx;
}
fn main() {
env_logger::init();
let f: IteratedFunction<AccFnContext, (u64, u64), u64> =
IteratedFunction::new(multiply_iterated);
let x = f.run_to_end((5, 4));
println!("computed x -> {x}");
// async
let f = IteratedFunction::new_threaded(multiply_iterated, (5, 5));
println!("This is running");
println!("result is {}", f.output());
let f = IteratedFunction::new_threaded(multiply_iterated, (5, 50));
// pause the function
f.pause();
// stop the function and get state
let state = f.stop();
println!("This was running");
println!("state is {state:?}");
println!("reviving function");
// continue with previous computed state
let f = IteratedFunction::new_threaded_from_state(multiply_iterated, state);
// get output
println!("result is {}", f.output());
}

View file

@ -1,7 +1,7 @@
use std::time::Duration;
use comrade::{
job::{JobDispatcher, JobOrder},
job::{JobDispatcher, JobOrder, LabelPendingTaskIterator},
service::ServiceManager,
worker,
};
@ -22,6 +22,18 @@ pub fn multiply(a: i32, b: i32) -> i32 {
a * b
}
pub fn batch_work() {
let mut work = Vec::new();
for i in 0..10 {
work.push((i.to_string(), multiply_async(i, i)));
}
for (label, res) in LabelPendingTaskIterator(work) {
println!("Finished task {label} -> {res}");
}
}
fn do_work(multiply: multiply_Scoped, myfn: myfn_Scoped) {
for i in 0..10 {
let x = multiply.call(i, i);
@ -41,17 +53,19 @@ fn main() {
let s = s.spawn();
do_work(multiply, myfn_fn);
s.join().unwrap();
let s = ServiceManager::new().mode(comrade::service::ServiceMode::Decay);
let s = myfn_init(s);
let s = multiply_init(s);
let s = take_time_init(s);
let s = s.spawn();
let x = myfn(55);
println!("myfn {x}");
batch_work();
// decoupled
let e = take_time_async(1500);
println!("This will run right after!");

16
src/defer.rs Normal file
View file

@ -0,0 +1,16 @@
pub struct Defer {
f: Box<dyn Fn()>,
}
impl Defer {
pub fn new<F: Fn() + 'static>(f: F) -> Self {
Self { f: Box::new(f) }
}
}
impl Drop for Defer {
fn drop(&mut self) {
log::debug!("Calling defer function");
self.f.as_ref()();
}
}

240
src/iterated.rs Normal file
View file

@ -0,0 +1,240 @@
// TODO : docs
// TODO : measure avg iteration time
use std::{sync::Arc, time::Duration};
use crate::{
job::{JobDispatch, JobDispatcher},
retry,
};
#[derive(Debug, Clone)]
pub struct FunctionContext<C, I, O> {
pub context: C,
pub initial: I,
pub output: Option<O>,
done: bool,
}
impl<C: Default, I, O> FunctionContext<C, I, O> {
pub fn new(input: I) -> FunctionContext<C, I, O> {
FunctionContext {
context: C::default(),
initial: input,
output: None,
done: false,
}
}
}
impl<C, I, O> FunctionContext<C, I, O> {
pub fn state<F: Fn(&mut C)>(&mut self, f: F) {
f(&mut self.context);
}
pub fn done(&mut self, output: O) {
self.done = true;
self.output = Some(output);
}
}
pub struct RunningIteratedFunction<C, I, O> {
output: std::marker::PhantomData<O>,
context: FunctionContext<C, I, O>,
f: Arc<Box<dyn Fn(FunctionContext<C, I, O>) -> FunctionContext<C, I, O>>>,
}
impl<C, I, O> RunningIteratedFunction<C, I, O> {
pub fn next(mut self) -> Self {
let new_ctx = self.f.as_ref()(self.context);
self.context = new_ctx;
return self;
}
pub fn return_value(&self) -> Option<&O> {
self.context.output.as_ref()
}
pub fn take_return_value(self) -> O {
self.context.output.unwrap()
}
pub fn is_done(&self) -> bool {
self.context.done
}
}
pub struct ThreadRunningIteratedFunction<C: Send + 'static, I: Send + 'static, O: Send + 'static>(
JobDispatcher<IteratedFnQuery, IteratedFnOutput<C, I, O>>,
);
impl<C: Send + 'static, I: Send + 'static, O: Send + 'static>
ThreadRunningIteratedFunction<C, I, O>
{
pub fn start(&self) {
let _ = self.0.try_send(IteratedFnQuery::Start);
}
pub fn pause(&self) {
let _ = self.0.try_send(IteratedFnQuery::Pause);
}
pub fn stop(&self) -> FunctionContext<C, I, O> {
match self.0.try_send(IteratedFnQuery::Stop) {
Some(IteratedFnOutput::Context(function_context)) => function_context,
_ => unreachable!(),
}
}
pub fn try_output(&self) -> Option<O> {
match self.0.send(IteratedFnQuery::GetOutput) {
IteratedFnOutput::Out(out) => Some(out),
_ => None,
}
}
pub fn output(&self) -> O {
retry(|| self.try_output())
}
}
#[derive(Debug, Clone)]
pub enum IteratedFnQuery {
Pause,
Start,
Stop,
GetOutput,
}
pub enum IteratedFnOutput<C, I, O> {
Out(O),
Context(FunctionContext<C, I, O>),
Ok,
}
pub struct IteratedFunction<C, I, O> {
output: std::marker::PhantomData<O>,
f: Arc<Box<dyn Fn(FunctionContext<C, I, O>) -> FunctionContext<C, I, O>>>,
}
impl<C: Default + Clone + Send, I: Clone + Send + 'static, O: Send + Clone>
IteratedFunction<C, I, O>
{
pub fn new<F: Fn(FunctionContext<C, I, O>) -> FunctionContext<C, I, O> + 'static>(
f: F,
) -> Self {
Self {
f: Arc::new(Box::new(f)),
output: std::marker::PhantomData,
}
}
pub fn new_threaded<
F: Fn(FunctionContext<C, I, O>) -> FunctionContext<C, I, O> + Send + 'static,
>(
f: F,
input: I,
) -> ThreadRunningIteratedFunction<C, I, O> {
Self::new_threaded_from_state(f, FunctionContext::new(input))
}
pub fn new_threaded_from_state<
F: Fn(FunctionContext<C, I, O>) -> FunctionContext<C, I, O> + Send + 'static,
>(
f: F,
context: FunctionContext<C, I, O>,
) -> ThreadRunningIteratedFunction<C, I, O> {
let (dispatch, recv) = JobDispatcher::<IteratedFnQuery, IteratedFnOutput<C, I, O>>::new();
let _ = std::thread::spawn(move || {
let f = Self::new(f);
let mut f = f.call_with_context(context);
let mut counter = 0;
let mut sleep = false;
while !f.is_done() {
if sleep {
std::thread::sleep(Duration::from_secs(3));
}
if counter == 5 || sleep {
if let Ok(request) = recv.recv_timeout(Duration::from_millis(300)) {
match request.param {
IteratedFnQuery::Pause => {
log::info!("Paused threaded iterative function");
sleep = true;
}
IteratedFnQuery::Start => {
log::info!("Restarted threaded iterative function");
sleep = false;
}
IteratedFnQuery::Stop => {
log::info!("Ending threaded iterative function");
request.done(IteratedFnOutput::Context(f.context));
return;
}
_ => {}
}
request.done(IteratedFnOutput::Ok);
}
counter = 0;
}
if !sleep {
f = f.next();
}
counter += 1;
}
if f.is_done() {
while let Ok(request) = recv.recv() {
match request.param {
IteratedFnQuery::Stop => {
log::warn!("Function was asked to stop but was already done");
request.done(IteratedFnOutput::Context(f.context.clone()));
}
IteratedFnQuery::GetOutput => {
request.done(IteratedFnOutput::Out(f.context.output.clone().unwrap()));
break;
}
_ => {
request.done(IteratedFnOutput::Out(f.context.output.clone().unwrap()));
}
}
}
}
});
ThreadRunningIteratedFunction(dispatch)
}
pub fn call_with_context(
&self,
ctx: FunctionContext<C, I, O>,
) -> RunningIteratedFunction<C, I, O> {
RunningIteratedFunction {
output: std::marker::PhantomData,
context: ctx,
f: self.f.clone(),
}
}
pub fn call(&self, input: I) -> RunningIteratedFunction<C, I, O> {
RunningIteratedFunction {
output: std::marker::PhantomData,
context: FunctionContext::new(input),
f: self.f.clone(),
}
}
pub fn run_to_end(&self, input: I) -> O {
let mut f = self.call(input);
while !f.is_done() {
f = f.next();
}
return f.take_return_value();
}
}

View file

@ -1,8 +1,12 @@
use std::{sync::mpsc, thread, time::Instant};
mod defer;
pub mod iterated;
pub mod job;
pub mod service;
pub use comrade_macro::worker;
pub use defer::Defer;
pub use comrade_macro::{defer, worker};
pub use crossbeam;
use dashmap::DashMap;
use once_cell::sync::Lazy;
@ -10,16 +14,6 @@ pub use serde_json;
// TODO : worker docs + refactor
// TODO : functions which can be stopped, paused, etc
/*
Example:
let myf = Function::new(|| do_something());
// stop fn
myf.stop();
*/
pub static UNION: Lazy<
DashMap<&'static str, job::JobMultiplexer<serde_json::Value, serde_json::Value>>,
> = Lazy::new(DashMap::new);
@ -63,40 +57,15 @@ where
(fastest_item, fastest_result)
}
// TODO : async version
/*
pub fn rally_async<T: Send + Sync + 'static, F, X: Send + 'static>(items: Vec<T>, f: F) -> (T, X)
where
F: AsyncFn(&T) -> X + Send + Sync + Copy + 'static,
{
let (tx, rx) = mpsc::channel();
let mut handles = Vec::new();
for item in items {
let tx = tx.clone();
let item_ref = item;
let f = f;
tokio::task::spawn()
let handle = thread::spawn(move || {
let start = Instant::now();
let result = f(&item_ref);
let elapsed = start.elapsed();
let _ = tx.send((item_ref, result, elapsed));
});
handles.push(handle);
pub fn retry<O, F: Fn() -> Option<O>>(f: F) -> O {
loop {
match f() {
Some(resp) => {
return resp;
}
None => {
log::info!("Got nothing, retrying...");
}
}
}
drop(tx);
let (fastest_item, fastest_result, _) = rx.recv().unwrap();
for handle in handles {
handle.thread().unpark();
}
(fastest_item, fastest_result)
}
*/