grafos_batch/

lib.rs

//! grafos-batch — Task graph executor for DAG-structured batch jobs.
//!
//! This crate provides a mini-Spark/Airflow-style batch job executor built on
//! fabricBIOS. Tasks are organized into a directed acyclic graph (DAG) with
//! explicit data dependencies. The executor computes a topological sort,
//! groups independent tasks into execution waves, and runs them with automatic
//! retry and failure propagation.
//!
//! # Quick start
//!
//! ```rust
//! use grafos_batch::*;
//!
//! let mut graph = TaskGraph::new();
//! let a = graph.add_task(TaskDef {
//!     name: "produce".into(),
//!     task_fn: Box::new(|_ctx| {
//!         let mut data = std::collections::HashMap::new();
//!         data.insert("out".into(), b"hello".to_vec());
//!         Ok(TaskOutput { data })
//!     }),
//!     resource_req: ResourceReq::default(),
//!     inputs: vec![],
//!     outputs: vec![DataRef { name: "out".into(), format: DataFormat::RawBytes }],
//!     retries: 0,
//! });
//! let b = graph.add_task(TaskDef {
//!     name: "consume".into(),
//!     task_fn: Box::new(|ctx| {
//!         assert_eq!(ctx.inputs.get("out").unwrap(), b"hello");
//!         Ok(TaskOutput { data: std::collections::HashMap::new() })
//!     }),
//!     resource_req: ResourceReq::default(),
//!     inputs: vec![DataRef { name: "out".into(), format: DataFormat::RawBytes }],
//!     outputs: vec![],
//!     retries: 0,
//! });
//! graph.add_dependency(a, b).unwrap();
//! let plan = graph.build().unwrap();
//! assert_eq!(plan.waves().len(), 2);
//! ```

#![cfg_attr(not(feature = "std"), no_std)]

extern crate alloc;

use alloc::boxed::Box;
use alloc::collections::BTreeMap;
use alloc::collections::BTreeSet;
use alloc::string::String;
use alloc::vec::Vec;
use core::fmt;

#[cfg(feature = "std")]
use std::collections::HashMap;

#[cfg(not(feature = "std"))]
use alloc::collections::BTreeMap as HashMap;

use grafos_std::error::FabricError;

// ── TaskId ──────────────────────────────────────────────────────────────

/// Unique identifier for a task within a [`TaskGraph`].
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct TaskId(pub u64);

impl fmt::Display for TaskId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "Task({})", self.0)
    }
}

// ── DataRef / DataFormat ────────────────────────────────────────────────

/// Format of data exchanged between tasks.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum DataFormat {
    /// Structured data serialized via postcard.
    Postcard,
    /// Opaque byte blob.
    RawBytes,
}

/// Reference to a named data artifact produced or consumed by a task.
#[derive(Debug, Clone)]
pub struct DataRef {
    /// Name used to match outputs to inputs across tasks.
    pub name: String,
    /// Serialization format.
    pub format: DataFormat,
}

// ── ResourceReq ─────────────────────────────────────────────────────────

/// Resource requirements for a task.
#[derive(Debug, Clone)]
pub struct ResourceReq {
    /// Minimum memory in bytes.
    pub min_memory: u64,
    /// Minimum CPU cores.
    pub min_cpu_cores: u32,
    /// Minimum block storage in bytes.
    pub min_block: u64,
    /// Fuel budget for execution.
    pub fuel: u64,
}

impl Default for ResourceReq {
    fn default() -> Self {
        ResourceReq {
            min_memory: 0,
            min_cpu_cores: 1,
            min_block: 0,
            fuel: 1_000_000,
        }
    }
}

// ── TaskContext / TaskOutput ─────────────────────────────────────────────

/// Execution context provided to a task function.
pub struct TaskContext {
    /// The ID of the currently executing task.
    pub task_id: TaskId,
    /// Input data keyed by [`DataRef::name`].
    pub inputs: HashMap<String, Vec<u8>>,
}

/// Output produced by a task function.
pub struct TaskOutput {
    /// Output data keyed by [`DataRef::name`].
    pub data: HashMap<String, Vec<u8>>,
}

