463 lines
15 KiB
Rust
463 lines
15 KiB
Rust
//! Tests for multi-threading and concurrent execution
|
|
|
|
use crate::common::task_fn;
|
|
use dagx::{DagResult, DagRunner, TaskHandle};
|
|
use futures::FutureExt;
|
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
|
use std::sync::{Arc, Mutex};
|
|
use std::time::{Duration, Instant};
|
|
use tokio::time::sleep;
|
|
|
|
#[tokio::test]
|
|
async fn test_tasks_run_on_different_threads() -> DagResult<()> {
|
|
// Prove that tasks CAN run on multiple different threads (true parallelism capability)
|
|
// Note: Tokio may choose to run fast tasks on the same thread for efficiency
|
|
use std::thread;
|
|
|
|
let dag = DagRunner::new();
|
|
let thread_ids = Arc::new(Mutex::new(std::collections::HashSet::new()));
|
|
|
|
// Create 50 tasks with actual work to encourage multi-threading
|
|
let tasks: Vec<_> = (0..50)
|
|
.map(|i| {
|
|
let ids = thread_ids.clone();
|
|
dag.add_task(task_fn(move |_: ()| {
|
|
let ids = ids.clone();
|
|
async move {
|
|
let tid = thread::current().id();
|
|
ids.lock().unwrap().insert(tid);
|
|
|
|
// Do some actual async work to give tokio reason to distribute
|
|
sleep(Duration::from_millis(5)).await;
|
|
|
|
// Some CPU work too
|
|
let mut sum = 0;
|
|
for j in 0..1000 {
|
|
sum += j;
|
|
}
|
|
i + sum
|
|
}
|
|
}))
|
|
})
|
|
.collect();
|
|
|
|
dag.run(|fut| tokio::spawn(fut).map(Result::unwrap)).await?;
|
|
|
|
// Verify all tasks completed
|
|
for task in tasks.into_iter() {
|
|
let _ = dag.get(task)?;
|
|
}
|
|
|
|
// With a multi-threaded runtime and sufficient work, we typically see multiple threads
|
|
// But we can't guarantee it (scheduler's choice), so we just verify the capability exists
|
|
let unique_threads = thread_ids.lock().unwrap().len();
|
|
|
|
// On a multi-core system with sufficient work, tokio should use multiple threads
|
|
// This is a soft assertion - if it fails, it means tasks CAN'T use multiple threads
|
|
if unique_threads == 1 {
|
|
eprintln!(
|
|
"Warning: All tasks ran on 1 thread. This may indicate parallelism isn't working, \
|
|
or tokio chose not to distribute work (valid for small/fast tasks)"
|
|
);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_concurrent_execution_with_atomic_counter() -> DagResult<()> {
|
|
// Prove concurrent execution using atomic counters to track simultaneous tasks
|
|
let dag = DagRunner::new();
|
|
let concurrent_count = Arc::new(AtomicUsize::new(0));
|
|
let max_concurrent = Arc::new(AtomicUsize::new(0));
|
|
|
|
// Create 20 tasks that track concurrent execution
|
|
let tasks: Vec<_> = (0..20)
|
|
.map(|i| {
|
|
let count = concurrent_count.clone();
|
|
let max = max_concurrent.clone();
|
|
|
|
dag.add_task(task_fn(move |_: ()| {
|
|
let count = count.clone();
|
|
let max = max.clone();
|
|
async move {
|
|
// Increment concurrent counter
|
|
let current = count.fetch_add(1, Ordering::SeqCst) + 1;
|
|
|
|
// Track maximum concurrency
|
|
let mut prev_max = max.load(Ordering::SeqCst);
|
|
while current > prev_max {
|
|
match max.compare_exchange_weak(
|
|
prev_max,
|
|
current,
|
|
Ordering::SeqCst,
|
|
Ordering::SeqCst,
|
|
) {
|
|
Ok(_) => break,
|
|
Err(x) => prev_max = x,
|
|
}
|
|
}
|
|
|
|
// Simulate work
|
|
sleep(Duration::from_millis(10)).await;
|
|
|
|
// Decrement counter
|
|
count.fetch_sub(1, Ordering::SeqCst);
|
|
|
|
i
|
|
}
|
|
}))
|
|
})
|
|
.collect();
|
|
|
|
dag.run(|fut| tokio::spawn(fut).map(Result::unwrap)).await?;
|
|
|
|
// Verify all tasks completed
|
|
for (i, task) in tasks.into_iter().enumerate() {
|
|
assert_eq!(dag.get(task)?, i);
|
|
}
|
|
|
|
// Check that we achieved real parallelism
|
|
let max_seen = max_concurrent.load(Ordering::SeqCst);
|
|
println!("Maximum concurrent tasks: {}", max_seen);
|
|
|
|
// Should see multiple tasks running concurrently
|
|
assert!(
|
|
max_seen > 1,
|
|
"Expected concurrent execution, max was {}",
|
|
max_seen
|
|
);
|
|
|
|
// Final count should be 0
|
|
assert_eq!(concurrent_count.load(Ordering::SeqCst), 0);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_diamond_parallel_execution() -> DagResult<()> {
|
|
// Test diamond pattern: source -> (parallel1, parallel2) -> sink
|
|
let dag = DagRunner::new();
|
|
|
|
let concurrent_middle = Arc::new(AtomicUsize::new(0));
|
|
|
|
let source: TaskHandle<_> = dag.add_task(task_fn(|_: ()| async { 100 })).into();
|
|
|
|
// These should run in parallel
|
|
let parallel1 = {
|
|
let counter = concurrent_middle.clone();
|
|
dag.add_task(task_fn(move |x: i32| {
|
|
let counter = counter.clone();
|
|
async move {
|
|
counter.fetch_add(1, Ordering::SeqCst);
|
|
sleep(Duration::from_millis(20)).await;
|
|
let count = counter.load(Ordering::SeqCst);
|
|
counter.fetch_sub(1, Ordering::SeqCst);
|
|
assert!(count > 0, "Should see concurrent execution");
|
|
x * 2
|
|
}
|
|
}))
|
|
.depends_on(source)
|
|
};
|
|
|
|
let parallel2 = {
|
|
let counter = concurrent_middle.clone();
|
|
dag.add_task(task_fn(move |x: i32| {
|
|
let counter = counter.clone();
|
|
async move {
|
|
counter.fetch_add(1, Ordering::SeqCst);
|
|
sleep(Duration::from_millis(20)).await;
|
|
let count = counter.load(Ordering::SeqCst);
|
|
counter.fetch_sub(1, Ordering::SeqCst);
|
|
// Both should be running
|
|
assert!(count > 0, "Should see concurrent execution");
|
|
x * 3
|
|
}
|
|
}))
|
|
.depends_on(source)
|
|
};
|
|
|
|
let sink = dag
|
|
.add_task(task_fn(|(a, b): (i32, i32)| async move { a + b }))
|
|
.depends_on((¶llel1, ¶llel2));
|
|
|
|
dag.run(|fut| tokio::spawn(fut).map(Result::unwrap)).await?;
|
|
|
|
assert_eq!(dag.get(sink)?, 500); // (100 * 2) + (100 * 3)
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_parallel_branches_with_dependencies() -> DagResult<()> {
|
|
// Test that parallel branches execute independently but respect their own dependencies
|
|
let dag = DagRunner::new();
|
|
|
|
let branch_progress = Arc::new(Mutex::new(Vec::new()));
|
|
|
|
// Source node
|
|
let source: TaskHandle<_> = dag.add_task(task_fn(|_: ()| async { 10 })).into();
|
|
|
|
// Branch 1: source -> b1_1 -> b1_2
|
|
let b1_1 = {
|
|
let progress = branch_progress.clone();
|
|
dag.add_task(task_fn(move |x: i32| {
|
|
let progress = progress.clone();
|
|
async move {
|
|
progress
|
|
.lock()
|
|
.unwrap()
|
|
.push(("b1_1_start", Instant::now()));
|
|
sleep(Duration::from_millis(30)).await;
|
|
progress.lock().unwrap().push(("b1_1_end", Instant::now()));
|
|
x * 2
|
|
}
|
|
}))
|
|
.depends_on(source)
|
|
};
|
|
|
|
let b1_2 = {
|
|
let progress = branch_progress.clone();
|
|
dag.add_task(task_fn(move |x: i32| {
|
|
let progress = progress.clone();
|
|
async move {
|
|
progress
|
|
.lock()
|
|
.unwrap()
|
|
.push(("b1_2_start", Instant::now()));
|
|
sleep(Duration::from_millis(10)).await;
|
|
progress.lock().unwrap().push(("b1_2_end", Instant::now()));
|
|
x + 1
|
|
}
|
|
}))
|
|
.depends_on(b1_1)
|
|
};
|
|
|
|
// Branch 2: source -> b2_1 -> b2_2
|
|
let b2_1 = {
|
|
let progress = branch_progress.clone();
|
|
dag.add_task(task_fn(move |x: i32| {
|
|
let progress = progress.clone();
|
|
async move {
|
|
progress
|
|
.lock()
|
|
.unwrap()
|
|
.push(("b2_1_start", Instant::now()));
|
|
sleep(Duration::from_millis(10)).await;
|
|
progress.lock().unwrap().push(("b2_1_end", Instant::now()));
|
|
x * 3
|
|
}
|
|
}))
|
|
.depends_on(source)
|
|
};
|
|
|
|
let b2_2 = {
|
|
let progress = branch_progress.clone();
|
|
dag.add_task(task_fn(move |x: i32| {
|
|
let progress = progress.clone();
|
|
async move {
|
|
progress
|
|
.lock()
|
|
.unwrap()
|
|
.push(("b2_2_start", Instant::now()));
|
|
sleep(Duration::from_millis(30)).await;
|
|
progress.lock().unwrap().push(("b2_2_end", Instant::now()));
|
|
x + 2
|
|
}
|
|
}))
|
|
.depends_on(b2_1)
|
|
};
|
|
|
|
dag.run(|fut| tokio::spawn(fut).map(Result::unwrap)).await?;
|
|
|
|
// Verify results
|
|
assert_eq!(dag.get(b1_2)?, 21); // (10 * 2) + 1
|
|
assert_eq!(dag.get(b2_2)?, 32); // (10 * 3) + 2
|
|
|
|
// Analyze timing to prove parallel execution
|
|
let progress = branch_progress.lock().unwrap();
|
|
|
|
// Find timestamps
|
|
let find_time = |name: &str| -> Instant {
|
|
progress
|
|
.iter()
|
|
.find(|(n, _)| *n == name)
|
|
.map(|(_, t)| *t)
|
|
.unwrap_or_else(|| panic!("Couldn't find {}", name))
|
|
};
|
|
|
|
let b1_1_start = find_time("b1_1_start");
|
|
let b1_1_end = find_time("b1_1_end");
|
|
let b1_2_start = find_time("b1_2_start");
|
|
let b2_1_start = find_time("b2_1_start");
|
|
|
|
// Branch 1 and Branch 2 first tasks should start nearly simultaneously
|
|
let branch_start_diff = if b1_1_start > b2_1_start {
|
|
b1_1_start - b2_1_start
|
|
} else {
|
|
b2_1_start - b1_1_start
|
|
};
|
|
|
|
// Use a more generous threshold to account for coverage instrumentation overhead
|
|
// 50ms is enough to prove they're parallel (not sequential) while being robust to overhead
|
|
assert!(
|
|
branch_start_diff < Duration::from_millis(50),
|
|
"Branches should start in parallel, diff was {:?}",
|
|
branch_start_diff
|
|
);
|
|
|
|
// b1_2 should only start after b1_1 ends
|
|
assert!(b1_2_start >= b1_1_end, "b1_2 started before b1_1 finished");
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_wide_fanout_parallel_execution() -> DagResult<()> {
|
|
// Test that a single source can fan out to many parallel tasks
|
|
let dag = DagRunner::new();
|
|
|
|
let max_concurrent = Arc::new(AtomicUsize::new(0));
|
|
let current_concurrent = Arc::new(AtomicUsize::new(0));
|
|
|
|
let source: TaskHandle<_> = dag.add_task(task_fn(|_: ()| async { 42 })).into();
|
|
|
|
// Create 50 tasks that all depend on the source
|
|
let dependents: Vec<_> = (0..50)
|
|
.map(|i| {
|
|
let max_c = max_concurrent.clone();
|
|
let curr_c = current_concurrent.clone();
|
|
|
|
dag.add_task(task_fn(move |x: i32| {
|
|
let max_c = max_c.clone();
|
|
let curr_c = curr_c.clone();
|
|
async move {
|
|
// Track concurrency
|
|
let current = curr_c.fetch_add(1, Ordering::SeqCst) + 1;
|
|
|
|
// Update max
|
|
let mut prev_max = max_c.load(Ordering::SeqCst);
|
|
while current > prev_max {
|
|
match max_c.compare_exchange_weak(
|
|
prev_max,
|
|
current,
|
|
Ordering::SeqCst,
|
|
Ordering::SeqCst,
|
|
) {
|
|
Ok(_) => break,
|
|
Err(v) => prev_max = v,
|
|
}
|
|
}
|
|
|
|
// Do work
|
|
sleep(Duration::from_millis(5)).await;
|
|
|
|
curr_c.fetch_sub(1, Ordering::SeqCst);
|
|
x + i
|
|
}
|
|
}))
|
|
.depends_on(source)
|
|
})
|
|
.collect();
|
|
|
|
dag.run(|fut| tokio::spawn(fut).map(Result::unwrap)).await?;
|
|
|
|
// Verify results
|
|
for (i, task) in dependents.iter().enumerate() {
|
|
assert_eq!(dag.get(task)?, 42 + i as i32);
|
|
}
|
|
|
|
// Check that we saw significant parallelism
|
|
let max_seen = max_concurrent.load(Ordering::SeqCst);
|
|
println!("Max concurrent in fanout: {}", max_seen);
|
|
assert!(
|
|
max_seen > 5,
|
|
"Expected significant parallelism in fanout, got {}",
|
|
max_seen
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_massive_parallel_fanout() -> DagResult<()> {
|
|
// Test extreme parallelism with 1000 parallel tasks
|
|
let dag = DagRunner::new();
|
|
|
|
let completed = Arc::new(AtomicUsize::new(0));
|
|
let max_concurrent = Arc::new(AtomicUsize::new(0));
|
|
let current = Arc::new(AtomicUsize::new(0));
|
|
|
|
let source: TaskHandle<_> = dag.add_task(task_fn(|_: ()| async { 1 })).into();
|
|
|
|
// Create 1000 tasks all depending on the same source
|
|
let tasks: Vec<_> = (0..1000)
|
|
.map(|i| {
|
|
let comp = completed.clone();
|
|
let max = max_concurrent.clone();
|
|
let curr = current.clone();
|
|
|
|
dag.add_task(task_fn(move |x: i32| {
|
|
let comp = comp.clone();
|
|
let max = max.clone();
|
|
let curr = curr.clone();
|
|
async move {
|
|
// Track concurrency
|
|
let c = curr.fetch_add(1, Ordering::SeqCst) + 1;
|
|
|
|
let mut prev_max = max.load(Ordering::SeqCst);
|
|
while c > prev_max {
|
|
match max.compare_exchange_weak(
|
|
prev_max,
|
|
c,
|
|
Ordering::SeqCst,
|
|
Ordering::SeqCst,
|
|
) {
|
|
Ok(_) => break,
|
|
Err(v) => prev_max = v,
|
|
}
|
|
}
|
|
|
|
// Quick work
|
|
sleep(Duration::from_micros(100)).await;
|
|
|
|
curr.fetch_sub(1, Ordering::SeqCst);
|
|
comp.fetch_add(1, Ordering::SeqCst);
|
|
x + i
|
|
}
|
|
}))
|
|
.depends_on(source)
|
|
})
|
|
.collect();
|
|
|
|
let start = Instant::now();
|
|
dag.run(|fut| tokio::spawn(fut).map(Result::unwrap)).await?;
|
|
let elapsed = start.elapsed();
|
|
|
|
// All should complete
|
|
assert_eq!(completed.load(Ordering::SeqCst), 1000);
|
|
|
|
// Check some results
|
|
assert_eq!(dag.get(tasks[0])?, 1);
|
|
assert_eq!(dag.get(tasks[500])?, 501);
|
|
assert_eq!(dag.get(tasks[999])?, 1000);
|
|
|
|
let max_seen = max_concurrent.load(Ordering::SeqCst);
|
|
println!(
|
|
"Massive fanout: {} concurrent tasks, completed in {:?}",
|
|
max_seen, elapsed
|
|
);
|
|
|
|
// Should see significant parallelism
|
|
assert!(max_seen > 10, "Expected high parallelism, got {}", max_seen);
|
|
|
|
// Should complete quickly (not sequentially)
|
|
assert!(
|
|
elapsed < Duration::from_secs(2),
|
|
"1000 tasks took too long: {:?}",
|
|
elapsed
|
|
);
|
|
|
|
Ok(())
|
|
}
|