106 lines
3.2 KiB
Rust
106 lines
3.2 KiB
Rust
//! Tests for spawner functionality
|
|
|
|
use crate::common::task_fn;
|
|
use dagx::{DagResult, DagRunner};
|
|
use futures::FutureExt;
|
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
|
use std::sync::Arc;
|
|
use std::sync::Mutex;
|
|
use std::time::Duration;
|
|
use tokio::time::sleep;
|
|
|
|
#[tokio::test]
|
|
async fn test_spawner_actually_spawns_tasks() -> DagResult<()> {
|
|
// Verify that the spawner function is actually called for each task
|
|
let dag = DagRunner::new();
|
|
let spawn_count = Arc::new(AtomicUsize::new(0));
|
|
|
|
// Create 10 independent tasks
|
|
let tasks: Vec<_> = (0..10)
|
|
.map(|i| {
|
|
dag.add_task(task_fn({
|
|
let value = spawn_count.clone();
|
|
move |_: ()| {
|
|
let value = value.clone();
|
|
async move {
|
|
value.fetch_add(1, Ordering::SeqCst);
|
|
i
|
|
}
|
|
}
|
|
}))
|
|
})
|
|
.collect();
|
|
|
|
// Custom spawner that counts invocations
|
|
dag.run(|fut| tokio::spawn(fut).map(Result::unwrap)).await?;
|
|
|
|
// Verify all tasks were spawned
|
|
for (i, task) in tasks.into_iter().enumerate() {
|
|
assert_eq!(dag.get(task)?, i as i32);
|
|
}
|
|
|
|
// The spawner should have been called once for each task
|
|
let spawns = spawn_count.load(Ordering::SeqCst);
|
|
assert_eq!(
|
|
spawns, 10,
|
|
"Expected spawner to be called 10 times, but was called {} times",
|
|
spawns
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_tasks_not_running_inline() -> DagResult<()> {
|
|
// Verify that tasks are spawned and CAN run on different threads than the caller
|
|
// Note: Small/fast tasks may run on the same thread if the scheduler chooses
|
|
use std::thread;
|
|
|
|
let dag = DagRunner::new();
|
|
let calling_thread = thread::current().id();
|
|
let task_threads = Arc::new(Mutex::new(Vec::new()));
|
|
|
|
// Create tasks with some work to encourage distribution
|
|
let tasks: Vec<_> = (0..20)
|
|
.map(|i| {
|
|
let threads = task_threads.clone();
|
|
dag.add_task(task_fn(move |_: ()| {
|
|
let threads = threads.clone();
|
|
async move {
|
|
threads.lock().unwrap().push(thread::current().id());
|
|
|
|
// Add some async work
|
|
sleep(Duration::from_millis(2)).await;
|
|
|
|
i
|
|
}
|
|
}))
|
|
})
|
|
.collect();
|
|
|
|
dag.run(|fut| tokio::spawn(fut).map(Result::unwrap)).await?;
|
|
|
|
// Verify results
|
|
for (i, task) in tasks.into_iter().enumerate() {
|
|
assert_eq!(dag.get(task)?, i as i32);
|
|
}
|
|
|
|
// Check if tasks ran on different threads than the caller
|
|
// This is a capability test - the scheduler MAY choose to use the same thread
|
|
let task_thread_ids = task_threads.lock().unwrap();
|
|
let different_threads = task_thread_ids
|
|
.iter()
|
|
.filter(|&&tid| tid != calling_thread)
|
|
.count();
|
|
|
|
if different_threads == 0 {
|
|
eprintln!(
|
|
"Warning: All tasks ran on the calling thread. This may indicate tasks aren't \
|
|
being spawned, or the scheduler chose not to distribute (valid for very fast tasks)"
|
|
);
|
|
}
|
|
|
|
// The key thing is that tasks completed successfully via spawner
|
|
Ok(())
|
|
}
|