impl fmt::Debug for TaskOutput {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("TaskOutput")
            .field("data_keys", &self.data.keys().collect::<Vec<_>>())
            .finish()
    }
}

// ── TaskDef ─────────────────────────────────────────────────────────────

/// Definition of a task within the graph.
///
/// The `task_fn` closure uses `FnMut` so it can be retried on failure.
/// For tasks with `retries > 0`, the closure is called again with the
/// same [`TaskContext`] on each retry attempt.
#[allow(clippy::type_complexity)]
pub struct TaskDef {
    /// Human-readable task name.
    pub name: String,
    /// The function to execute. Called with a [`TaskContext`] on each
    /// attempt (initial + retries).
    pub task_fn: Box<dyn FnMut(&TaskContext) -> Result<TaskOutput, FabricError>>,
    /// Resource requirements.
    pub resource_req: ResourceReq,
    /// Named data inputs this task consumes.
    pub inputs: Vec<DataRef>,
    /// Named data outputs this task produces.
    pub outputs: Vec<DataRef>,
    /// Number of retry attempts on failure (0 = no retries).
    pub retries: u32,
}

// ── TaskGraph ───────────────────────────────────────────────────────────

/// Builder for constructing a task DAG.
///
/// Tasks are added with [`add_task`](TaskGraph::add_task) and dependency
/// edges with [`add_dependency`](TaskGraph::add_dependency). Call
/// [`build`](TaskGraph::build) to validate the graph and compute an
/// [`ExecutionPlan`].
pub struct TaskGraph {
    tasks: BTreeMap<TaskId, TaskDef>,
    /// Forward edges: from → set of to (meaning "to" depends on "from").
    deps: BTreeMap<TaskId, BTreeSet<TaskId>>,
    /// Reverse edges: to → set of from (meaning "to" is blocked by "from").
    rev_deps: BTreeMap<TaskId, BTreeSet<TaskId>>,
    next_id: u64,
}

impl TaskGraph {
    /// Create an empty task graph.
    pub fn new() -> Self {
        TaskGraph {
            tasks: BTreeMap::new(),
            deps: BTreeMap::new(),
            rev_deps: BTreeMap::new(),
            next_id: 0,
        }
    }

    /// Add a task and return its assigned [`TaskId`].
    pub fn add_task(&mut self, def: TaskDef) -> TaskId {
        let id = TaskId(self.next_id);
        self.next_id += 1;
        self.tasks.insert(id, def);
        id
    }

    /// Declare that `to` depends on `from` (i.e. `from` must complete before
    /// `to` can start).
    ///
    /// # Errors
    ///
    /// Returns `FabricError::IoError(-100)` if either task ID does not
    /// exist in the graph.
    pub fn add_dependency(&mut self, from: TaskId, to: TaskId) -> Result<(), FabricError> {
        if !self.tasks.contains_key(&from) {
            return Err(FabricError::IoError(-100));
        }
        if !self.tasks.contains_key(&to) {
            return Err(FabricError::IoError(-100));
        }
        self.deps.entry(from).or_default().insert(to);
        self.rev_deps.entry(to).or_default().insert(from);
        Ok(())
    }

    /// Validate the graph: check for cycles and unresolved data references.
    ///
    /// # Errors
    ///
    /// Returns `FabricError::IoError(-101)` if a cycle is detected.
    /// Returns `FabricError::IoError(-102)` if a task has an input with no
    /// upstream producer.
    pub fn validate(&self) -> Result<(), FabricError> {
        self.detect_cycle()?;
        self.check_data_refs()?;
        Ok(())
    }

    /// Validate and compute an [`ExecutionPlan`] with topologically sorted
    /// waves. Consumes the graph.
    ///
    /// # Errors
    ///
    /// Returns the same errors as [`validate`](TaskGraph::validate).
    pub fn build(self) -> Result<ExecutionPlan, FabricError> {
        self.validate()?;
        let waves = self.compute_waves();
        Ok(ExecutionPlan {
            waves,
            tasks: self.tasks,
            rev_deps: self.rev_deps,
        })
    }

