//! Agent Runtime 模块 //! //! 简化版的 Agent Worker 管理器,利用 SACP 的 Send trait 支持。 //! //! ## 新架构设计 //! //! - 移除独立 OS 线程,使用 `tokio::spawn` //! - 简化 sender 管理,移除 ArcSwap //! - 保留自动重启功能 //! - 保留心跳检测(僵尸检测) //! - 简化状态机(使用原子操作) //! //! ## 与旧架构对比 //! //! | 组件 | 旧设计 | 新设计 | //! |------|--------|--------| //! | 运行环境 | 独立 OS 线程 + 独立运行时 | 主运行时 + tokio::spawn | //! | Sender 管理 | ArcSwap> | mpsc::Sender (固定) | //! | Worker 生命周期 | 手动管理线程 | JoinHandle | //! | Ready 信号 | oneshot::channel | JoinHandle 完成 | //! | 状态机 | watch::Sender + Mutex | Arc | //! | 重启 | 替换 sender | abort + spawn | use std::collections::HashMap; use std::sync::Arc; use std::sync::atomic::{AtomicI64, AtomicU8, AtomicUsize, Ordering}; use std::time::Duration; use chrono::Utc; use tokio::sync::{Mutex, mpsc}; use tokio::task::JoinHandle; use tracing::{info, warn}; use crate::proxy_agent::AgentRequest; /// Worker 状态 #[repr(u8)] #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum WorkerState { /// 启动中 Starting = 0, /// 运行中 Running = 1, /// 停止中 Stopping = 2, /// 已停止 Stopped = 3, } /// 原子状态包装器 (无需 Mutex) pub struct AtomicState(AtomicU8); impl AtomicState { pub fn new(state: WorkerState) -> Self { Self(AtomicU8::new(state as u8)) } pub fn get(&self) -> WorkerState { match self.0.load(Ordering::Acquire) { 0 => WorkerState::Starting, 1 => WorkerState::Running, 2 => WorkerState::Stopping, 3 => WorkerState::Stopped, invalid => { tracing::error!( "[AtomicState] Invalid state value: {}, falling back to Stopped", invalid ); WorkerState::Stopped } } } pub fn set(&self, state: WorkerState) { self.0.store(state as u8, Ordering::Release); } } /// 心跳包 #[derive(Clone, Debug)] pub struct Heartbeat { /// 心跳时间戳 pub timestamp: chrono::DateTime, } /// Worker 就绪信号 /// /// Worker 在初始化完成后发送此信号 #[derive(Clone, Debug)] pub struct WorkerReady { /// 就绪时间戳 pub timestamp: chrono::DateTime, } /// 并发控制配置 /// /// 工作线程池大小 - 决定可以并发处理的 Agent 会话数量 /// 🔥 已改为运行时可配置的全局变量,使用 get_concurrency_limit() 获取 /// 全局并发限制(运行时可配置) pub static WORKER_THREAD_POOL_SIZE: AtomicUsize = AtomicUsize::new(10); /// 初始化并发限制(在应用启动时调用) pub fn init_concurrency_limit(limit: usize) { WORKER_THREAD_POOL_SIZE.store(limit, Ordering::Release); info!("🔧 Concurrency limit initialized: {}", limit); } /// 获取当前并发限制 pub fn get_concurrency_limit() -> usize { WORKER_THREAD_POOL_SIZE.load(Ordering::Acquire) } /// Agent 运行时 /// /// 替代 AgentWorkerManager,使用简化的架构: /// - 直接在主运行时中运行 (SACP 支持 Send) /// - 使用原子操作管理状态 /// - 使用 JoinHandle 管理生命周期 pub struct AgentRuntime { /// 请求发送端 (固定不变) request_tx: mpsc::Sender, /// 当前 Worker 的 JoinHandle worker_handle: Arc>>>, /// 当前状态 state: Arc, /// 🔥 P1 修复: 最后心跳时间戳(毫秒,Unix timestamp) /// 使用 AtomicI64 替代 Mutex,避免频繁的锁竞争 /// - 0 表示从未收到心跳 /// - 正数表示最后一次心跳的 timestamp_millis() last_heartbeat_ts: Arc, /// 活跃请求追踪: request_id -> 开始时间 active_requests: Arc>>>, /// 心跳超时阈值 heartbeat_timeout: Duration, /// 首次启动宽限期 initial_grace_period: Duration, } impl AgentRuntime { /// 创建新的 AgentRuntime /// /// 返回 (runtime, request_receiver) pub fn new(request_buffer: usize) -> (Self, mpsc::Receiver) { let (request_tx, request_rx) = mpsc::channel(request_buffer); let runtime = Self { request_tx, worker_handle: Arc::new(Mutex::new(None)), state: Arc::new(AtomicState::new(WorkerState::Starting)), last_heartbeat_ts: Arc::new(AtomicI64::new(0)), active_requests: Arc::new(Mutex::new(HashMap::new())), heartbeat_timeout: Duration::from_secs(15), initial_grace_period: Duration::from_secs(30), }; (runtime, request_rx) } /// 启动 Worker (在主运行时中) /// /// SACP 支持 Send,直接在主运行时中运行,无需独立线程 pub async fn start(&self, receiver: mpsc::Receiver) { let state = self.state.clone(); let last_heartbeat_ts = self.last_heartbeat_ts.clone(); let active_requests = self.active_requests.clone(); let handle = tokio::spawn(async move { // SACP 支持 Send,直接在主运行时中运行 if let Err(e) = crate::proxy_agent::agent_worker_with_heartbeat( receiver, state.clone(), last_heartbeat_ts.clone(), active_requests.clone(), ) .await { tracing::error!("Agent worker failed: {}", e); } }); *self.worker_handle.lock().await = Some(handle); self.state.set(WorkerState::Running); info!("AgentRuntime: worker started"); } /// 重启 Worker pub async fn restart(&self, new_receiver: mpsc::Receiver) { warn!("AgentRuntime: preparing to restart worker..."); // 1. 停止旧 worker if let Some(handle) = self.worker_handle.lock().await.take() { handle.abort(); info!("AgentRuntime: previous worker terminated"); } // 2. 重置状态 self.state.set(WorkerState::Starting); self.last_heartbeat_ts.store(0, Ordering::Release); *self.active_requests.lock().await = HashMap::new(); // 3. 启动新 worker self.start(new_receiver).await; info!("AgentRuntime: worker restart completed"); } /// 发送请求 pub async fn send(&self, request: AgentRequest) -> anyhow::Result<()> { self.request_tx .send(request) .await .map_err(|_| anyhow::anyhow!("Worker is closed"))?; Ok(()) } /// 健康检查 pub async fn check_health(&self) -> bool { let state = self.state.get(); // 检查状态 if state == WorkerState::Stopped { return false; } // 🔥 P1 修复: 使用原子操作检查心跳(无锁) let last_ts = self.last_heartbeat_ts.load(Ordering::Acquire); if last_ts > 0 { let elapsed_ms = Utc::now().timestamp_millis() - last_ts; elapsed_ms < self.heartbeat_timeout.as_millis() as i64 } else { // 首次启动宽限期 true } } /// 获取当前状态 pub fn state(&self) -> WorkerState { self.state.get() } /// 🔥 P1 修复: 检查心跳是否超时(无锁) /// /// ## 返回值 /// /// - `true`: 心跳超时(超过 15 秒未收到心跳) /// - `false`: 心跳正常或在宽限期内 pub fn check_heartbeat_timeout(&self) -> bool { let last_ts = self.last_heartbeat_ts.load(Ordering::Acquire); if last_ts > 0 { // 有心跳记录,检查是否超过 15 秒 let elapsed_ms = Utc::now().timestamp_millis() - last_ts; elapsed_ms > 15_000 } else { // 从未收到心跳,使用首次启动宽限期 false } } /// 🔥 P1 修复: 获取最后心跳时间(无锁) /// /// ## 返回值 /// /// - `Some(timestamp)`: 最后心跳时间 /// - `None`: 从未收到心跳 pub fn last_heartbeat_time(&self) -> Option> { let last_ts = self.last_heartbeat_ts.load(Ordering::Acquire); if last_ts > 0 { // 将毫秒时间戳转换为 DateTime use chrono::TimeZone; // timestamp_millis_opt 返回 LocalResult,使用 single() 转换为 Option match chrono::Utc.timestamp_millis_opt(last_ts).single() { Some(dt) => Some(dt), None => { tracing::warn!( "[WorkerInfo] Invalid timestamp: {}, using current time", last_ts ); Some(chrono::Utc::now()) } } } else { None } } /// 获取活跃请求句柄 pub fn active_requests(&self) -> Arc>>> { self.active_requests.clone() } /// 检查请求通道是否已关闭 pub fn is_closed(&self) -> bool { self.request_tx.is_closed() } } #[cfg(test)] mod tests { use super::*; #[test] fn test_atomic_state() { let state = AtomicState::new(WorkerState::Starting); assert_eq!(state.get(), WorkerState::Starting); state.set(WorkerState::Running); assert_eq!(state.get(), WorkerState::Running); state.set(WorkerState::Stopped); assert_eq!(state.get(), WorkerState::Stopped); } #[tokio::test] async fn test_agent_runtime_creation() { let (runtime, _rx) = AgentRuntime::new(100); assert_eq!(runtime.state(), WorkerState::Starting); assert!(!runtime.is_closed()); } #[tokio::test] async fn test_heartbeat_timeout_detection() { let (runtime, _rx) = AgentRuntime::new(100); // 初始状态:从未收到心跳 assert!(!runtime.check_heartbeat_timeout()); // 模拟心跳超时(设置20秒前的时间戳) let timestamp_20s_ago = Utc::now().timestamp_millis() - (20 * 1000); runtime .last_heartbeat_ts .store(timestamp_20s_ago, std::sync::atomic::Ordering::Release); // 心跳超过 15 秒,应检测到超时 assert!(runtime.check_heartbeat_timeout()); } #[tokio::test] async fn test_health_check() { let (runtime, _rx) = AgentRuntime::new(100); // 初始状态应该是健康的 assert!(runtime.check_health().await); // 设置状态为 Stopped runtime.state.set(WorkerState::Stopped); assert!(!runtime.check_health().await); } }