consortium/dag/

executor.rs

//! DAG executor — schedules and runs tasks with maximum parallelism.
//!
//! The executor is a single-threaded scheduler that dispatches tasks to
//! worker threads. It blocks on an mpsc channel waiting for completions,
//! then immediately dispatches any newly-ready tasks. This gives maximum
//! parallelism without busy-waiting: a task starts the instant its
//! dependencies are satisfied and a pool slot is available.

use std::collections::{HashMap, HashSet, VecDeque};
use std::sync::mpsc;
use std::thread;

use crate::dag::cache::CacheStrategy;
use crate::dag::context::DagContext;
use crate::dag::error::{DagError, Result};
use crate::dag::graph::DagGraph;
use crate::dag::pool::WorkerPool;
use crate::dag::types::{ErrorPolicy, TaskId, TaskOutcome};

/// Completion message from a worker thread.
struct CompletionEvent {
    task_id: TaskId,
    outcome: TaskOutcome,
}

/// Events emitted during DAG execution for monitoring.
#[derive(Debug, Clone)]
pub enum DagEvent {
    TaskStarted(TaskId),
    TaskCompleted(TaskId, TaskOutcome),
    TaskSkipped(TaskId),
    TaskCancelled(TaskId),
}

/// Callback trait for monitoring DAG execution progress.
pub trait DagMonitor: Send {
    fn on_event(&mut self, event: &DagEvent);
}

/// Report from a DAG execution run.
#[derive(Debug)]
pub struct DagReport {
    /// Tasks that completed successfully.
    pub completed: HashSet<TaskId>,
    /// Tasks that were skipped (cache hit).
    pub skipped: HashSet<TaskId>,
    /// Tasks that failed, with error messages.
    pub failed: HashMap<TaskId, String>,
    /// Tasks cancelled because a dependency failed.
    pub cancelled: HashSet<TaskId>,
}

impl DagReport {
    /// Whether the entire DAG succeeded (no failures).
    pub fn is_success(&self) -> bool {
        self.failed.is_empty()
    }

    /// Total number of tasks processed (completed + skipped + failed + cancelled).
    pub fn total(&self) -> usize {
        self.completed.len() + self.skipped.len() + self.failed.len() + self.cancelled.len()
    }
}

/// The DAG executor. Constructed via `DagBuilder` or `StageBuilder`.
pub struct DagExecutor {
    pub(crate) graph: DagGraph,
    pub(crate) context: DagContext,
    pub(crate) error_policy: ErrorPolicy,
    pub(crate) pool: Box<dyn WorkerPool>,
    pub(crate) global_cache: Option<Box<dyn CacheStrategy>>,
    pub(crate) monitor: Option<Box<dyn DagMonitor>>,
}