    fn detect_cycle(&self) -> Result<(), FabricError> {
        const WHITE: u8 = 0;

        let mut color: BTreeMap<TaskId, u8> = BTreeMap::new();
        for &id in self.tasks.keys() {
            color.insert(id, WHITE);
        }

        for &id in self.tasks.keys() {
            if color[&id] == WHITE && self.dfs_has_cycle(id, &mut color) {
                return Err(FabricError::IoError(-101));
            }
        }
        Ok(())
    }

    fn dfs_has_cycle(&self, node: TaskId, color: &mut BTreeMap<TaskId, u8>) -> bool {
        color.insert(node, 1); // GRAY
        if let Some(neighbors) = self.deps.get(&node) {
            for &next in neighbors {
                match color.get(&next) {
                    Some(1) => return true, // back edge = cycle
                    Some(0) => {
                        if self.dfs_has_cycle(next, color) {
                            return true;
                        }
                    }
                    _ => {}
                }
            }
        }
        color.insert(node, 2); // BLACK
        false
    }

    fn check_data_refs(&self) -> Result<(), FabricError> {
        let mut all_outputs: BTreeSet<&str> = BTreeSet::new();
        for def in self.tasks.values() {
            for out_ref in &def.outputs {
                all_outputs.insert(&out_ref.name);
            }
        }
        for def in self.tasks.values() {
            for in_ref in &def.inputs {
                if !all_outputs.contains(in_ref.name.as_str()) {
                    return Err(FabricError::IoError(-102));
                }
            }
        }
        Ok(())
    }

    fn compute_waves(&self) -> Vec<Vec<TaskId>> {
        // Kahn's algorithm producing level-grouped topological order.
        let mut in_degree: BTreeMap<TaskId, usize> = BTreeMap::new();
        for &id in self.tasks.keys() {
            in_degree.insert(id, 0);
        }
        for dependents in self.deps.values() {
            for &dep in dependents {
                *in_degree.entry(dep).or_insert(0) += 1;
            }
        }

        let mut waves = Vec::new();
        let mut ready: Vec<TaskId> = in_degree
            .iter()
            .filter(|(_, &deg)| deg == 0)
            .map(|(&id, _)| id)
            .collect();
        ready.sort();

        while !ready.is_empty() {
            waves.push(ready.clone());
            let mut next_ready = Vec::new();
            for &id in &ready {
                if let Some(dependents) = self.deps.get(&id) {
                    for &dep in dependents {
                        let deg = in_degree.get_mut(&dep).unwrap();
                        *deg -= 1;
                        if *deg == 0 {
                            next_ready.push(dep);
                        }
                    }
                }
            }
            next_ready.sort();
            ready = next_ready;
        }

        waves
    }
}

impl Default for TaskGraph {
    fn default() -> Self {
        Self::new()
    }
}

// ── ExecutionPlan ───────────────────────────────────────────────────────

/// A validated execution plan with tasks grouped into waves.
///
/// Wave N contains only tasks whose dependencies all appear in waves < N.
/// Tasks within the same wave are independent and could be executed in
/// parallel (though the current executor runs them sequentially).
pub struct ExecutionPlan {
    waves: Vec<Vec<TaskId>>,
    tasks: BTreeMap<TaskId, TaskDef>,
    rev_deps: BTreeMap<TaskId, BTreeSet<TaskId>>,
}

impl fmt::Debug for ExecutionPlan {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("ExecutionPlan")
            .field("waves", &self.waves)
            .field("task_count", &self.task_count())
            .finish()
    }
}

impl ExecutionPlan {
    /// Get the execution waves. Each wave is a list of independent task IDs.
    pub fn waves(&self) -> &[Vec<TaskId>] {
        &self.waves
    }

    /// Total number of tasks across all waves.
    pub fn task_count(&self) -> usize {
        self.waves.iter().map(|w| w.len()).sum()
    }
}

// ── TaskResult / TaskStatus / ExecutionResult ──────────────────────────

/// Status of a completed task.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TaskStatus {
    /// Task completed successfully.
    Success,
    /// Task failed after exhausting retries.
    Failed(String),
    /// Task was skipped because an upstream dependency failed.
    Skipped,
}

/// Result of executing a single task.
#[derive(Debug)]
pub struct TaskResult {
    /// Final status.
    pub status: TaskStatus,
    /// Number of retries consumed before success or final failure.
    pub retries_used: u32,
    /// Output data if the task succeeded.
    pub output: Option<TaskOutput>,
}

