dag-executor 0.1.0

//! The DAG engine: graph, scheduler, worker pool, and the executor that drives
//! them.
//!
//! [`DagExecutor`] is the entry point. It validates a [`Dag`], optionally
//! recovers persisted state, then runs tasks concurrently — respecting
//! dependencies, priorities, retries, the circuit breaker, and the dead-letter
//! queue — and returns an [`ExecutionReport`].

mod graph;
mod scheduler;
mod worker_pool;

pub use graph::Dag;
pub use scheduler::Scheduler;
pub use worker_pool::{TaskResult, WorkerPool};

use crate::advanced::{CircuitBreaker, DeadLetterQueue, RetryPolicy};
use crate::context::Context;
use crate::error::{DagExecutorError, Result, TaskError};
use crate::metrics::{MetricsCollector, MetricsSnapshot};
use crate::state::{StateValidator, TaskRecord, TaskState};
use crate::storage::{FileStorage, MemoryStorage, Storage};
use crate::utils::Config;
use futures::stream::{FuturesUnordered, StreamExt};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;

/// Summary of a completed (or cancelled) run.
#[derive(Debug, Clone)]
pub struct ExecutionReport {
    /// Unique id of the run.
    pub run_id: String,
    /// Final record for every task.
    pub records: HashMap<String, TaskRecord>,
    /// Metrics snapshot taken at the end of the run.
    pub metrics: MetricsSnapshot,
}

impl ExecutionReport {
    /// Whether every task completed successfully.
    pub fn is_success(&self) -> bool {
        !self.records.is_empty()
            && self
                .records
                .values()
                .all(|r| r.state == TaskState::Completed)
    }

    /// Ids of tasks that ended in a failure state.
    pub fn failed_tasks(&self) -> Vec<String> {
        self.records
            .values()
            .filter(|r| r.state.is_failure())
            .map(|r| r.id.clone())
            .collect()
    }

    /// Count of tasks in a given state.
    pub fn count_in(&self, state: TaskState) -> usize {
        self.records.values().filter(|r| r.state == state).count()
    }
}

/// Fluent builder for [`DagExecutor`].
pub struct DagExecutorBuilder {
    config: Config,
    storage: Option<Arc<dyn Storage>>,
    retry: Option<RetryPolicy>,
    breaker: Option<Arc<CircuitBreaker>>,
}

impl DagExecutorBuilder {
    fn new() -> Self {
        DagExecutorBuilder {
            config: Config::default(),
            storage: None,
            retry: None,
            breaker: None,
        }
    }

    /// Use a fully custom configuration.
    pub fn config(mut self, config: Config) -> Self {
        self.config = config;
        self
    }

    /// Set the maximum number of concurrently executing tasks.
    pub fn concurrency(mut self, n: usize) -> Self {
        self.config.max_concurrency = n;
        self
    }

    /// Provide a custom storage backend (defaults to [`FileStorage`] at the
    /// configured `storage_dir`, or [`MemoryStorage`] when persistence is off).
    pub fn storage(mut self, storage: Arc<dyn Storage>) -> Self {
        self.storage = Some(storage);
        self
    }

    /// Set the retry policy applied to every task.
    pub fn retry(mut self, retry: RetryPolicy) -> Self {
        self.retry = Some(retry);
        self
    }

    /// Enable a shared circuit breaker guarding task execution.
    pub fn circuit_breaker(mut self, breaker: CircuitBreaker) -> Self {
        self.breaker = Some(Arc::new(breaker));
        self
    }

    /// Toggle state persistence.
    pub fn persist(mut self, persist: bool) -> Self {
        self.config.persist_state = persist;
        self
    }

    /// Finalize the executor.
    pub fn build(self) -> DagExecutor {
        let config = Arc::new(self.config);
        let storage: Arc<dyn Storage> = self.storage.unwrap_or_else(|| {
            if config.persist_state {
                Arc::new(
                    FileStorage::open(&config.storage_dir)
                        .expect("failed to open storage directory"),
                )
            } else {
                Arc::new(MemoryStorage::new())
            }
        });
        let retry = self.retry.unwrap_or(RetryPolicy {
            max_attempts: config.max_attempts,
            ..RetryPolicy::default()
        });

        DagExecutor {
            metrics: Arc::new(MetricsCollector::new()),
            dead_letter: DeadLetterQueue::new(storage.clone()),
            validator: StateValidator::new(),
            breaker: self.breaker,
            timeout: config.task_timeout,
            retry,
            storage,
            config,
        }
    }
}