impl DagExecutor {
    /// Execute the DAG to completion.
    ///
    /// Blocks until all tasks finish, are cancelled, or an error policy
    /// aborts execution.
    pub fn run(mut self) -> Result<DagReport> {
        self.graph.validate()?;

        let (tx, rx) = mpsc::channel::<CompletionEvent>();

        // Track unmet dependency count per task
        let mut pending_deps: HashMap<TaskId, usize> = HashMap::new();
        let all_task_ids: Vec<TaskId> = self.graph.tasks.keys().cloned().collect();
        for id in &all_task_ids {
            let dep_count = self
                .graph
                .dependencies
                .get(id)
                .map(|d| d.len())
                .unwrap_or(0);
            pending_deps.insert(id.clone(), dep_count);
        }

        // Concurrency group tracking: group_name -> active count
        let mut group_active: HashMap<String, usize> = HashMap::new();

        let mut ready: VecDeque<TaskId> = VecDeque::new();
        let mut deferred: VecDeque<TaskId> = VecDeque::new(); // tasks waiting for pool/group slots
        let mut in_flight: HashSet<TaskId> = HashSet::new();

        let mut report = DagReport {
            completed: HashSet::new(),
            skipped: HashSet::new(),
            failed: HashMap::new(),
            cancelled: HashSet::new(),
        };

        // Seed with root tasks
        for id in &all_task_ids {
            if pending_deps.get(id).copied().unwrap_or(0) == 0 {
                ready.push_back(id.clone());
            }
        }

        loop {
            // 1. Try to dispatch ready tasks
            let mut retry_ready: VecDeque<TaskId> = VecDeque::new();

            while let Some(task_id) = ready.pop_front() {
                if report.cancelled.contains(&task_id) {
                    continue;
                }
                // Check cache (per-task or global)
                if let Some(cached) = self.check_cache(&task_id) {
                    self.context.set_output(task_id.clone(), cached);
                    self.notify(DagEvent::TaskSkipped(task_id.clone()));
                    report.skipped.insert(task_id.clone());
                    self.mark_completed(&task_id, &mut pending_deps, &mut retry_ready);
                    continue;
                }

                // Check concurrency group limit
                if !self.can_dispatch_group(&task_id, &group_active) {
                    deferred.push_back(task_id);
                    continue;
                }

                // Check worker pool
                let resource_tag = self
                    .graph
                    .tasks
                    .get(&task_id)
                    .and_then(|t| t.resource_tag().map(|s| s.to_string()));
                if !self.pool.acquire(resource_tag.as_deref()) {
                    deferred.push_back(task_id);
                    continue;
                }

                // Dispatch to worker thread
                let task = match self.graph.take_task(&task_id) {
                    Some(t) => t,
                    None => continue, // already taken (shouldn't happen)
                };

                // Update group tracking
                if let Some(group) = self.graph.task_groups.get(&task_id) {
                    *group_active.entry(group.clone()).or_insert(0) += 1;
                }

                let ctx = self.context.clone();
                let tx = tx.clone();
                let id = task_id.clone();
                in_flight.insert(task_id.clone());
                self.notify(DagEvent::TaskStarted(task_id));

                thread::spawn(move || {
                    let outcome = task.execute(&ctx);
                    let _ = tx.send(CompletionEvent {
                        task_id: id,
                        outcome,
                    });
                });
            }

            // Move retry_ready items back to ready
            ready.extend(retry_ready);

            // 2. If nothing in flight and nothing ready/deferred, we're done
            if in_flight.is_empty() && ready.is_empty() && deferred.is_empty() {
                break;
            }

            // 3. If nothing in flight but tasks are deferred, that's a deadlock
            if in_flight.is_empty() && ready.is_empty() && !deferred.is_empty() {
                // This can happen if pool limits are too tight
                return Err(DagError::General(format!(
                    "deadlock: {} deferred tasks but nothing in flight",
                    deferred.len()
                )));
            }

            // 4. Wait for a completion event
            let event = rx.recv().map_err(|_| DagError::ChannelClosed)?;
            let CompletionEvent { task_id, outcome } = event;

            in_flight.remove(&task_id);

            // Release pool slot
            self.pool.release(None);

            // Release concurrency group slot
            if let Some(group) = self.graph.task_groups.get(&task_id) {
                if let Some(count) = group_active.get_mut(group) {
                    *count = count.saturating_sub(1);
                }
            }

            let is_success = outcome.is_success();
            let fail_msg = match &outcome {
                TaskOutcome::Failed(msg) => Some(msg.clone()),
                _ => None,
            };

            self.notify(DagEvent::TaskCompleted(task_id.clone(), outcome));

            if is_success {
                report.completed.insert(task_id.clone());
                self.mark_completed(&task_id, &mut pending_deps, &mut ready);

                // Re-check deferred tasks (pool slot freed)
                let deferred_tasks: Vec<_> = deferred.drain(..).collect();
                for dt in deferred_tasks {
                    ready.push_back(dt);
                }
            } else if let Some(msg) = fail_msg {
                match &self.error_policy {
                    ErrorPolicy::FailFast => {
                        report.failed.insert(task_id.clone(), msg);
                        // Wait for in-flight tasks to finish, then return
                        drop(tx);
                        for event in rx {
                            in_flight.remove(&event.task_id);
                            report.cancelled.insert(event.task_id);
                        }
                        return Ok(report);
                    }
                    ErrorPolicy::ContinueIndependent => {
                        report.failed.insert(task_id.clone(), msg);
                        self.cancel_dependents(
                            &task_id,
                            &mut pending_deps,
                            &mut report.cancelled,
                            &mut ready,
                        );
                        let deferred_tasks: Vec<_> = deferred.drain(..).collect();
                        for dt in deferred_tasks {
                            if !report.cancelled.contains(&dt) {
                                ready.push_back(dt);
                            }
                        }
                    }
                    ErrorPolicy::Retry { .. } => {
                        report.failed.insert(task_id.clone(), msg);
                        self.cancel_dependents(
                            &task_id,
                            &mut pending_deps,
                            &mut report.cancelled,
                            &mut ready,
                        );
                    }
                }
            }
        }

        Ok(report)
    }

    /// Check cache for a task (per-task cache or global cache).
    fn check_cache(&self, task_id: &TaskId) -> Option<Box<dyn std::any::Any + Send>> {
        // Per-task cache takes priority
        if let Some(cache) = self.graph.task_cache.get(task_id) {
            if let Some(output) = cache.check(task_id, &self.context) {
                return Some(output);
            }
        }
        // Fall back to global cache
        if let Some(ref cache) = self.global_cache {
            return cache.check(task_id, &self.context);
        }
        None
    }

    /// Check if a task can be dispatched given its concurrency group limit.
    fn can_dispatch_group(&self, task_id: &TaskId, group_active: &HashMap<String, usize>) -> bool {
        if let Some(group_name) = self.graph.task_groups.get(task_id) {
            if let Some(limit) = self.graph.concurrency_groups.get(group_name) {
                if let Some(max) = limit.max_concurrent {
                    let active = group_active.get(group_name).copied().unwrap_or(0);
                    return active < max;
                }
            }
        }
        true
    }