/// Aggregate result of executing an entire [`ExecutionPlan`].
pub struct ExecutionResult {
    /// Per-task results keyed by [`TaskId`].
    pub task_results: BTreeMap<TaskId, TaskResult>,
    /// Total tasks in the plan.
    pub total_tasks: u32,
    /// Tasks that completed successfully.
    pub succeeded: u32,
    /// Tasks that failed after exhausting retries.
    pub failed: u32,
    /// Tasks that were skipped due to upstream failures.
    pub skipped: u32,
}

// ── Executor ────────────────────────────────────────────────────────────

/// Single-threaded task graph executor.
///
/// Executes an [`ExecutionPlan`] wave by wave, propagating outputs from
/// completed tasks to downstream consumers via [`DataRef`] name matching.
/// Failed tasks are retried up to [`TaskDef::retries`] times; if retries
/// are exhausted, all transitive downstream tasks are marked
/// [`TaskStatus::Skipped`].
pub struct Executor;

impl Executor {
    /// Create a new executor.
    pub fn new() -> Self {
        Executor
    }

    /// Execute the plan, consuming it.
    ///
    /// Tasks are executed sequentially within each wave. Outputs from
    /// completed tasks are stored and made available as inputs to downstream
    /// tasks matched by [`DataRef::name`].
    pub fn run(mut plan: ExecutionPlan) -> Result<ExecutionResult, FabricError> {
        let total_tasks = plan.task_count() as u32;
        let mut results: BTreeMap<TaskId, TaskResult> = BTreeMap::new();
        let mut data_store: HashMap<String, Vec<u8>> = HashMap::new();
        let mut failed_ids: BTreeSet<TaskId> = BTreeSet::new();

        let mut succeeded: u32 = 0;
        let mut failed: u32 = 0;
        let mut skipped: u32 = 0;

        let waves: Vec<Vec<TaskId>> = plan.waves.clone();

        for wave in &waves {
            for &task_id in wave {
                // Check if any upstream dependency failed or was skipped.
                let should_skip = plan
                    .rev_deps
                    .get(&task_id)
                    .map(|upstream| upstream.iter().any(|u| failed_ids.contains(u)))
                    .unwrap_or(false);

                if should_skip {
                    failed_ids.insert(task_id);
                    results.insert(
                        task_id,
                        TaskResult {
                            status: TaskStatus::Skipped,
                            retries_used: 0,
                            output: None,
                        },
                    );
                    skipped += 1;
                    continue;
                }

                let mut def = plan.tasks.remove(&task_id).unwrap();

                // Gather inputs from data store.
                let mut inputs = HashMap::new();
                for in_ref in &def.inputs {
                    if let Some(data) = data_store.get(&in_ref.name) {
                        inputs.insert(in_ref.name.clone(), data.clone());
                    }
                }

                let max_attempts = def.retries + 1;
                let mut last_err = None;
                let mut retries_used = 0u32;

                for attempt in 0..max_attempts {
                    let ctx = TaskContext {
                        task_id,
                        inputs: inputs.clone(),
                    };

                    match (def.task_fn)(&ctx) {
                        Ok(output) => {
                            for out_ref in &def.outputs {
                                if let Some(data) = output.data.get(&out_ref.name) {
                                    data_store.insert(out_ref.name.clone(), data.clone());
                                }
                            }
                            retries_used = attempt;
                            results.insert(
                                task_id,
                                TaskResult {
                                    status: TaskStatus::Success,
                                    retries_used,
                                    output: Some(output),
                                },
                            );
                            succeeded += 1;
                            last_err = None;
                            break;
                        }
                        Err(e) => {
                            last_err = Some(e);
                            retries_used = attempt;
                        }
                    }
                }

                if let Some(e) = last_err {
                    failed_ids.insert(task_id);
                    results.insert(
                        task_id,
                        TaskResult {
                            status: TaskStatus::Failed(alloc::format!("{}", e)),
                            retries_used,
                            output: None,
                        },
                    );
                    failed += 1;
                }
            }
        }

        Ok(ExecutionResult {
            task_results: results,
            total_tasks,
            succeeded,
            failed,
            skipped,
        })
    }
}