/// Executes DAGs of tasks with persistence, fault-tolerance and observability.
pub struct DagExecutor {
    config: Arc<Config>,
    storage: Arc<dyn Storage>,
    metrics: Arc<MetricsCollector>,
    dead_letter: DeadLetterQueue,
    validator: StateValidator,
    breaker: Option<Arc<CircuitBreaker>>,
    retry: RetryPolicy,
    timeout: Option<Duration>,
}

impl DagExecutor {
    /// Start building an executor.
    pub fn builder() -> DagExecutorBuilder {
        DagExecutorBuilder::new()
    }

    /// Build an executor with all defaults.
    pub fn new() -> Self {
        DagExecutorBuilder::new().build()
    }

    /// Access the metrics collector (live during a run).
    pub fn metrics(&self) -> &Arc<MetricsCollector> {
        &self.metrics
    }

    /// Access the dead-letter queue.
    pub fn dead_letter(&self) -> &DeadLetterQueue {
        &self.dead_letter
    }

    fn record_key(id: &str) -> String {
        format!("record:{id}")
    }

    /// Run `dag` to completion with a fresh context.
    pub async fn run(&self, dag: Dag) -> Result<ExecutionReport> {
        let ctx = Arc::new(Context::new(self.config.clone()));
        self.run_with_context(dag, ctx).await
    }

    /// Run `dag` using a caller-supplied [`Context`].
    ///
    /// This is the hook for graceful shutdown: hold a clone of `ctx` and call
    /// [`Context::cancel`] (e.g. on `SIGINT`) to stop scheduling new work.
    pub async fn run_with_context(&self, dag: Dag, ctx: Arc<Context>) -> Result<ExecutionReport> {
        dag.validate()?;

        let mut scheduler = self.recover_scheduler(&dag, &ctx).await?;

        // remaining[id] = number of dependencies not yet completed.
        let mut remaining: HashMap<String, usize> = HashMap::new();
        for id in dag.task_ids() {
            let pending_deps = dag
                .dependencies_of(&id)
                .into_iter()
                .filter(|d| scheduler.state(d) != Some(TaskState::Completed))
                .count();
            remaining.insert(id, pending_deps);
        }

        let pool = WorkerPool::new(
            self.config.max_concurrency,
            self.retry,
            self.timeout,
            self.metrics.clone(),
        );

        // Seed the schedule: cascade-skip anything blocked by an already-failed
        // dependency, then enqueue everything whose deps are all satisfied.
        let initial_failures: Vec<String> = dag
            .task_ids()
            .into_iter()
            .filter(|id| scheduler.state(id).map(|s| s.is_failure()).unwrap_or(false))
            .collect();
        for id in initial_failures {
            self.cascade_skip(&dag, &mut scheduler, &id).await?;
        }
        for id in dag.task_ids() {
            if scheduler.state(&id) == Some(TaskState::Pending)
                && remaining.get(&id).copied().unwrap_or(0) == 0
            {
                let prio = dag.task(&id).map(|t| t.priority()).unwrap_or(0);
                scheduler.mark_ready(&id, prio);
            }
        }

        let mut in_flight: FuturesUnordered<_> = FuturesUnordered::new();

        loop {
            // Launch every currently-ready task; the worker pool's semaphore
            // bounds how many actually run at once.
            while let Some(id) = scheduler.next_ready() {
                let task = match dag.task(&id) {
                    Some(t) => t,
                    None => continue,
                };
                scheduler.transition(&id, TaskState::Running);
                self.persist(&scheduler, &id).await?;
                in_flight.push(pool.spawn(task, ctx.clone(), self.breaker.clone()));
            }

            if in_flight.is_empty() {
                break;
            }

            let joined = match in_flight.next().await {
                Some(Ok(result)) => result,
                Some(Err(join_err)) => {
                    return Err(DagExecutorError::Executor(format!(
                        "worker task panicked: {join_err}"
                    )))
                }
                None => break,
            };

            self.handle_result(&dag, &mut scheduler, &mut remaining, &ctx, joined)
                .await?;
        }

        Ok(ExecutionReport {
            run_id: ctx.run_id.clone(),
            records: scheduler.records().clone(),
            metrics: self.metrics.snapshot(),
        })
    }