    /// Mark a task as completed and update dependency counts.
    /// Newly-ready tasks are pushed into the ready queue.
    fn mark_completed(
        &self,
        task_id: &TaskId,
        pending_deps: &mut HashMap<TaskId, usize>,
        ready: &mut VecDeque<TaskId>,
    ) {
        if let Some(dependents) = self.graph.dependents.get(task_id) {
            for dep_id in dependents {
                if let Some(count) = pending_deps.get_mut(dep_id) {
                    *count = count.saturating_sub(1);
                    if *count == 0 {
                        ready.push_back(dep_id.clone());
                    }
                }
            }
        }
    }

    /// Cancel all transitive dependents of a failed task.
    fn cancel_dependents(
        &mut self,
        task_id: &TaskId,
        pending_deps: &mut HashMap<TaskId, usize>,
        cancelled: &mut HashSet<TaskId>,
        ready: &mut VecDeque<TaskId>,
    ) {
        let mut to_cancel = VecDeque::from([task_id.clone()]);
        let mut settled = Vec::new();
        while let Some(parent) = to_cancel.pop_front() {
            settled.push(parent.clone());
            let dependents: Vec<_> = self
                .graph
                .dependents
                .get(&parent)
                .into_iter()
                .flatten()
                .cloned()
                .collect();
            for dependent in dependents {
                if self
                    .graph
                    .completion_dependencies
                    .contains(&(dependent.clone(), parent.clone()))
                    || cancelled.contains(&dependent)
                {
                    continue;
                }
                cancelled.insert(dependent.clone());
                self.notify(DagEvent::TaskCancelled(dependent.clone()));
                pending_deps.remove(&dependent);
                to_cancel.push_back(dependent);
            }
        }
        for terminal in settled {
            self.mark_completed(&terminal, pending_deps, ready);
        }
    }