impl Default for Executor {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn empty_graph() {
        let graph = TaskGraph::new();
        let plan = graph.build().unwrap();
        assert_eq!(plan.task_count(), 0);
        assert!(plan.waves().is_empty());
        let result = Executor::run(plan).unwrap();
        assert_eq!(result.total_tasks, 0);
        assert_eq!(result.succeeded, 0);
        assert_eq!(result.failed, 0);
        assert_eq!(result.skipped, 0);
    }

    #[test]
    fn single_task() {
        let mut graph = TaskGraph::new();
        graph.add_task(TaskDef {
            name: "solo".into(),
            task_fn: Box::new(|_ctx| {
                Ok(TaskOutput {
                    data: HashMap::new(),
                })
            }),
            resource_req: ResourceReq::default(),
            inputs: vec![],
            outputs: vec![],
            retries: 0,
        });

        let plan = graph.build().unwrap();
        assert_eq!(plan.task_count(), 1);
        assert_eq!(plan.waves().len(), 1);

        let result = Executor::run(plan).unwrap();
        assert_eq!(result.succeeded, 1);
    }

    #[test]
    fn linear_dag_a_b_c() {
        let mut graph = TaskGraph::new();
        let a = graph.add_task(TaskDef {
            name: "A".into(),
            task_fn: Box::new(|_ctx| {
                let mut data = HashMap::new();
                data.insert("a_out".into(), b"from_a".to_vec());
                Ok(TaskOutput { data })
            }),
            resource_req: ResourceReq::default(),
            inputs: vec![],
            outputs: vec![DataRef {
                name: "a_out".into(),
                format: DataFormat::RawBytes,
            }],
            retries: 0,
        });
        let b = graph.add_task(TaskDef {
            name: "B".into(),
            task_fn: Box::new(|ctx| {
                let a_data = ctx.inputs.get("a_out").unwrap();
                assert_eq!(a_data, b"from_a");
                let mut data = HashMap::new();
                data.insert("b_out".into(), b"from_b".to_vec());
                Ok(TaskOutput { data })
            }),
            resource_req: ResourceReq::default(),
            inputs: vec![DataRef {
                name: "a_out".into(),
                format: DataFormat::RawBytes,
            }],
            outputs: vec![DataRef {
                name: "b_out".into(),
                format: DataFormat::RawBytes,
            }],
            retries: 0,
        });
        let c = graph.add_task(TaskDef {
            name: "C".into(),
            task_fn: Box::new(|ctx| {
                let b_data = ctx.inputs.get("b_out").unwrap();
                assert_eq!(b_data, b"from_b");
                Ok(TaskOutput {
                    data: HashMap::new(),
                })
            }),
            resource_req: ResourceReq::default(),
            inputs: vec![DataRef {
                name: "b_out".into(),
                format: DataFormat::RawBytes,
            }],
            outputs: vec![],
            retries: 0,
        });

        graph.add_dependency(a, b).unwrap();
        graph.add_dependency(b, c).unwrap();

        let plan = graph.build().unwrap();
        assert_eq!(plan.waves().len(), 3);
        assert_eq!(plan.waves()[0], vec![a]);
        assert_eq!(plan.waves()[1], vec![b]);
        assert_eq!(plan.waves()[2], vec![c]);

        let result = Executor::run(plan).unwrap();
        assert_eq!(result.succeeded, 3);
        assert_eq!(result.failed, 0);
        assert_eq!(result.skipped, 0);
    }