    /// Build a scheduler, recovering and repairing any persisted records.
    async fn recover_scheduler(&self, dag: &Dag, ctx: &Arc<Context>) -> Result<Scheduler> {
        let mut records: HashMap<String, TaskRecord> = HashMap::new();

        if self.config.persist_state {
            for id in dag.task_ids() {
                // A record that fails to load (e.g. a checksum mismatch from a
                // crash during a Fast in-place write) is treated as absent, so
                // the task simply re-runs rather than aborting recovery.
                let value = match self.storage.load(&Self::record_key(&id)).await {
                    Ok(v) => v,
                    Err(e) => {
                        tracing::warn!(task = %id, error = %e, "ignoring unreadable record during recovery");
                        None
                    }
                };
                if let Some(value) = value {
                    if let Ok(record) = serde_json::from_value::<TaskRecord>(value) {
                        // Re-publish recovered outputs so dependents can read them.
                        if record.state == TaskState::Completed {
                            if let Some(out) = &record.output {
                                ctx.set(record.id.clone(), out.clone());
                            }
                        }
                        records.insert(id, record);
                    }
                }
            }
            self.validator.repair(&mut records, self.retry.max_attempts);
        }

        let mut scheduler = Scheduler::with_records(records);
        // Make sure every task has a record (recovered runs may add new tasks).
        for id in dag.task_ids() {
            scheduler.ensure_record(&id);
        }
        Ok(scheduler)
    }

    /// Apply the outcome of a finished task and unblock/skip dependents.
    async fn handle_result(
        &self,
        dag: &Dag,
        scheduler: &mut Scheduler,
        remaining: &mut HashMap<String, usize>,
        ctx: &Arc<Context>,
        result: TaskResult,
    ) -> Result<()> {
        let TaskResult {
            id,
            attempts,
            outcome,
        } = result;

        if let Some(record) = scheduler.records_mut().get_mut(&id) {
            record.attempts = attempts;
        }

        match outcome {
            Ok(output) => {
                // Publish output for downstream consumers.
                ctx.set(id.clone(), output.clone());
                if let Some(record) = scheduler.records_mut().get_mut(&id) {
                    record.output = Some(output);
                    record.transition(TaskState::Completed);
                }
                let duration = scheduler
                    .record(&id)
                    .and_then(|r| r.duration_millis())
                    .unwrap_or(0);
                self.metrics.task_completed(&id, duration);
                self.persist(scheduler, &id).await?;

                // Unblock dependents whose last dependency just completed.
                for dep in dag.dependents_of(&id) {
                    let count = remaining.entry(dep.clone()).or_insert(0);
                    *count = count.saturating_sub(1);
                    if *count == 0 && scheduler.state(&dep) == Some(TaskState::Pending) {
                        let prio = dag.task(&dep).map(|t| t.priority()).unwrap_or(0);
                        scheduler.mark_ready(&dep, prio);
                    }
                }
            }
            Err(TaskError::Cancelled) => {
                if let Some(record) = scheduler.records_mut().get_mut(&id) {
                    record.transition(TaskState::Cancelled);
                }
                self.persist(scheduler, &id).await?;
                self.cascade_skip(dag, scheduler, &id).await?;
            }
            Err(err) => {
                let msg = err.to_string();
                if let Some(record) = scheduler.records_mut().get_mut(&id) {
                    record.error = Some(msg.clone());
                    record.transition(TaskState::Failed);
                }
                self.metrics.task_failed();

                // Retries are exhausted by the time we get here: dead-letter it.
                self.dead_letter.push(&id, attempts, msg).await?;
                if let Some(record) = scheduler.records_mut().get_mut(&id) {
                    record.transition(TaskState::DeadLettered);
                }
                self.metrics.task_dead_lettered();
                self.persist(scheduler, &id).await?;

                self.cascade_skip(dag, scheduler, &id).await?;
            }
        }
        Ok(())
    }

    /// Mark every (transitive) dependent of a failed/cancelled task as skipped.
    async fn cascade_skip(
        &self,
        dag: &Dag,
        scheduler: &mut Scheduler,
        failed_id: &str,
    ) -> Result<()> {
        let mut stack: Vec<String> = dag.dependents_of(failed_id);
        while let Some(id) = stack.pop() {
            if scheduler.state(&id) == Some(TaskState::Pending)
                && scheduler.transition(&id, TaskState::Skipped)
            {
                self.metrics.task_skipped();
                self.persist(scheduler, &id).await?;
                stack.extend(dag.dependents_of(&id));
            }
        }
        Ok(())
    }

    /// Persist a single task's record if persistence is enabled.
    async fn persist(&self, scheduler: &Scheduler, id: &str) -> Result<()> {
        if !self.config.persist_state {
            return Ok(());
        }
        if let Some(record) = scheduler.record(id) {
            let value = serde_json::to_value(record).map_err(crate::error::StorageError::from)?;
            self.storage.save(&Self::record_key(id), &value).await?;
        }
        Ok(())
    }
}

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