    fn notify(&mut self, event: DagEvent) {
        if let Some(ref mut monitor) = self.monitor {
            monitor.on_event(&event);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::dag::builder::{DagBuilder, StageBuilder};
    use crate::dag::pool::FixedPool;
    use crate::dag::types::{FnTask, ShellTask, TaskOutcome};
    use std::sync::atomic::{AtomicUsize, Ordering};
    use std::sync::{Arc, Mutex};
    use std::time::{Duration, Instant};

    #[test]
    fn test_single_task_success() {
        let mut dag = DagBuilder::new();
        dag.add_task("a", FnTask::new("task a", |_| TaskOutcome::Success));
        let report = dag.build().unwrap().run().unwrap();
        assert!(report.is_success());
        assert!(report.completed.contains(&TaskId::from("a")));
    }

    #[test]
    fn test_single_task_failure() {
        let mut dag = DagBuilder::new();
        dag.add_task(
            "a",
            FnTask::new("task a", |_| TaskOutcome::Failed("boom".into())),
        );
        dag.error_policy(ErrorPolicy::FailFast);
        let report = dag.build().unwrap().run().unwrap();
        assert!(!report.is_success());
        assert!(report.failed.contains_key(&TaskId::from("a")));
    }

    #[test]
    fn test_dependency_ordering() {
        let order = Arc::new(Mutex::new(Vec::new()));
        let o1 = order.clone();
        let o2 = order.clone();
        let o3 = order.clone();

        let mut dag = DagBuilder::new();
        dag.add_task(
            "a",
            FnTask::new("a", move |_| {
                o1.lock().unwrap().push("a");
                TaskOutcome::Success
            }),
        );
        dag.add_task(
            "b",
            FnTask::new("b", move |_| {
                o2.lock().unwrap().push("b");
                TaskOutcome::Success
            }),
        );
        dag.add_task(
            "c",
            FnTask::new("c", move |_| {
                o3.lock().unwrap().push("c");
                TaskOutcome::Success
            }),
        );
        dag.add_dep("b", "a");
        dag.add_dep("c", "b");
        // Force sequential with fanout 1
        dag.pool(FixedPool::new(1));

        let report = dag.build().unwrap().run().unwrap();
        assert!(report.is_success());

        let executed = order.lock().unwrap();
        assert_eq!(*executed, vec!["a", "b", "c"]);
    }

    #[test]
    fn test_dependency_ordering_with_high_parallelism() {
        // Verify dependencies are respected even with unlimited parallelism.
        // Task B reads A's output — if B runs before A, it will fail.
        let mut dag = DagBuilder::new();
        dag.add_task(
            "a",
            FnTask::new("a", |ctx| {
                std::thread::sleep(Duration::from_millis(20));
                ctx.set_output(TaskId::from("a"), "from_a".to_string());
                TaskOutcome::Success
            }),
        );
        dag.add_task(
            "b",
            FnTask::new("b", |ctx| {
                // This MUST see A's output — if deps are broken, this fails
                let val: Option<String> = ctx.get_output(&TaskId::from("a"));
                match val {
                    Some(v) if v == "from_a" => {
                        ctx.set_output(TaskId::from("b"), "from_b".to_string());
                        TaskOutcome::Success
                    }
                    Some(v) => TaskOutcome::Failed(format!("unexpected value from a: {}", v)),
                    None => {
                        TaskOutcome::Failed("dependency a has no output — ran out of order!".into())
                    }
                }
            }),
        );
        dag.add_task(
            "c",
            FnTask::new("c", |ctx| {
                let val: Option<String> = ctx.get_output(&TaskId::from("b"));
                match val {
                    Some(v) if v == "from_b" => TaskOutcome::Success,
                    _ => {
                        TaskOutcome::Failed("dependency b has no output — ran out of order!".into())
                    }
                }
            }),
        );
        dag.add_dep("b", "a");
        dag.add_dep("c", "b");
        // High parallelism — deps must still be respected
        dag.pool(FixedPool::new(100));

        let report = dag.build().unwrap().run().unwrap();
        assert!(
            report.is_success(),
            "dependencies violated with high parallelism: failed={:?}",
            report.failed
        );
    }

    #[test]
    fn test_parallel_independent_tasks() {
        let counter = Arc::new(AtomicUsize::new(0));
        let start = Instant::now();

        let mut dag = DagBuilder::new();
        for i in 0..5 {
            let c = counter.clone();
            dag.add_task(
                format!("t{}", i),
                FnTask::new(format!("task {}", i), move |_| {
                    std::thread::sleep(Duration::from_millis(50));
                    c.fetch_add(1, Ordering::SeqCst);
                    TaskOutcome::Success
                }),
            );
        }

        let report = dag.build().unwrap().run().unwrap();
        let elapsed = start.elapsed();

        assert!(report.is_success());
        assert_eq!(counter.load(Ordering::SeqCst), 5);
        // Functional check: if tasks ran sequentially, it would take >=250ms (5 × 50ms).
        // With parallelism it should be ~50-100ms. Use a generous bound to avoid CI flakiness.
        assert!(
            elapsed < Duration::from_secs(2),
            "tasks took {:?}, likely not running in parallel (sequential would be >=250ms)",
            elapsed
        );
    }

    #[test]
    fn test_continue_independent_on_failure() {
        let mut dag = DagBuilder::new();
        dag.add_task("good", FnTask::new("good", |_| TaskOutcome::Success));
        dag.add_task(
            "bad",
            FnTask::new("bad", |_| TaskOutcome::Failed("fail".into())),
        );
        dag.add_task(
            "depends_on_bad",
            FnTask::new("after bad", |_| TaskOutcome::Success),
        );
        dag.add_dep("depends_on_bad", "bad");
        dag.error_policy(ErrorPolicy::ContinueIndependent);

        let report = dag.build().unwrap().run().unwrap();
        assert!(!report.is_success());
        assert!(report.completed.contains(&TaskId::from("good")));
        assert!(report.failed.contains_key(&TaskId::from("bad")));
        assert!(report.cancelled.contains(&TaskId::from("depends_on_bad")));
    }

    #[test]
    fn test_output_passing_between_tasks() {
        let mut dag = DagBuilder::new();
        dag.add_task(
            "producer",
            FnTask::new("produce", |ctx| {
                ctx.set_output(TaskId::from("producer"), "hello from producer".to_string());
                TaskOutcome::Success
            }),
        );
        dag.add_task(
            "consumer",
            FnTask::new("consume", |ctx| {
                let val: String = ctx.get_output(&TaskId::from("producer")).unwrap();
                assert_eq!(val, "hello from producer");
                TaskOutcome::Success
            }),
        );
        dag.add_dep("consumer", "producer");

        let report = dag.build().unwrap().run().unwrap();
        assert!(report.is_success());
    }

    #[test]
    fn test_stage_builder_basic() {
        let report = StageBuilder::new()
            .resources(vec!["h1".into(), "h2".into(), "h3".into()])
            .stage("build", None, |host| {
                let h = host.to_string();
                Box::new(FnTask::new(format!("build {}", h), move |ctx| {
                    ctx.set_output(TaskId(format!("build:{}", h)), format!("built-{}", h));
                    TaskOutcome::Success
                }))
            })
            .stage("deploy", None, |host| {
                let h = host.to_string();
                Box::new(FnTask::new(format!("deploy {}", h), move |ctx| {
                    // Should be able to read build output
                    let built: String = ctx.get_output(&TaskId(format!("build:{}", h))).unwrap();
                    assert!(built.starts_with("built-"));
                    TaskOutcome::Success
                }))
            })
            .build()
            .unwrap()
            .run()
            .unwrap();

        assert!(report.is_success());
        assert_eq!(report.completed.len(), 6); // 3 hosts × 2 stages
    }

    #[test]
    fn test_stage_builder_per_host_pipelining() {
        // Verify that host1 can be in "deploy" while host2 is still in "build"
        let start = Instant::now();

        let report = StageBuilder::new()
            .resources(vec!["fast".into(), "slow".into()])
            .stage("build", None, |host| {
                let h = host.to_string();
                Box::new(FnTask::new(format!("build {}", h), move |_| {
                    let delay = if h == "slow" { 100 } else { 10 };
                    std::thread::sleep(Duration::from_millis(delay));
                    TaskOutcome::Success
                }))
            })
            .stage("deploy", None, |host| {
                let h = host.to_string();
                Box::new(FnTask::new(format!("deploy {}", h), move |_| {
                    std::thread::sleep(Duration::from_millis(10));
                    TaskOutcome::Success
                }))
            })
            .build()
            .unwrap()
            .run()
            .unwrap();

        let elapsed = start.elapsed();
        assert!(report.is_success());
        assert_eq!(report.completed.len(), 4); // 2 hosts × 2 stages
                                               // Pipelining means fast.deploy starts while slow is still building.
                                               // Without pipelining (global barrier), total = build_all(100ms) + deploy_all(10ms) = 110ms+
                                               // With pipelining, total ≈ max(slow_build + slow_deploy, fast_build + fast_deploy) = 110ms
                                               // Use generous bound to avoid CI flakiness.
        assert!(
            elapsed < Duration::from_secs(2),
            "pipelining took {:?}, expected well under 2s",
            elapsed
        );
    }

    #[test]
    fn test_stage_builder_failure_cancels_later_stages() {
        let report = StageBuilder::new()
            .resources(vec!["good".into(), "bad".into()])
            .stage("build", None, |host| {
                let h = host.to_string();
                Box::new(FnTask::new(format!("build {}", h), move |_| {
                    if h == "bad" {
                        TaskOutcome::Failed("build failed".into())
                    } else {
                        TaskOutcome::Success
                    }
                }))
            })
            .stage("deploy", None, |host| {
                let h = host.to_string();
                Box::new(FnTask::new(format!("deploy {}", h), move |_| {
                    TaskOutcome::Success
                }))
            })
            .error_policy(ErrorPolicy::ContinueIndependent)
            .build()
            .unwrap()
            .run()
            .unwrap();

        assert!(!report.is_success());
        // good: build + deploy both succeed
        assert!(report.completed.contains(&TaskId::from("build:good")));
        assert!(report.completed.contains(&TaskId::from("deploy:good")));
        // bad: build fails, deploy cancelled
        assert!(report.failed.contains_key(&TaskId::from("build:bad")));
        assert!(report.cancelled.contains(&TaskId::from("deploy:bad")));
    }

    #[test]
    fn test_concurrency_group_limit() {
        // Use context to pass the counters through, avoiding 'static closure issues
        let ctx = DagContext::new();
        let max_concurrent = Arc::new(AtomicUsize::new(0));
        let current = Arc::new(AtomicUsize::new(0));
        ctx.set_state("max_concurrent", max_concurrent.clone());
        ctx.set_state("current", current.clone());

        let report = StageBuilder::new()
            .resources((0..10).map(|i| format!("n{}", i)).collect())
            .stage("work", Some(2), |host| {
                let h = host.to_string();
                Box::new(FnTask::new(format!("work {}", h), |ctx| {
                    let cur: Arc<AtomicUsize> = ctx.get_state("current").unwrap();
                    let mc: Arc<AtomicUsize> = ctx.get_state("max_concurrent").unwrap();
                    let c = cur.fetch_add(1, Ordering::SeqCst) + 1;
                    loop {
                        let prev = mc.load(Ordering::SeqCst);
                        if c <= prev
                            || mc
                                .compare_exchange(prev, c, Ordering::SeqCst, Ordering::SeqCst)
                                .is_ok()
                        {
                            break;
                        }
                    }
                    std::thread::sleep(Duration::from_millis(30));
                    cur.fetch_sub(1, Ordering::SeqCst);
                    TaskOutcome::Success
                }))
            })
            .context(ctx)
            .build()
            .unwrap()
            .run()
            .unwrap();

        assert!(report.is_success());
        assert_eq!(report.completed.len(), 10);
        assert!(
            max_concurrent.load(Ordering::SeqCst) <= 2,
            "max concurrent was {}, expected <= 2",
            max_concurrent.load(Ordering::SeqCst)
        );
    }

    // ─── Functional proof tests ────────────────────────────────────────
    // These tests prove that parallelism, pipelining, and concurrency
    // limits are working via observable state (not timing assertions).

    #[test]
    fn test_parallelism_proven_by_overlapping_execution() {
        // PROOF: If tasks run in parallel, their execution windows overlap.
        // Each task records its (start_time, end_time). If any two tasks
        // have overlapping windows, parallelism is proven.
        use std::time::SystemTime;

        let ctx = DagContext::new();
        let timestamps: Arc<Mutex<Vec<(String, u128, u128)>>> = Arc::new(Mutex::new(Vec::new()));
        ctx.set_state("timestamps", timestamps.clone());

        let mut dag = DagBuilder::new();
        for i in 0..5 {
            dag.add_task(
                format!("t{}", i),
                FnTask::new(format!("task {}", i), move |ctx| {
                    let ts: Arc<Mutex<Vec<(String, u128, u128)>>> =
                        ctx.get_state("timestamps").unwrap();
                    let start = SystemTime::now()
                        .duration_since(SystemTime::UNIX_EPOCH)
                        .unwrap()
                        .as_millis();
                    std::thread::sleep(Duration::from_millis(50));
                    let end = SystemTime::now()
                        .duration_since(SystemTime::UNIX_EPOCH)
                        .unwrap()
                        .as_millis();
                    ts.lock().unwrap().push((format!("t{}", i), start, end));
                    TaskOutcome::Success
                }),
            );
        }
        // No dependencies — all should run in parallel
        dag.context(ctx);

        let report = dag.build().unwrap().run().unwrap();
        assert!(report.is_success());

        let ts = timestamps.lock().unwrap();
        assert_eq!(ts.len(), 5);

        // Check for overlapping execution windows
        let mut found_overlap = false;
        for i in 0..ts.len() {
            for j in (i + 1)..ts.len() {
                let (_, s1, e1) = &ts[i];
                let (_, s2, e2) = &ts[j];
                // Overlap: one starts before the other ends
                if s1 < e2 && s2 < e1 {
                    found_overlap = true;
                    break;
                }
            }
            if found_overlap {
                break;
            }
        }
        assert!(
            found_overlap,
            "no overlapping execution windows found — tasks may not be running in parallel. timestamps: {:?}",
            *ts
        );
    }

    #[test]
    fn test_pipelining_proven_by_stage_overlap() {
        // PROOF: With per-host pipelining, fast_host's stage 2 should START
        // before slow_host's stage 1 ENDS. We record timestamps per task
        // and check that deploy:fast starts before build:slow ends.
        use std::time::SystemTime;

        let ctx = DagContext::new();
        let timestamps: Arc<Mutex<HashMap<String, (u128, u128)>>> =
            Arc::new(Mutex::new(HashMap::new()));
        ctx.set_state("timestamps", timestamps.clone());

        let report = StageBuilder::new()
            .resources(vec!["fast".into(), "slow".into()])
            .stage("build", None, |host| {
                let h = host.to_string();
                Box::new(FnTask::new(format!("build {}", h), move |ctx| {
                    let ts: Arc<Mutex<HashMap<String, (u128, u128)>>> =
                        ctx.get_state("timestamps").unwrap();
                    let start = SystemTime::now()
                        .duration_since(SystemTime::UNIX_EPOCH)
                        .unwrap()
                        .as_millis();
                    let delay = if h == "slow" { 200 } else { 10 };
                    std::thread::sleep(Duration::from_millis(delay));
                    let end = SystemTime::now()
                        .duration_since(SystemTime::UNIX_EPOCH)
                        .unwrap()
                        .as_millis();
                    ts.lock()
                        .unwrap()
                        .insert(format!("build:{}", h), (start, end));
                    TaskOutcome::Success
                }))
            })
            .stage("deploy", None, |host| {
                let h = host.to_string();
                Box::new(FnTask::new(format!("deploy {}", h), move |ctx| {
                    let ts: Arc<Mutex<HashMap<String, (u128, u128)>>> =
                        ctx.get_state("timestamps").unwrap();
                    let start = SystemTime::now()
                        .duration_since(SystemTime::UNIX_EPOCH)
                        .unwrap()
                        .as_millis();
                    std::thread::sleep(Duration::from_millis(10));
                    let end = SystemTime::now()
                        .duration_since(SystemTime::UNIX_EPOCH)
                        .unwrap()
                        .as_millis();
                    ts.lock()
                        .unwrap()
                        .insert(format!("deploy:{}", h), (start, end));
                    TaskOutcome::Success
                }))
            })
            .context(ctx)
            .build()
            .unwrap()
            .run()
            .unwrap();

        assert!(report.is_success());

        let ts = timestamps.lock().unwrap();
        let build_slow = ts.get("build:slow").expect("build:slow missing");
        let deploy_fast = ts.get("deploy:fast").expect("deploy:fast missing");

        // Pipelining proof: deploy:fast should START before build:slow ENDS
        assert!(
            deploy_fast.0 < build_slow.1,
            "pipelining not working: deploy:fast started at {} but build:slow ended at {} \
             (deploy should start BEFORE build:slow finishes)",
            deploy_fast.0,
            build_slow.1
        );
    }

    #[test]
    fn test_concurrency_limit_proven_by_max_concurrent() {
        // PROOF: With concurrency limit=2 and 10 tasks, the maximum number
        // of simultaneously executing tasks must never exceed 2.
        // We track this with an atomic counter that increments on start
        // and decrements on end, recording the peak.
        let ctx = DagContext::new();
        let peak = Arc::new(AtomicUsize::new(0));
        let active = Arc::new(AtomicUsize::new(0));
        ctx.set_state("peak", peak.clone());
        ctx.set_state("active", active.clone());

        let report = StageBuilder::new()
            .resources((0..10).map(|i| format!("n{}", i)).collect())
            .stage("work", Some(2), |host| {
                let h = host.to_string();
                Box::new(FnTask::new(format!("work {}", h), |ctx| {
                    let active: Arc<AtomicUsize> = ctx.get_state("active").unwrap();
                    let peak: Arc<AtomicUsize> = ctx.get_state("peak").unwrap();

                    let current = active.fetch_add(1, Ordering::SeqCst) + 1;
                    // Update peak
                    peak.fetch_max(current, Ordering::SeqCst);

                    // Hold the slot for a bit so concurrency is observable
                    std::thread::sleep(Duration::from_millis(30));

                    active.fetch_sub(1, Ordering::SeqCst);
                    TaskOutcome::Success
                }))
            })
            .context(ctx)
            .build()
            .unwrap()
            .run()
            .unwrap();

        assert!(report.is_success());
        assert_eq!(report.completed.len(), 10);

        let observed_peak = peak.load(Ordering::SeqCst);
        assert!(
            observed_peak <= 2,
            "concurrency limit violated: peak concurrent tasks was {}, limit is 2",
            observed_peak
        );
        // Also prove tasks DID run concurrently (peak > 1)
        assert!(
            observed_peak >= 2,
            "tasks never ran concurrently: peak was {} (expected 2 with limit=2 and 10 tasks)",
            observed_peak
        );
    }

    #[test]
    fn test_concurrency_limit_different_groups() {
        // PROOF: Two concurrency groups with different limits operate independently.
        // Group "fast" has limit=4, group "slow" has limit=1.
        // Fast tasks should have peak=4, slow tasks should have peak=1.
        let ctx = DagContext::new();
        let fast_peak = Arc::new(AtomicUsize::new(0));
        let fast_active = Arc::new(AtomicUsize::new(0));
        let slow_peak = Arc::new(AtomicUsize::new(0));
        let slow_active = Arc::new(AtomicUsize::new(0));
        ctx.set_state("fast_peak", fast_peak.clone());
        ctx.set_state("fast_active", fast_active.clone());
        ctx.set_state("slow_peak", slow_peak.clone());
        ctx.set_state("slow_active", slow_active.clone());

        let mut dag = DagBuilder::new();
        dag.concurrency_group("fast", 4);
        dag.concurrency_group("slow", 1);

        for i in 0..8 {
            let id = format!("fast-{}", i);
            dag.add_task(
                id.clone(),
                FnTask::new(format!("fast {}", i), |ctx| {
                    let active: Arc<AtomicUsize> = ctx.get_state("fast_active").unwrap();
                    let peak: Arc<AtomicUsize> = ctx.get_state("fast_peak").unwrap();
                    let c = active.fetch_add(1, Ordering::SeqCst) + 1;
                    peak.fetch_max(c, Ordering::SeqCst);
                    std::thread::sleep(Duration::from_millis(30));
                    active.fetch_sub(1, Ordering::SeqCst);
                    TaskOutcome::Success
                }),
            );
            dag.assign_group(id, "fast");
        }

        for i in 0..4 {
            let id = format!("slow-{}", i);
            dag.add_task(
                id.clone(),
                FnTask::new(format!("slow {}", i), |ctx| {
                    let active: Arc<AtomicUsize> = ctx.get_state("slow_active").unwrap();
                    let peak: Arc<AtomicUsize> = ctx.get_state("slow_peak").unwrap();
                    let c = active.fetch_add(1, Ordering::SeqCst) + 1;
                    peak.fetch_max(c, Ordering::SeqCst);
                    std::thread::sleep(Duration::from_millis(20));
                    active.fetch_sub(1, Ordering::SeqCst);
                    TaskOutcome::Success
                }),
            );
            dag.assign_group(id, "slow");
        }

        dag.context(ctx);
        let report = dag.build().unwrap().run().unwrap();

        assert!(report.is_success());
        assert_eq!(report.completed.len(), 12);

        let fp = fast_peak.load(Ordering::SeqCst);
        let sp = slow_peak.load(Ordering::SeqCst);

        assert!(fp <= 4, "fast group exceeded limit: peak={}, limit=4", fp);
        assert!(sp <= 1, "slow group exceeded limit: peak={}, limit=1", sp);
        // Prove fast group actually used concurrency
        assert!(
            fp >= 2,
            "fast group didn't use concurrency: peak={}, expected >=2",
            fp
        );
    }

    #[test]
    fn test_dependencies_prevent_early_execution() {
        // PROOF: A dependent task NEVER executes before its dependency,
        // even under high parallelism. We verify this by having the
        // dependent task check a flag that the dependency sets.
        // If the dependent runs first, the flag won't be set and the task fails.
        let ctx = DagContext::new();
        let flag = Arc::new(AtomicUsize::new(0));
        ctx.set_state("flag", flag.clone());

        let mut dag = DagBuilder::new();

        // 10 independent "setup" tasks that all set flags
        for i in 0..10 {
            dag.add_task(
                format!("setup-{}", i),
                FnTask::new(format!("setup {}", i), move |ctx| {
                    let f: Arc<AtomicUsize> = ctx.get_state("flag").unwrap();
                    std::thread::sleep(Duration::from_millis(10));
                    f.fetch_add(1, Ordering::SeqCst);
                    TaskOutcome::Success
                }),
            );
        }

        // "check" task depends on ALL setup tasks
        dag.add_task(
            "check",
            FnTask::new("check all flags", move |ctx| {
                let f: Arc<AtomicUsize> = ctx.get_state("flag").unwrap();
                let count = f.load(Ordering::SeqCst);
                if count == 10 {
                    TaskOutcome::Success
                } else {
                    TaskOutcome::Failed(format!(
                        "only {}/10 setup tasks completed before check ran!",
                        count
                    ))
                }
            }),
        );

        for i in 0..10 {
            dag.add_dep("check", format!("setup-{}", i));
        }

        // High parallelism — check must still wait for all 10
        dag.pool(FixedPool::new(100));
        dag.context(ctx);

        let report = dag.build().unwrap().run().unwrap();
        assert!(
            report.is_success(),
            "dependency violation: {:?}",
            report.failed
        );
    }

    #[test]
    fn test_shell_task() {
        let mut dag = DagBuilder::new();
        dag.add_task(
            "echo",
            ShellTask {
                command: "echo hello".into(),
                description: "echo test".into(),
                resource: None,
            },
        );
        let report = dag.build().unwrap().run().unwrap();
        assert!(report.is_success());
    }

    #[test]
    fn test_shell_task_failure() {
        let mut dag = DagBuilder::new();
        dag.add_task(
            "fail",
            ShellTask {
                command: "exit 42".into(),
                description: "fail test".into(),
                resource: None,
            },
        );
        dag.error_policy(ErrorPolicy::FailFast);
        let report = dag.build().unwrap().run().unwrap();
        assert!(!report.is_success());
        assert!(report
            .failed
            .get(&TaskId::from("fail"))
            .unwrap()
            .contains("42"));
    }
    #[test]
    fn completion_barrier_waits_for_success_failed_and_cancelled_parents() {
        for _ in 0..24 {
            let healthy_finished = Arc::new(std::sync::atomic::AtomicBool::new(false));
            let mut dag = DagBuilder::new();
            dag.add_task(
                "failed",
                FnTask::new("fail", |_| TaskOutcome::Failed("injected".into())),
            );
            let finished = healthy_finished.clone();
            dag.add_task(
                "healthy",
                FnTask::new("healthy", move |_| {
                    finished.store(true, Ordering::SeqCst);
                    TaskOutcome::Success
                }),
            );
            for id in ["cancelled", "cancelled-child"] {
                dag.add_task(
                    id,
                    FnTask::new("must not run", |_| {
                        TaskOutcome::Failed("cancelled task dispatched".into())
                    }),
                );
            }
            dag.add_dep("cancelled", "failed");
            dag.add_dep("cancelled-child", "cancelled");
            let finished = healthy_finished.clone();
            dag.add_task(
                "settled",
                FnTask::new("barrier", move |_| {
                    if finished.load(Ordering::SeqCst) {
                        TaskOutcome::Success
                    } else {
                        TaskOutcome::Failed("barrier ran before healthy predecessor".into())
                    }
                }),
            );
            for id in ["failed", "healthy", "cancelled", "cancelled-child"] {
                dag.add_completion_dep("settled", id);
            }
            dag.add_task(
                "after",
                FnTask::new("after barrier", |_| TaskOutcome::Success),
            );
            dag.add_dep("after", "settled");
            dag.error_policy(ErrorPolicy::ContinueIndependent);
            let report = dag.build().unwrap().run().unwrap();
            assert_eq!(report.completed.len(), 3);
            assert!(report.completed.contains(&"settled".into()));
            assert!(report.completed.contains(&"after".into()));
            assert_eq!(report.failed.len(), 1);
            assert_eq!(report.cancelled.len(), 2);
        }
    }

    #[test]
    fn completion_dependency_does_not_override_another_required_success() {
        let mut dag = DagBuilder::new();
        for id in ["failure-one", "failure-two"] {
            dag.add_task(
                id,
                FnTask::new("fail", |_| TaskOutcome::Failed("injected".into())),
            );
        }
        dag.add_task(
            "mixed",
            FnTask::new("must not run", |_| {
                TaskOutcome::Failed("success dependency failed".into())
            }),
        );
        dag.add_completion_dep("mixed", "failure-one");
        dag.add_dep("mixed", "failure-two");
        dag.error_policy(ErrorPolicy::ContinueIndependent);
        let report = dag.build().unwrap().run().unwrap();
        assert!(report.cancelled.contains(&"mixed".into()));
        assert_eq!(report.failed.len(), 2);
    }
}