    #[test]
    fn diamond_dag() {
        let mut graph = TaskGraph::new();
        let a = graph.add_task(TaskDef {
            name: "A".into(),
            task_fn: Box::new(|_ctx| {
                let mut data = HashMap::new();
                data.insert("shared".into(), b"from_a".to_vec());
                Ok(TaskOutput { data })
            }),
            resource_req: ResourceReq::default(),
            inputs: vec![],
            outputs: vec![DataRef {
                name: "shared".into(),
                format: DataFormat::RawBytes,
            }],
            retries: 0,
        });
        let b = graph.add_task(TaskDef {
            name: "B".into(),
            task_fn: Box::new(|ctx| {
                assert!(ctx.inputs.contains_key("shared"));
                let mut data = HashMap::new();
                data.insert("b_out".into(), b"from_b".to_vec());
                Ok(TaskOutput { data })
            }),
            resource_req: ResourceReq::default(),
            inputs: vec![DataRef {
                name: "shared".into(),
                format: DataFormat::RawBytes,
            }],
            outputs: vec![DataRef {
                name: "b_out".into(),
                format: DataFormat::RawBytes,
            }],
            retries: 0,
        });
        let c = graph.add_task(TaskDef {
            name: "C".into(),
            task_fn: Box::new(|ctx| {
                assert!(ctx.inputs.contains_key("shared"));
                let mut data = HashMap::new();
                data.insert("c_out".into(), b"from_c".to_vec());
                Ok(TaskOutput { data })
            }),
            resource_req: ResourceReq::default(),
            inputs: vec![DataRef {
                name: "shared".into(),
                format: DataFormat::RawBytes,
            }],
            outputs: vec![DataRef {
                name: "c_out".into(),
                format: DataFormat::RawBytes,
            }],
            retries: 0,
        });
        let d = graph.add_task(TaskDef {
            name: "D".into(),
            task_fn: Box::new(|ctx| {
                assert!(ctx.inputs.contains_key("b_out"));
                assert!(ctx.inputs.contains_key("c_out"));
                Ok(TaskOutput {
                    data: HashMap::new(),
                })
            }),
            resource_req: ResourceReq::default(),
            inputs: vec![
                DataRef {
                    name: "b_out".into(),
                    format: DataFormat::RawBytes,
                },
                DataRef {
                    name: "c_out".into(),
                    format: DataFormat::RawBytes,
                },
            ],
            outputs: vec![],
            retries: 0,
        });

        graph.add_dependency(a, b).unwrap();
        graph.add_dependency(a, c).unwrap();
        graph.add_dependency(b, d).unwrap();
        graph.add_dependency(c, d).unwrap();

        let plan = graph.build().unwrap();
        assert_eq!(plan.waves().len(), 3);
        assert_eq!(plan.waves()[0], vec![a]);
        assert!(plan.waves()[1].contains(&b));
        assert!(plan.waves()[1].contains(&c));
        assert_eq!(plan.waves()[2], vec![d]);

        let result = Executor::run(plan).unwrap();
        assert_eq!(result.succeeded, 4);
    }

    #[test]
    fn cycle_detection() {
        let mut graph = TaskGraph::new();
        let a = graph.add_task(TaskDef {
            name: "A".into(),
            task_fn: Box::new(|_| {
                Ok(TaskOutput {
                    data: HashMap::new(),
                })
            }),
            resource_req: ResourceReq::default(),
            inputs: vec![],
            outputs: vec![],
            retries: 0,
        });
        let b = graph.add_task(TaskDef {
            name: "B".into(),
            task_fn: Box::new(|_| {
                Ok(TaskOutput {
                    data: HashMap::new(),
                })
            }),
            resource_req: ResourceReq::default(),
            inputs: vec![],
            outputs: vec![],
            retries: 0,
        });
        let c = graph.add_task(TaskDef {
            name: "C".into(),
            task_fn: Box::new(|_| {
                Ok(TaskOutput {
                    data: HashMap::new(),
                })
            }),
            resource_req: ResourceReq::default(),
            inputs: vec![],
            outputs: vec![],
            retries: 0,
        });

        graph.add_dependency(a, b).unwrap();
        graph.add_dependency(b, c).unwrap();
        graph.add_dependency(c, a).unwrap();

        let err = graph.build().unwrap_err();
        assert_eq!(err, FabricError::IoError(-101));
    }

