DevFlow/crates/df-workflow/src/state.rs

//! 节点状态机 — 管理节点的状态转换
//!
//! 合法转换：Pending → Running，Running → Completed / Failed，
//! 非法转换返回错误，避免状态被随意覆盖。
//!
//! 共享语义：内部用 `Arc<Mutex<HashMap>>`，`clone()` 为浅拷贝（Arc 引用计数 +1），
//! 所有 clone 共享同一底层 HashMap。DagExecutor 的 `state_machine` 与其下沉到
//! `NodeContext.node_status` 的 clone 共享底层；run_workflow 把执行器状态机注册到
//! AppState 全局表后，`cancel_workflow_node` IPC 的 `set_cancelled` 可直达运行中
//! 阻塞节点（如 HumanNode）的 `is_cancelled` 轮询。

use std::collections::HashMap;
use std::sync::{Arc, Mutex};

use df_types::types::{NodeId, NodeStatus};

/// 节点状态机
///
/// # ⚠️ 并发安全警告
///
/// `state_machine` 方法（`get`、`transition`、`set_running`、`set_completed`、
/// `set_failed`、`set_cancelled`、`snapshot`、`is_cancelled`）内部持锁
/// （`self.states.lock()`），**不得在 `.await` 期间持锁**。
/// 若在异步上下文中调用，请确保获取结果后立即释放锁（即不要将锁跨越
/// `.await` 点）。当前所有方法均为同步且及时释放，符合此规则。
///
/// 共享语义：内部用 `Arc<Mutex<HashMap>>`，`clone()` 为浅拷贝（Arc 引用计数 +1），
/// 所有 clone 共享同一底层 HashMap。
#[derive(Debug, Clone)]
pub struct StateMachine {
    /// 节点状态映射（Arc 共享，clone 浅拷贝同一 HashMap）
    states: Arc<Mutex<HashMap<NodeId, NodeStatus>>>,
}

