374 lines
11 KiB
Rust
374 lines
11 KiB
Rust
//! Tests for execution correctness and parallelism
|
|
|
|
use crate::common::task_fn;
|
|
use dagx::{DagResult, DagRunner};
|
|
use std::sync::{Arc, Mutex};
|
|
use std::time::{Duration, Instant};
|
|
use tokio::time::sleep;
|
|
|
|
#[tokio::test]
|
|
async fn test_layers_execute_in_parallel() -> DagResult<()> {
|
|
// Test that tasks in the same topological layer execute in parallel
|
|
let dag = DagRunner::new();
|
|
|
|
let layer_timing = Arc::new(Mutex::new(Vec::new()));
|
|
|
|
// Layer 0: 3 independent sources
|
|
let sources: Vec<_> = (0..3)
|
|
.map(|i| {
|
|
let timing = layer_timing.clone();
|
|
dag.add_task(task_fn(move |_: ()| {
|
|
let timing = timing.clone();
|
|
async move {
|
|
let start = Instant::now();
|
|
timing.lock().unwrap().push((format!("L0_{}", i), start));
|
|
sleep(Duration::from_millis(20)).await;
|
|
i * 10
|
|
}
|
|
}))
|
|
})
|
|
.collect();
|
|
|
|
// Layer 1: Tasks depending on layer 0
|
|
let layer1: Vec<_> = sources
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(i, source)| {
|
|
let timing = layer_timing.clone();
|
|
dag.add_task(task_fn(move |x: i32| {
|
|
let timing = timing.clone();
|
|
async move {
|
|
let start = Instant::now();
|
|
timing.lock().unwrap().push((format!("L1_{}", i), start));
|
|
sleep(Duration::from_millis(20)).await;
|
|
x + 1
|
|
}
|
|
}))
|
|
.depends_on(source)
|
|
})
|
|
.collect();
|
|
|
|
// Layer 2: Final aggregation
|
|
let final_task = {
|
|
let timing = layer_timing.clone();
|
|
dag.add_task(task_fn(move |(a, b, c): (i32, i32, i32)| {
|
|
let timing = timing.clone();
|
|
async move {
|
|
let start = Instant::now();
|
|
timing.lock().unwrap().push(("L2_final".to_string(), start));
|
|
a + b + c
|
|
}
|
|
}))
|
|
.depends_on((&layer1[0], &layer1[1], &layer1[2]))
|
|
};
|
|
|
|
dag.run(|fut| {
|
|
tokio::spawn(fut);
|
|
})
|
|
.await?;
|
|
|
|
assert_eq!(dag.get(final_task)?, 33); // (0+1) + (10+1) + (20+1)
|
|
|
|
// Analyze timing
|
|
let timing = layer_timing.lock().unwrap();
|
|
|
|
// Get layer 0 start times
|
|
let l0_times: Vec<Instant> = (0..3)
|
|
.map(|i| {
|
|
timing
|
|
.iter()
|
|
.find(|(name, _)| name == &format!("L0_{}", i))
|
|
.map(|(_, t)| *t)
|
|
.unwrap()
|
|
})
|
|
.collect();
|
|
|
|
// Get layer 1 start times
|
|
let l1_times: Vec<Instant> = (0..3)
|
|
.map(|i| {
|
|
timing
|
|
.iter()
|
|
.find(|(name, _)| name == &format!("L1_{}", i))
|
|
.map(|(_, t)| *t)
|
|
.unwrap()
|
|
})
|
|
.collect();
|
|
|
|
// Layer 0 tasks should all start nearly simultaneously
|
|
let l0_spread = l0_times
|
|
.iter()
|
|
.max()
|
|
.unwrap()
|
|
.duration_since(*l0_times.iter().min().unwrap());
|
|
assert!(
|
|
l0_spread < Duration::from_millis(5),
|
|
"Layer 0 tasks should start together, spread was {:?}",
|
|
l0_spread
|
|
);
|
|
|
|
// Layer 1 tasks should also start nearly simultaneously (after layer 0 completes)
|
|
let l1_spread = l1_times
|
|
.iter()
|
|
.max()
|
|
.unwrap()
|
|
.duration_since(*l1_times.iter().min().unwrap());
|
|
assert!(
|
|
l1_spread < Duration::from_millis(5),
|
|
"Layer 1 tasks should start together, spread was {:?}",
|
|
l1_spread
|
|
);
|
|
|
|
// Layer 1 should start after layer 0
|
|
let earliest_l1 = *l1_times.iter().min().unwrap();
|
|
let earliest_l0 = *l0_times.iter().min().unwrap();
|
|
assert!(
|
|
earliest_l1 >= earliest_l0 + Duration::from_millis(15),
|
|
"Layer 1 should start after layer 0 completes"
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_completion_order_matches_dependencies() -> DagResult<()> {
|
|
// Verify that tasks complete in dependency order
|
|
let dag = DagRunner::new();
|
|
|
|
let completion_order = Arc::new(Mutex::new(Vec::new()));
|
|
|
|
// Create a complex DAG with clear ordering requirements
|
|
let a = {
|
|
let order = completion_order.clone();
|
|
dag.add_task(task_fn(move |_: ()| {
|
|
let order = order.clone();
|
|
async move {
|
|
sleep(Duration::from_millis(10)).await;
|
|
order.lock().unwrap().push("A");
|
|
1
|
|
}
|
|
}))
|
|
};
|
|
|
|
let b = {
|
|
let order = completion_order.clone();
|
|
dag.add_task(task_fn(move |_: ()| {
|
|
let order = order.clone();
|
|
async move {
|
|
sleep(Duration::from_millis(5)).await;
|
|
order.lock().unwrap().push("B");
|
|
2
|
|
}
|
|
}))
|
|
};
|
|
|
|
let c = {
|
|
let order = completion_order.clone();
|
|
dag.add_task(task_fn(move |(x, y): (i32, i32)| {
|
|
let order = order.clone();
|
|
async move {
|
|
order.lock().unwrap().push("C");
|
|
x + y
|
|
}
|
|
}))
|
|
.depends_on((&a, &b))
|
|
};
|
|
|
|
let d = {
|
|
let order = completion_order.clone();
|
|
dag.add_task(task_fn(move |x: i32| {
|
|
let order = order.clone();
|
|
async move {
|
|
order.lock().unwrap().push("D");
|
|
x * 2
|
|
}
|
|
}))
|
|
.depends_on(c)
|
|
};
|
|
|
|
dag.run(|fut| {
|
|
tokio::spawn(fut);
|
|
})
|
|
.await?;
|
|
|
|
assert_eq!(dag.get(d)?, 6); // (1 + 2) * 2
|
|
|
|
// Check completion order
|
|
let order = completion_order.lock().unwrap();
|
|
|
|
// A and B can complete in any order (parallel)
|
|
assert!(order.contains(&"A"));
|
|
assert!(order.contains(&"B"));
|
|
|
|
// C must come after both A and B
|
|
let c_pos = order.iter().position(|name| *name == "C").unwrap();
|
|
let a_pos = order.iter().position(|name| *name == "A").unwrap();
|
|
let b_pos = order.iter().position(|name| *name == "B").unwrap();
|
|
assert!(c_pos > a_pos, "C should complete after A");
|
|
assert!(c_pos > b_pos, "C should complete after B");
|
|
|
|
// D must come after C
|
|
let d_pos = order.iter().position(|name| *name == "D").unwrap();
|
|
assert!(d_pos > c_pos, "D should complete after C");
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_dependent_tasks_respect_ordering() -> DagResult<()> {
|
|
// Prove that dependent tasks wait for their dependencies
|
|
let dag = DagRunner::new();
|
|
|
|
let execution_order = Arc::new(Mutex::new(Vec::new()));
|
|
|
|
// Create a chain: A -> B -> C -> D
|
|
let a = {
|
|
let order = execution_order.clone();
|
|
dag.add_task(task_fn(move |_: ()| {
|
|
let order = order.clone();
|
|
async move {
|
|
sleep(Duration::from_millis(20)).await;
|
|
order.lock().unwrap().push("A");
|
|
1
|
|
}
|
|
}))
|
|
};
|
|
|
|
let b = {
|
|
let order = execution_order.clone();
|
|
dag.add_task(task_fn(move |x: i32| {
|
|
let order = order.clone();
|
|
async move {
|
|
order.lock().unwrap().push("B");
|
|
sleep(Duration::from_millis(10)).await;
|
|
x + 1
|
|
}
|
|
}))
|
|
.depends_on(&a)
|
|
};
|
|
|
|
let c = {
|
|
let order = execution_order.clone();
|
|
dag.add_task(task_fn(move |x: i32| {
|
|
let order = order.clone();
|
|
async move {
|
|
order.lock().unwrap().push("C");
|
|
sleep(Duration::from_millis(10)).await;
|
|
x + 1
|
|
}
|
|
}))
|
|
.depends_on(b)
|
|
};
|
|
|
|
let d = {
|
|
let order = execution_order.clone();
|
|
dag.add_task(task_fn(move |x: i32| {
|
|
let order = order.clone();
|
|
async move {
|
|
order.lock().unwrap().push("D");
|
|
x + 1
|
|
}
|
|
}))
|
|
.depends_on(c)
|
|
};
|
|
|
|
dag.run(|fut| {
|
|
tokio::spawn(fut);
|
|
})
|
|
.await?;
|
|
|
|
// Verify results
|
|
assert_eq!(dag.get(a)?, 1);
|
|
assert_eq!(dag.get(b)?, 2);
|
|
assert_eq!(dag.get(c)?, 3);
|
|
assert_eq!(dag.get(d)?, 4);
|
|
|
|
// Verify execution order
|
|
let order = execution_order.lock().unwrap();
|
|
let expected = vec!["A", "B", "C", "D"];
|
|
assert_eq!(order.as_slice(), expected.as_slice());
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_independent_tasks_dont_wait() -> DagResult<()> {
|
|
// Prove that independent tasks don't wait for each other
|
|
let dag = DagRunner::new();
|
|
|
|
let task_started = Arc::new(Mutex::new(Vec::new()));
|
|
let task_finished = Arc::new(Mutex::new(Vec::new()));
|
|
|
|
// Create tasks with different durations
|
|
let fast1 = {
|
|
let started = task_started.clone();
|
|
let finished = task_finished.clone();
|
|
dag.add_task(task_fn(move |_: ()| {
|
|
let started = started.clone();
|
|
let finished = finished.clone();
|
|
async move {
|
|
started.lock().unwrap().push(("fast1", Instant::now()));
|
|
sleep(Duration::from_millis(10)).await;
|
|
finished.lock().unwrap().push(("fast1", Instant::now()));
|
|
"fast1"
|
|
}
|
|
}))
|
|
};
|
|
|
|
let slow = {
|
|
let started = task_started.clone();
|
|
let finished = task_finished.clone();
|
|
dag.add_task(task_fn(move |_: ()| {
|
|
let started = started.clone();
|
|
let finished = finished.clone();
|
|
async move {
|
|
started.lock().unwrap().push(("slow", Instant::now()));
|
|
sleep(Duration::from_millis(100)).await;
|
|
finished.lock().unwrap().push(("slow", Instant::now()));
|
|
"slow"
|
|
}
|
|
}))
|
|
};
|
|
|
|
let fast2 = {
|
|
let started = task_started.clone();
|
|
let finished = task_finished.clone();
|
|
dag.add_task(task_fn(move |_: ()| {
|
|
let started = started.clone();
|
|
let finished = finished.clone();
|
|
async move {
|
|
started.lock().unwrap().push(("fast2", Instant::now()));
|
|
sleep(Duration::from_millis(10)).await;
|
|
finished.lock().unwrap().push(("fast2", Instant::now()));
|
|
"fast2"
|
|
}
|
|
}))
|
|
};
|
|
|
|
dag.run(|fut| {
|
|
tokio::spawn(fut);
|
|
})
|
|
.await?;
|
|
|
|
// Check results
|
|
assert_eq!(dag.get(fast1)?, "fast1");
|
|
assert_eq!(dag.get(slow)?, "slow");
|
|
assert_eq!(dag.get(fast2)?, "fast2");
|
|
|
|
// Verify that fast tasks finished before slow task
|
|
let finished = task_finished.lock().unwrap();
|
|
let fast1_finish = finished.iter().find(|x| x.0 == "fast1").unwrap().1;
|
|
let fast2_finish = finished.iter().find(|x| x.0 == "fast2").unwrap().1;
|
|
let slow_finish = finished.iter().find(|x| x.0 == "slow").unwrap().1;
|
|
|
|
assert!(
|
|
fast1_finish < slow_finish,
|
|
"fast1 should finish before slow"
|
|
);
|
|
assert!(
|
|
fast2_finish < slow_finish,
|
|
"fast2 should finish before slow"
|
|
);
|
|
|
|
Ok(())
|
|
}
|