    #[test]
    fn task_failure_skips_downstream() {
        let mut graph = TaskGraph::new();
        let a = graph.add_task(TaskDef {
            name: "A-fails".into(),
            task_fn: Box::new(|_ctx| Err(FabricError::Disconnected)),
            resource_req: ResourceReq::default(),
            inputs: vec![],
            outputs: vec![DataRef {
                name: "a_out".into(),
                format: DataFormat::RawBytes,
            }],
            retries: 0,
        });
        let b = graph.add_task(TaskDef {
            name: "B-skipped".into(),
            task_fn: Box::new(|_ctx| panic!("should not be called")),
            resource_req: ResourceReq::default(),
            inputs: vec![DataRef {
                name: "a_out".into(),
                format: DataFormat::RawBytes,
            }],
            outputs: vec![],
            retries: 0,
        });

        graph.add_dependency(a, b).unwrap();

        let plan = graph.build().unwrap();
        let result = Executor::run(plan).unwrap();

        assert_eq!(result.failed, 1);
        assert_eq!(result.skipped, 1);
        assert_eq!(result.succeeded, 0);

        assert!(matches!(
            result.task_results.get(&a).unwrap().status,
            TaskStatus::Failed(_)
        ));
        assert_eq!(
            result.task_results.get(&b).unwrap().status,
            TaskStatus::Skipped
        );
    }

    #[test]
    fn task_retry_succeeds() {
        use std::cell::Cell;

        let counter = std::rc::Rc::new(Cell::new(0u32));
        let counter_clone = counter.clone();

        let mut graph = TaskGraph::new();
        graph.add_task(TaskDef {
            name: "flaky".into(),
            task_fn: Box::new(move |_ctx| {
                let attempt = counter_clone.get();
                counter_clone.set(attempt + 1);
                if attempt == 0 {
                    Err(FabricError::Disconnected)
                } else {
                    Ok(TaskOutput {
                        data: HashMap::new(),
                    })
                }
            }),
            resource_req: ResourceReq::default(),
            inputs: vec![],
            outputs: vec![],
            retries: 2,
        });

        let plan = graph.build().unwrap();
        let result = Executor::run(plan).unwrap();
        assert_eq!(result.succeeded, 1);
        assert_eq!(result.failed, 0);
        assert_eq!(counter.get(), 2); // called twice: fail then succeed
        let task_result = result.task_results.values().next().unwrap();
        assert_eq!(task_result.status, TaskStatus::Success);
        assert_eq!(task_result.retries_used, 1);
    }

    #[test]
    fn retries_exhausted_marks_failed_and_skips_downstream() {
        let mut graph = TaskGraph::new();
        let a = graph.add_task(TaskDef {
            name: "always-fails".into(),
            task_fn: Box::new(|_ctx| Err(FabricError::Disconnected)),
            resource_req: ResourceReq::default(),
            inputs: vec![],
            outputs: vec![],
            retries: 2,
        });
        let b = graph.add_task(TaskDef {
            name: "downstream".into(),
            task_fn: Box::new(|_ctx| panic!("should not run")),
            resource_req: ResourceReq::default(),
            inputs: vec![],
            outputs: vec![],
            retries: 0,
        });

        graph.add_dependency(a, b).unwrap();

        let plan = graph.build().unwrap();
        let result = Executor::run(plan).unwrap();
        assert_eq!(result.failed, 1);
        assert_eq!(result.skipped, 1);

        let a_result = result.task_results.get(&a).unwrap();
        assert!(matches!(a_result.status, TaskStatus::Failed(_)));
        assert_eq!(a_result.retries_used, 2); // 0, 1, 2 = 3 attempts, last index is 2
    }

    #[test]
    fn data_passing_raw_bytes() {
        let mut graph = TaskGraph::new();
        let a = graph.add_task(TaskDef {
            name: "producer".into(),
            task_fn: Box::new(|_ctx| {
                let mut data = HashMap::new();
                data.insert("payload".into(), vec![1, 2, 3, 4, 5]);
                Ok(TaskOutput { data })
            }),
            resource_req: ResourceReq::default(),
            inputs: vec![],
            outputs: vec![DataRef {
                name: "payload".into(),
                format: DataFormat::RawBytes,
            }],
            retries: 0,
        });
        let b = graph.add_task(TaskDef {
            name: "consumer".into(),
            task_fn: Box::new(|ctx| {
                let payload = ctx.inputs.get("payload").unwrap();
                assert_eq!(payload, &[1, 2, 3, 4, 5]);
                Ok(TaskOutput {
                    data: HashMap::new(),
                })
            }),
            resource_req: ResourceReq::default(),
            inputs: vec![DataRef {
                name: "payload".into(),
                format: DataFormat::RawBytes,
            }],
            outputs: vec![],
            retries: 0,
        });

        graph.add_dependency(a, b).unwrap();
        let plan = graph.build().unwrap();
        let result = Executor::run(plan).unwrap();
        assert_eq!(result.succeeded, 2);
    }