impl StateMachine {
    /// 创建状态机
    pub fn new() -> Self {
        Self {
            states: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    /// 获取节点状态
    pub fn get(&self, node_id: &NodeId) -> NodeStatus {
        self.states
            .lock()
            .expect("状态机锁中毒")
            .get(node_id)
            .cloned()
            .unwrap_or(NodeStatus::Pending)
    }

    /// 判断状态转换是否合法
    fn is_legal(from: &NodeStatus, to: &NodeStatus) -> bool {
        matches!(
            (from, to),
            (NodeStatus::Pending, NodeStatus::Running)
                | (NodeStatus::Running, NodeStatus::Completed)
                | (NodeStatus::Running, NodeStatus::Failed)
        )
    }

    /// 状态转换 — 校验合法性后更新，非法转换返回错误
    pub fn transition(&self, node_id: NodeId, target: NodeStatus) -> anyhow::Result<()> {
        let mut states = self.states.lock().expect("状态机锁中毒");
        let current = states
            .get(&node_id)
            .cloned()
            .unwrap_or(NodeStatus::Pending);
        if !Self::is_legal(&current, &target) {
            anyhow::bail!(
                "节点 {} 状态转换非法：{} -> {}",
                node_id,
                current.as_str(),
                target.as_str()
            );
        }
        states.insert(node_id, target);
        Ok(())
    }

    /// 设置节点为运行中（仅允许 Pending -> Running）
    pub fn set_running(&self, node_id: NodeId) -> anyhow::Result<()> {
        self.transition(node_id, NodeStatus::Running)
    }

    /// 设置节点为已完成（仅允许 Running -> Completed）
    pub fn set_completed(&self, node_id: NodeId) -> anyhow::Result<()> {
        self.transition(node_id, NodeStatus::Completed)
    }

    /// 设置节点为失败（仅允许 Running -> Failed）
    pub fn set_failed(&self, node_id: NodeId) -> anyhow::Result<()> {
        self.transition(node_id, NodeStatus::Failed)
    }

    /// 设置节点为已取消 — 唯一受控旁路：不经 transition 合法性校验
    ///
    /// 设计理由：人工审批取消场景由前端 cancel_workflow_node IPC 异步触发，
    /// 此时节点可能处于 Running 之外的任意态（如 Pending 排队中、阻塞等待审批中），
    /// 若走 transition 校验会被 is_legal(Pending→Cancelled) 拒绝。
    /// 故 set_cancelled 作为唯一受控旁路直接置位 Cancelled，
    /// 阻塞节点(如 HumanNode)在 select! 轮询分支通过 is_cancelled 检测后返回 Err。
    ///
    /// 注：原 set_waiting/set_skipped 同为旁路置位但全仓零调用，已删除。
    pub fn set_cancelled(&self, node_id: NodeId) {
        self.states
            .lock()
            .expect("状态机锁中毒")
            .insert(node_id, NodeStatus::Cancelled);
    }

    /// 获取所有状态快照（clone 返回，调用方持独立副本）
    pub fn snapshot(&self) -> HashMap<NodeId, NodeStatus> {
        self.states.lock().expect("状态机锁中毒").clone()
    }

    /// 检查节点是否被取消
    pub fn is_cancelled(&self, node_id: &NodeId) -> bool {
        matches!(self.get(node_id), NodeStatus::Cancelled)
    }
}

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

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

    #[test]
    fn test_legal_transitions() {
        let sm = StateMachine::new();
        let id = "node-a".to_string();

        // Pending -> Running -> Completed 全程合法
        assert!(sm.set_running(id.clone()).is_ok());
        assert_eq!(sm.get(&id), NodeStatus::Running);
        assert!(sm.set_completed(id.clone()).is_ok());
        assert_eq!(sm.get(&id), NodeStatus::Completed);

        // Running -> Failed 合法
        let id2 = "node-b".to_string();
        sm.set_running(id2.clone()).unwrap();
        assert!(sm.set_failed(id2.clone()).is_ok());
        assert_eq!(sm.get(&id2), NodeStatus::Failed);
    }

    #[test]
    fn test_illegal_transitions_rejected() {
        let sm = StateMachine::new();
        let id = "node-a".to_string();

        // Pending -> Completed 非法
        let err = sm.set_completed(id.clone()).unwrap_err();
        assert!(err.to_string().contains("状态转换非法"));
        assert!(err.to_string().contains("node-a"));
        // 状态保持不变
        assert_eq!(sm.get(&id), NodeStatus::Pending);

        // Pending -> Failed 非法
        assert!(sm.set_failed(id.clone()).is_err());

        // Completed 为终态，不允许再转 Running
        sm.set_running(id.clone()).unwrap();
        sm.set_completed(id.clone()).unwrap();
        assert!(sm.set_running(id.clone()).is_err());

        // Running -> Running 重复置位非法
        let id2 = "node-b".to_string();
        sm.set_running(id2.clone()).unwrap();
        assert!(sm.set_running(id2.clone()).is_err());
    }

    #[test]
    fn test_set_cancelled_marks_node_cancelled() {
        // set_cancelled 为唯一受控旁路，不经转换校验直接置位 Cancelled
        let sm = StateMachine::new();
        let id = "human-1".to_string();

        // 初始为 Pending，is_cancelled = false
        assert_eq!(sm.get(&id), NodeStatus::Pending);
        assert!(!sm.is_cancelled(&id));

        // 取消后置位 Cancelled，is_cancelled = true
        sm.set_cancelled(id.clone());
        assert_eq!(sm.get(&id), NodeStatus::Cancelled);
        assert!(sm.is_cancelled(&id));
    }

    #[test]
    fn test_clone_shares_underlying_state() {
        // 核心机制：clone 浅拷贝 Arc，写入对另一副本可见。
        // 这让 cancel IPC 写 AppState 注册表的 StateMachine，
        // 直达 HumanNode 持有的 ctx.node_status（executor.state_machine 的 clone）。
        let sm = StateMachine::new();
        let sm_clone = sm.clone();

        // 原实例写入，clone 可见
        sm.set_cancelled("human-1".to_string());
        assert!(sm_clone.is_cancelled(&"human-1".to_string()));

        // clone 写入，原实例亦可见（双向）
        sm_clone.set_running("node-a".to_string()).unwrap();
        assert_eq!(sm.get(&"node-a".to_string()), NodeStatus::Running);

        // 第三副本同样共享
        let sm3 = sm_clone.clone();
        assert!(sm3.is_cancelled(&"human-1".to_string()));
    }
}