258 lines
9.6 KiB
Rust
258 lines
9.6 KiB
Rust
//! 节点状态机 — 管理节点的状态转换
|
||
//!
|
||
//! 合法转换：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())),
|
||
}
|
||
}
|
||
|
||
/// 获取节点状态
|
||
///
|
||
/// 锁中毒时降级返回 `NodeStatus::Pending`（保守默认：视为未启动，
|
||
/// 执行器不会误判为已完成/失败），并 `tracing::error!` 记录，不 panic。
|
||
pub fn get(&self, node_id: &NodeId) -> NodeStatus {
|
||
match self.states.lock() {
|
||
Ok(states) => states.get(node_id).cloned().unwrap_or(NodeStatus::Pending),
|
||
Err(poisoned) => {
|
||
tracing::error!(
|
||
"状态机锁中毒，get({}) 降级返回 Pending：{}",
|
||
node_id,
|
||
poisoned
|
||
);
|
||
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)
|
||
)
|
||
}
|
||
|
||
/// 状态转换 — 校验合法性后更新，非法转换返回错误
|
||
///
|
||
/// 锁中毒时降级返回 `Err`（携带「状态机锁中毒」上下文）并 `tracing::error!` 记录，
|
||
/// 由调用方决定如何处理（通常是 `set_running`/`set_completed`/`set_failed`
|
||
/// 把 Err 上抛 → 节点执行流捕获后置 Failed），不 panic 拖垮 runtime。
|
||
pub fn transition(&self, node_id: NodeId, target: NodeStatus) -> anyhow::Result<()> {
|
||
let mut states = match self.states.lock() {
|
||
Ok(guard) => guard,
|
||
Err(poisoned) => {
|
||
tracing::error!(
|
||
"状态机锁中毒，transition({}, {}) 降级返 Err：{}",
|
||
node_id,
|
||
target.as_str(),
|
||
poisoned
|
||
);
|
||
anyhow::bail!("状态机锁中毒，节点 {} 状态转换失败", node_id);
|
||
}
|
||
};
|
||
let current = states
|
||
.get(&node_id)
|
||
.cloned()
|
||
.unwrap_or(NodeStatus::Pending);
|
||
if !Self::is_legal(¤t, &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) {
|
||
match self.states.lock() {
|
||
Ok(mut states) => {
|
||
states.insert(node_id, NodeStatus::Cancelled);
|
||
}
|
||
Err(poisoned) => {
|
||
tracing::error!(
|
||
"状态机锁中毒，set_cancelled({}) 降级丢弃取消信号：{}",
|
||
node_id,
|
||
poisoned
|
||
);
|
||
}
|
||
}
|
||
}
|
||
|
||
/// 获取所有状态快照（clone 返回，调用方持独立副本）
|
||
///
|
||
/// 锁中毒时降级返回空 HashMap（调用方遍历视为无已完成节点，保守安全），
|
||
/// 并 `tracing::error!` 记录，不 panic。
|
||
pub fn snapshot(&self) -> HashMap<NodeId, NodeStatus> {
|
||
match self.states.lock() {
|
||
Ok(states) => states.clone(),
|
||
Err(poisoned) => {
|
||
tracing::error!(
|
||
"状态机锁中毒，snapshot() 降级返回空 HashMap：{}",
|
||
poisoned
|
||
);
|
||
HashMap::new()
|
||
}
|
||
}
|
||
}
|
||
|
||
/// 检查节点是否被取消
|
||
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()));
|
||
}
|
||
}
|