    #[test]
    fn unknown_task_id_in_dependency() {
        let mut graph = TaskGraph::new();
        let a = graph.add_task(TaskDef {
            name: "A".into(),
            task_fn: Box::new(|_| {
                Ok(TaskOutput {
                    data: HashMap::new(),
                })
            }),
            resource_req: ResourceReq::default(),
            inputs: vec![],
            outputs: vec![],
            retries: 0,
        });

        let err = graph.add_dependency(a, TaskId(999));
        assert_eq!(err.unwrap_err(), FabricError::IoError(-100));
    }

    #[test]
    fn unresolved_input_ref() {
        let mut graph = TaskGraph::new();
        graph.add_task(TaskDef {
            name: "needs-missing".into(),
            task_fn: Box::new(|_| {
                Ok(TaskOutput {
                    data: HashMap::new(),
                })
            }),
            resource_req: ResourceReq::default(),
            inputs: vec![DataRef {
                name: "nonexistent".into(),
                format: DataFormat::RawBytes,
            }],
            outputs: vec![],
            retries: 0,
        });

        let err = graph.build().unwrap_err();
        assert_eq!(err, FabricError::IoError(-102));
    }

    #[test]
    fn task_id_display() {
        let id = TaskId(42);
        assert_eq!(alloc::format!("{}", id), "Task(42)");
    }

    #[test]
    fn resource_req_default() {
        let req = ResourceReq::default();
        assert_eq!(req.min_cpu_cores, 1);
        assert_eq!(req.min_memory, 0);
        assert_eq!(req.min_block, 0);
        assert_eq!(req.fuel, 1_000_000);
    }

    #[test]
    fn wide_parallel_wave() {
        let mut graph = TaskGraph::new();
        for i in 0..5 {
            graph.add_task(TaskDef {
                name: alloc::format!("task-{}", i),
                task_fn: Box::new(|_| {
                    Ok(TaskOutput {
                        data: HashMap::new(),
                    })
                }),
                resource_req: ResourceReq::default(),
                inputs: vec![],
                outputs: vec![],
                retries: 0,
            });
        }

        let plan = graph.build().unwrap();
        assert_eq!(plan.waves().len(), 1);
        assert_eq!(plan.waves()[0].len(), 5);

        let result = Executor::run(plan).unwrap();
        assert_eq!(result.succeeded, 5);
    }

    #[test]
    fn transitive_skip_propagation() {
        let mut graph = TaskGraph::new();
        let a = graph.add_task(TaskDef {
            name: "A-fails".into(),
            task_fn: Box::new(|_| Err(FabricError::Disconnected)),
            resource_req: ResourceReq::default(),
            inputs: vec![],
            outputs: vec![],
            retries: 0,
        });
        let b = graph.add_task(TaskDef {
            name: "B".into(),
            task_fn: Box::new(|_| panic!("should not run")),
            resource_req: ResourceReq::default(),
            inputs: vec![],
            outputs: vec![],
            retries: 0,
        });
        let c = graph.add_task(TaskDef {
            name: "C".into(),
            task_fn: Box::new(|_| panic!("should not run")),
            resource_req: ResourceReq::default(),
            inputs: vec![],
            outputs: vec![],
            retries: 0,
        });

        graph.add_dependency(a, b).unwrap();
        graph.add_dependency(b, c).unwrap();

        let plan = graph.build().unwrap();
        let result = Executor::run(plan).unwrap();
        assert_eq!(result.failed, 1);
        assert_eq!(result.skipped, 2);
        assert_eq!(
            result.task_results.get(&a).unwrap().status,
            TaskStatus::Failed("disconnected".into())
        );
        assert_eq!(
            result.task_results.get(&b).unwrap().status,
            TaskStatus::Skipped
        );
        assert_eq!(
            result.task_results.get(&c).unwrap().status,
            TaskStatus::Skipped
        );
    }
}