use std::future::Future;
use std::pin::Pin;
use tokio::task::JoinSet;
use crate::error::Error;
use super::ctx::WorkflowCtx;
pub type BoxThunk<R> = Box<
dyn FnOnce() -> Pin<Box<dyn Future<Output = Result<R, Error>> + Send + 'static>>
+ Send
+ 'static,
>;
pub fn thunk<R, F, Fut>(f: F) -> BoxThunk<R>
where
F: FnOnce() -> Fut + Send + 'static,
Fut: Future<Output = Result<R, Error>> + Send + 'static,
R: Send + 'static,
{
Box::new(move || {
Box::pin(f()) as Pin<Box<dyn Future<Output = Result<R, Error>> + Send + 'static>>
})
}
pub async fn parallel<R, F, Fut>(_ctx: &WorkflowCtx, thunks: Vec<F>) -> Vec<Option<R>>
where
F: FnOnce() -> Fut + Send + 'static,
Fut: Future<Output = Result<R, Error>> + Send + 'static,
R: Send + 'static,
{
let n = thunks.len();
let mut set = JoinSet::new();
for (i, thunk) in thunks.into_iter().enumerate() {
set.spawn(async move { (i, thunk().await) });
}
let mut out: Vec<Option<R>> = (0..n).map(|_| None).collect();
while let Some(joined) = set.join_next().await {
match joined {
Ok((i, Ok(value))) => out[i] = Some(value),
Ok((_i, Err(_e))) => {}
Err(_join_err) => {}
}
}
out
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use std::time::Duration;
use super::*;
use crate::agent::test_helpers::MockProvider;
use crate::llm::BoxedProvider;
fn ctx() -> WorkflowCtx {
WorkflowCtx::builder(Arc::new(BoxedProvider::new(MockProvider::new(vec![]))))
.build()
.expect("build ctx")
}
#[tokio::test]
async fn results_in_submission_order_despite_completion_order() {
let ctx = ctx();
let thunks: Vec<BoxThunk<i32>> = vec![
thunk(|| async {
tokio::time::sleep(Duration::from_millis(40)).await;
Ok(10)
}),
thunk(|| async {
tokio::time::sleep(Duration::from_millis(20)).await;
Ok(20)
}),
thunk(|| async { Ok(30) }),
];
let out = parallel(&ctx, thunks).await;
assert_eq!(out, vec![Some(10), Some(20), Some(30)]);
}
#[tokio::test]
async fn errored_thunk_yields_none_others_unaffected() {
let ctx = ctx();
let thunks: Vec<BoxThunk<i32>> = vec![
thunk(|| async { Ok(1) }),
thunk(|| async { Err(Error::Agent("boom".into())) }),
thunk(|| async { Ok(3) }),
];
let out = parallel(&ctx, thunks).await;
assert_eq!(out, vec![Some(1), None, Some(3)]);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn panicking_thunk_is_isolated() {
let ctx = ctx();
let thunks: Vec<BoxThunk<i32>> = vec![
thunk(|| async { Ok(1) }),
thunk(|| async { panic!("intentional test panic") }),
thunk(|| async { Ok(3) }),
];
let out = parallel(&ctx, thunks).await;
assert_eq!(out, vec![Some(1), None, Some(3)]);
}
#[tokio::test]
async fn homogeneous_map_fanout_needs_no_boxing() {
let ctx = ctx();
let work = vec![1, 2, 3, 4];
let thunks: Vec<_> = work
.into_iter()
.map(|x| move || async move { Ok::<i32, Error>(x * 100) })
.collect();
let out = parallel(&ctx, thunks).await;
assert_eq!(out, vec![Some(100), Some(200), Some(300), Some(400)]);
}
#[tokio::test]
async fn empty_thunks_returns_empty() {
let ctx = ctx();
let thunks: Vec<BoxThunk<i32>> = Vec::new();
let out = parallel(&ctx, thunks).await;
assert!(out.is_empty());
}
}