diff --git a/packages/client-v2/Cargo.lock b/packages/client-v2/Cargo.lock index 7a007b4..844ba00 100644 --- a/packages/client-v2/Cargo.lock +++ b/packages/client-v2/Cargo.lock @@ -847,6 +847,7 @@ dependencies = [ "postcard", "serde", "symphonia", + "thiserror", "tokio", "tokio-util", ] diff --git a/packages/client-v2/Cargo.toml b/packages/client-v2/Cargo.toml index 8baba73..4b6de5b 100644 --- a/packages/client-v2/Cargo.toml +++ b/packages/client-v2/Cargo.toml @@ -20,6 +20,7 @@ serde = { version = "1.0", features = ["derive"] } postcard = { version = "1.0", features = ["alloc", "use-std"] } parking_lot = "0.12.5" dashmap = "6.1.0" +thiserror = "2.0.17" [target.'cfg(target_os = "linux")'.dependencies] alsa = "0.11" diff --git a/packages/client-v2/src/app/client/mod.rs b/packages/client-v2/src/app/client/mod.rs index 5587dd8..68cc196 100644 --- a/packages/client-v2/src/app/client/mod.rs +++ b/packages/client-v2/src/app/client/mod.rs @@ -31,14 +31,12 @@ mod session; pub use pipeline::{PipelineHandle, PlaybackPipeline, RecordPipeline}; pub use session::Session; -use crate::net::command::{ - AudioState, Command, CommandError, CommandResult, DeviceInfo, SetVolumeResponse, ShellRequest, - ShellResponse, -}; +use crate::net::command::CommandResult; use crate::net::discovery::Discovery; use crate::net::event::{EventBus, EventBusSubscription, EventData}; use crate::net::network::{AudioSocket, Connection}; use crate::net::protocol::ControlPacket; +use crate::net::rpc::RpcBuilder; use crate::net::sync::now_us; use anyhow::{Result, anyhow}; use session::handshake; @@ -254,11 +252,22 @@ impl Client { ControlPacket::RpcResponse { id, result } => { session.resolve_rpc(id, result); } - // todo 耗时任务，异步响应 - ControlPacket::RpcRequest { id, command } => { - let result = self.handle_command(session, command).await; + // 处理 RPC 请求（支持超时和异步控制） + ControlPacket::RpcRequest { + id, + run_async, + timeout, + command, + } => { + let session_clone = session.clone(); session - .send(&ControlPacket::RpcResponse { id, result }) + .rpc_manager + .handle_rpc_request(id, run_async, timeout, command, move |pck| { + let s = session_clone.clone(); + async move { + return s.send(&pck).await; + } + }) .await?; } ControlPacket::StartRecording { config } => { @@ -294,72 +303,6 @@ impl Client { Ok(()) } - /// 处理 RPC 命令 - async fn handle_command(&self, session: &Arc, command: Command) -> CommandResult { - match command { - Command::Shell(req) => self.handle_shell(req), - Command::GetInfo => self.handle_get_info(session), - Command::Ping { timestamp } => { - let now = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_millis() as u64; - CommandResult::Pong { - timestamp, - server_time: now, - } - } - Command::SetVolume(req) => { - let prev = session.set_volume(req.volume); - CommandResult::Volume(SetVolumeResponse { - previous: prev, - current: session.volume(), - }) - } - _ => CommandResult::Error(CommandError::not_implemented()), - } - } - - /// 处理 Shell 命令 - fn handle_shell(&self, req: ShellRequest) -> CommandResult { - let mut cmd = std::process::Command::new("sh"); - cmd.arg("-c").arg(&req.command); - - if let Some(cwd) = &req.cwd { - cmd.current_dir(cwd); - } - - if let Some(env) = &req.env { - for (k, v) in env { - cmd.env(k, v); - } - } - - match cmd.output() { - Ok(output) => CommandResult::Shell(ShellResponse { - stdout: String::from_utf8_lossy(&output.stdout).to_string(), - stderr: String::from_utf8_lossy(&output.stderr).to_string(), - exit_code: output.status.code().unwrap_or(-1), - }), - Err(e) => CommandResult::Error(CommandError::internal(e.to_string())), - } - } - - /// 处理获取设备信息 - fn handle_get_info(&self, session: &Arc) -> CommandResult { - CommandResult::Info(DeviceInfo { - model: self.config.model.clone(), - serial_number: self.config.serial_number.clone(), - version: env!("CARGO_PKG_VERSION").to_string(), - uptime_secs: session.uptime_secs(), - audio_state: AudioState { - is_recording: session.is_recording(), - is_playing: session.is_playing(), - volume: session.volume(), - }, - }) - } - /// 清理资源 async fn cleanup(&self) { let mut session_guard = self.session.write().await; @@ -372,25 +315,15 @@ impl Client { // ==================== 公开 API ==================== /// 执行 RPC 命令 - pub async fn execute(&self, command: Command) -> Result { + pub async fn rpc(&self, builder: &RpcBuilder) -> Result { let session_guard = self.session.read().await; if let Some(session) = session_guard.as_ref() { - session.execute(command).await + session.rpc(builder).await } else { Err(anyhow!("Not connected")) } } - /// 执行 Shell 命令 - pub async fn shell(&self, cmd: &str) -> Result { - let result = self.execute(Command::shell(cmd)).await?; - match result { - CommandResult::Shell(resp) => Ok(resp), - CommandResult::Error(e) => Err(anyhow!("{}", e)), - _ => Err(anyhow!("Unexpected response type")), - } - } - /// 发送事件 pub async fn send_event(&self, event: EventData) -> Result<()> { let session_guard = self.session.read().await; diff --git a/packages/client-v2/src/app/client/pipeline.rs b/packages/client-v2/src/app/client/pipeline.rs index fcd4ea6..a158c54 100644 --- a/packages/client-v2/src/app/client/pipeline.rs +++ b/packages/client-v2/src/app/client/pipeline.rs @@ -28,7 +28,6 @@ use crate::net::protocol::AudioPacket; use crate::net::sync::now_us; use std::net::SocketAddr; use std::sync::Arc; -use std::time::Duration; use tokio::sync::mpsc; use tokio_util::sync::CancellationToken; diff --git a/packages/client-v2/src/app/client/session.rs b/packages/client-v2/src/app/client/session.rs index 040712c..c9c0c73 100644 --- a/packages/client-v2/src/app/client/session.rs +++ b/packages/client-v2/src/app/client/session.rs @@ -6,10 +6,10 @@ //! - 会话生命周期 //! - 活动音频流追踪 -use crate::net::command::{Command, CommandResult}; +use crate::net::command::CommandResult; use crate::net::network::{AudioSocket, Connection}; use crate::net::protocol::{ClientInfo, ControlPacket}; -use crate::net::rpc::RpcManager; +use crate::net::rpc::{RpcBuilder, RpcManager}; use crate::net::sync::ClockSync; use anyhow::{Context, Result, anyhow}; use std::net::SocketAddr; @@ -90,7 +90,7 @@ pub struct Session { pub server_audio_addr: SocketAddr, /// RPC 管理器 - pub rpc: Arc, + pub rpc_manager: Arc, /// 会话取消令牌 pub cancel: CancellationToken, @@ -120,7 +120,7 @@ impl Session { conn, audio_socket, server_audio_addr, - rpc: Arc::new(RpcManager::new()), + rpc_manager: Arc::new(RpcManager::new()), cancel, clock: Arc::new(parking_lot::Mutex::new(ClockSync::new(100))), pipelines: parking_lot::Mutex::new(ActivePipelines::default()), @@ -161,18 +161,20 @@ impl Session { .update(client_send_ts, server_ts, client_recv_ts); } - /// 发起 RPC 调用（新版） - pub async fn execute(&self, command: Command) -> Result { - let (id, rx) = self.rpc.register(); - self.conn - .send(&ControlPacket::RpcRequest { id, command }) + /// 发起 RPC 调用（支持超时和异步控制） + pub async fn rpc(&self, request: &RpcBuilder) -> Result { + let result = self + .rpc_manager + .call(request, |pck| async move { + return self.conn.send(&pck).await; + }) .await?; - rx.await.context("RPC channel closed") + Ok(result) } /// 处理 RPC 响应 pub fn resolve_rpc(&self, id: u32, result: CommandResult) { - self.rpc.resolve(id, result); + self.rpc_manager.resolve(id, result); } /// 开始录音管道 diff --git a/packages/client-v2/src/app/server/audio_bus.rs b/packages/client-v2/src/app/server/audio_bus.rs index 4cc0fbc..fda8449 100644 --- a/packages/client-v2/src/app/server/audio_bus.rs +++ b/packages/client-v2/src/app/server/audio_bus.rs @@ -229,22 +229,3 @@ impl AudioBus { self.subscribers.contains_key(addr) } } - -#[cfg(test)] -mod tests { - use super::*; - - #[tokio::test] - async fn test_audio_bus_basic() { - let bus = AudioBus::new().await.unwrap(); - let addr: SocketAddr = "127.0.0.1:12345".parse().unwrap(); - - let id = bus.register(addr, true); - assert!(bus.is_registered(&addr)); - assert_eq!(bus.subscriber_count(), 1); - - bus.unregister(&addr); - assert!(!bus.is_registered(&addr)); - assert_eq!(bus.subscriber_count(), 0); - } -} diff --git a/packages/client-v2/src/app/server/mod.rs b/packages/client-v2/src/app/server/mod.rs index 6456b4d..7cbcd50 100644 --- a/packages/client-v2/src/app/server/mod.rs +++ b/packages/client-v2/src/app/server/mod.rs @@ -44,13 +44,12 @@ pub use stream::{FilePlaybackStream, RecorderStream, StreamHandle}; use crate::audio::config::AudioConfig; use crate::audio::wav::WavReader; -use crate::net::command::{ - AudioState, Command, CommandError, CommandResult, DeviceInfo, ShellResponse, -}; +use crate::net::command::CommandResult; use crate::net::discovery::Discovery; use crate::net::event::{EventBus, EventBusSubscription, EventData}; use crate::net::network::Connection; use crate::net::protocol::ControlPacket; +use crate::net::rpc::RpcBuilder; use crate::net::sync::now_us; use anyhow::{Context, Result, anyhow}; use std::net::SocketAddr; @@ -289,64 +288,28 @@ impl Server { ControlPacket::RpcResponse { id, result } => { session.resolve_rpc(id, result); } - // todo 耗时任务，异步响应 - ControlPacket::RpcRequest { id, command } => { - let result = self.handle_command(session, command).await; + ControlPacket::RpcRequest { + id, + run_async, + timeout, + command, + } => { + let session_clone = session.clone(); session - .send(&ControlPacket::RpcResponse { id, result }) + .rpc_manager + .handle_rpc_request(id, run_async, timeout, command, move |pck| { + let s = session_clone.clone(); + async move { + return s.send(&pck).await; + } + }) .await?; } - _ => {} } Ok(()) } - /// 处理 RPC 命令 - async fn handle_command(&self, session: &Arc, command: Command) -> CommandResult { - match command { - Command::Shell(req) => { - // 转发给客户端执行 - match session.execute(Command::Shell(req)).await { - Ok(result) => result, - Err(e) => CommandResult::Error(CommandError::internal(e.to_string())), - } - } - Command::GetInfo => { - // 返回服务器信息 - CommandResult::Info(DeviceInfo { - model: "XiaoAi-Server".to_string(), - serial_number: "SERVER-001".to_string(), - version: env!("CARGO_PKG_VERSION").to_string(), - uptime_secs: self.started_at.elapsed().as_secs(), - audio_state: AudioState { - is_recording: session.is_recording(), - is_playing: session.is_playing(), - volume: 100, - }, - }) - } - Command::Ping { timestamp } => { - let now = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_millis() as u64; - CommandResult::Pong { - timestamp, - server_time: now, - } - } - Command::SetVolume(_) => { - // 转发给客户端 - match session.execute(command).await { - Ok(result) => result, - Err(e) => CommandResult::Error(CommandError::internal(e.to_string())), - } - } - _ => CommandResult::Error(CommandError::not_implemented()), - } - } - // ==================== 公开 API ==================== /// 获取所有已连接的客户端地址 @@ -364,20 +327,10 @@ impl Server { self.started_at.elapsed().as_secs() } - /// 向客户端发起 RPC 调用（新版） - pub async fn execute(&self, addr: SocketAddr, command: Command) -> Result { + /// 向客户端发起 RPC + pub async fn rpc(&self, addr: SocketAddr, builder: &RpcBuilder) -> Result { let session = self.sessions.get(&addr).context("Session not found")?; - session.execute(command).await - } - - /// 执行 Shell 命令 - pub async fn shell(&self, addr: SocketAddr, cmd: &str) -> Result { - let result = self.execute(addr, Command::shell(cmd)).await?; - match result { - CommandResult::Shell(resp) => Ok(resp), - CommandResult::Error(e) => Err(anyhow!("{}", e)), - _ => Err(anyhow!("Unexpected response type")), - } + session.rpc(builder).await } /// 发送事件 diff --git a/packages/client-v2/src/app/server/session.rs b/packages/client-v2/src/app/server/session.rs index 05cc33e..1126307 100644 --- a/packages/client-v2/src/app/server/session.rs +++ b/packages/client-v2/src/app/server/session.rs @@ -8,10 +8,10 @@ //! 音频流的实际处理由 AudioBus 和 Stream 模块负责。 use crate::audio::config::AudioConfig; -use crate::net::command::{Command, CommandResult}; +use crate::net::command::CommandResult; use crate::net::network::Connection; use crate::net::protocol::{ClientInfo, ControlPacket}; -use crate::net::rpc::RpcManager; +use crate::net::rpc::{RpcBuilder, RpcManager}; use anyhow::{Context, Result}; use std::net::SocketAddr; use std::sync::Arc; @@ -98,7 +98,7 @@ pub struct Session { pub conn: Arc, /// RPC 管理器 - pub rpc: Arc, + pub rpc_manager: Arc, /// TCP 地址（用作会话 ID） pub tcp_addr: SocketAddr, @@ -131,7 +131,7 @@ impl Session { Self { info, conn, - rpc: Arc::new(RpcManager::new()), + rpc_manager: Arc::new(RpcManager::new()), tcp_addr, audio_addr, cancel, @@ -166,18 +166,20 @@ impl Session { self.conn.recv().await } - /// 执行 RPC 调用（新版，使用 Command） - pub async fn execute(&self, command: Command) -> Result { - let (id, rx) = self.rpc.register(); - self.conn - .send(&ControlPacket::RpcRequest { id, command }) + /// 发起 RPC 调用（支持超时和异步控制） + pub async fn rpc(&self, request: &RpcBuilder) -> Result { + let result = self + .rpc_manager + .call(request, |pck| async move { + return self.conn.send(&pck).await; + }) .await?; - rx.await.context("RPC channel closed") + Ok(result) } /// 处理 RPC 响应 pub fn resolve_rpc(&self, id: u32, result: CommandResult) { - self.rpc.resolve(id, result); + self.rpc_manager.resolve(id, result); } /// 开始录音流 diff --git a/packages/client-v2/src/bin/client.rs b/packages/client-v2/src/bin/client.rs index f1471d9..135f091 100644 --- a/packages/client-v2/src/bin/client.rs +++ b/packages/client-v2/src/bin/client.rs @@ -10,6 +10,7 @@ use std::sync::Arc; use xiao::app::client::{Client, ClientConfig}; use xiao::net::command::Command; use xiao::net::event::EventData; +use xiao::net::rpc::RpcBuilder; #[tokio::main] async fn main() -> anyhow::Result<()> { @@ -61,22 +62,20 @@ async fn main() -> anyhow::Result<()> { println!("Connected to server!"); println!("Running client-side tests...\n"); - // 1. 测试向服务器发送 Ping - println!("1️⃣ Testing Ping to server..."); - match client.execute(Command::ping()).await { - Ok(result) => println!(" ✅ Pong received: {:?}", result), - Err(e) => println!(" ❌ Ping failed: {}", e), + // 1. 测试简单的 Shell 命令 + println!("1️⃣ Testing Shell command..."); + match client + .rpc(&RpcBuilder::default(Command::Shell( + "echo 'Hello from client'".to_string(), + ))) + .await + { + Ok(resp) => println!(" ✅ Command executed: {:?}", resp), + Err(e) => println!(" ❌ Command failed: {}", e), } - // 2. 测试获取服务器信息 - println!("\n2️⃣ Testing GetInfo from server..."); - match client.execute(Command::GetInfo).await { - Ok(result) => println!(" ✅ Server info: {:?}", result), - Err(e) => println!(" ❌ GetInfo failed: {}", e), - } - - // 3. 发送客户端事件 - println!("\n3️⃣ Sending event to server..."); + // 2. 发送客户端事件 + println!("\n2️⃣ Sending event to server..."); match client .send_event(EventData::Hello { message: "from client!".to_string(), diff --git a/packages/client-v2/src/bin/server.rs b/packages/client-v2/src/bin/server.rs index d02aee3..85ee873 100644 --- a/packages/client-v2/src/bin/server.rs +++ b/packages/client-v2/src/bin/server.rs @@ -11,6 +11,7 @@ use xiao::app::server::{Server, ServerConfig}; use xiao::audio::config::AudioConfig; use xiao::net::command::Command; use xiao::net::event::EventData; +use xiao::net::rpc::RpcBuilder; #[tokio::main] async fn main() -> anyhow::Result<()> { @@ -65,32 +66,34 @@ async fn main() -> anyhow::Result<()> { println!("Client connected: {}", addr); println!("Running demo tests...\n"); - // 1. 测试 Ping - println!("1️⃣ Testing Ping..."); - match server.execute(addr, Command::ping()).await { - Ok(result) => println!(" ✅ Ping result: {:?}", result), - Err(e) => println!(" ❌ Ping failed: {}", e), + // 1. 测试 Shell 命令（同步） + println!("1️⃣ Testing Shell RPC (sync)..."); + match server + .rpc( + addr, + &RpcBuilder::default(Command::Shell("echo 'Hello from server'".to_string())), + ) + .await + { + Ok(resp) => println!(" ✅ Command executed: {:?}", resp), + Err(e) => println!(" ❌ Command failed: {}", e), } - // 2. 测试获取设备信息 - println!("\n2️⃣ Testing GetInfo..."); - match server.execute(addr, Command::GetInfo).await { - Ok(result) => println!(" ✅ Device info: {:?}", result), - Err(e) => println!(" ❌ GetInfo failed: {}", e), + // 2. 测试 Shell 命令（带超时） + println!("\n2️⃣ Testing Shell RPC with timeout (1s)..."); + match server + .rpc( + addr, + &RpcBuilder::default(Command::Shell("sleep 2".to_string())).set_timeout(1000), + ) + .await + { + Ok(resp) => println!(" ✅ Command executed: {:?}", resp), + Err(e) => println!(" ❌ Command failed: {}", e), } - // 3. 测试 Shell 命令 - println!("\n3️⃣ Testing Shell RPC..."); - match server.shell(addr, "echo 'Hello from XiaoAi!'").await { - Ok(resp) => { - println!(" ✅ stdout: {}", resp.stdout.trim()); - println!(" ✅ exit_code: {}", resp.exit_code); - } - Err(e) => println!(" ❌ Shell failed: {}", e), - } - - // 4. 测试事件 - println!("\n4️⃣ Broadcasting notification event..."); + // 3. 测试事件 + println!("\n3️⃣ Broadcasting notification event..."); match server .send_event( addr, diff --git a/packages/client-v2/src/lib.rs b/packages/client-v2/src/lib.rs index db47768..bc714b2 100644 --- a/packages/client-v2/src/lib.rs +++ b/packages/client-v2/src/lib.rs @@ -1,3 +1,4 @@ pub mod app; pub mod audio; pub mod net; +pub mod utils; diff --git a/packages/client-v2/src/net/command.rs b/packages/client-v2/src/net/command.rs index 33e9142..8f7f198 100644 --- a/packages/client-v2/src/net/command.rs +++ b/packages/client-v2/src/net/command.rs @@ -28,147 +28,10 @@ use serde::{Deserialize, Serialize}; +use crate::utils::shell::{ShellRequest, ShellResponse}; + // ==================== 命令请求类型 ==================== -/// Shell 命令请求 -#[derive(Serialize, Deserialize, Debug, Clone)] -pub struct ShellRequest { - /// 要执行的命令 - pub command: String, - /// 工作目录（可选） - pub cwd: Option, - /// 环境变量（可选） - pub env: Option>, - /// 超时秒数（可选） - pub timeout_secs: Option, -} - -/// Shell 命令响应 -#[derive(Serialize, Deserialize, Debug, Clone, Default)] -pub struct ShellResponse { - pub stdout: String, - pub stderr: String, - pub exit_code: i32, -} - -/// 获取设备信息请求 -#[derive(Serialize, Deserialize, Debug, Clone, Default)] -pub struct GetInfoRequest; - -/// 设备信息响应 -#[derive(Serialize, Deserialize, Debug, Clone)] -pub struct DeviceInfo { - pub model: String, - pub serial_number: String, - pub version: String, - pub uptime_secs: u64, - pub audio_state: AudioState, -} - -/// 音频状态 -#[derive(Serialize, Deserialize, Debug, Clone, Default)] -pub struct AudioState { - pub is_recording: bool, - pub is_playing: bool, - pub volume: u8, -} - -/// 设置音量请求 -#[derive(Serialize, Deserialize, Debug, Clone)] -pub struct SetVolumeRequest { - pub volume: u8, // 0-100 -} - -/// 设置音量响应 -#[derive(Serialize, Deserialize, Debug, Clone)] -pub struct SetVolumeResponse { - pub previous: u8, - pub current: u8, -} - -/// 文件操作请求 -#[derive(Serialize, Deserialize, Debug, Clone)] -pub enum FileRequest { - /// 读取文件 - Read { path: String }, - /// 写入文件 - Write { path: String, data: Vec }, - /// 删除文件 - Delete { path: String }, - /// 列出目录 - List { path: String }, - /// 获取文件信息 - Stat { path: String }, -} - -/// 文件操作响应 -#[derive(Serialize, Deserialize, Debug, Clone)] -pub enum FileResponse { - /// 读取结果 - Data(Vec), - /// 写入成功 - Written { bytes: usize }, - /// 删除成功 - Deleted, - /// 目录列表 - Entries(Vec), - /// 文件信息 - Stat(FileStat), -} - -/// 文件条目 -#[derive(Serialize, Deserialize, Debug, Clone)] -pub struct FileEntry { - pub name: String, - pub is_dir: bool, - pub size: u64, -} - -/// 文件状态 -#[derive(Serialize, Deserialize, Debug, Clone)] -pub struct FileStat { - pub size: u64, - pub is_dir: bool, - pub modified: u64, -} - -/// 系统控制请求 -#[derive(Serialize, Deserialize, Debug, Clone)] -pub enum SystemRequest { - /// 重启 - Reboot, - /// 关机 - Shutdown, - /// 获取系统负载 - GetLoad, - /// 获取内存使用 - GetMemory, -} - -/// 系统控制响应 -#[derive(Serialize, Deserialize, Debug, Clone)] -pub enum SystemResponse { - /// 操作已接受 - Accepted, - /// 系统负载 - Load { one: f32, five: f32, fifteen: f32 }, - /// 内存使用 - Memory { total: u64, used: u64, free: u64 }, -} - -/// 自定义命令请求（用于扩展） -#[derive(Serialize, Deserialize, Debug, Clone)] -pub struct CustomRequest { - pub name: String, - pub payload: Vec, -} - -/// 自定义命令响应 -#[derive(Serialize, Deserialize, Debug, Clone)] -pub struct CustomResponse { - pub payload: Vec, -} - // ==================== 统一命令枚举 ==================== /// RPC 命令 - 统一的请求类型 @@ -176,39 +39,15 @@ pub struct CustomResponse { pub enum Command { /// 执行 Shell 命令 Shell(ShellRequest), - /// 获取设备信息 - GetInfo, - /// 设置音量 - SetVolume(SetVolumeRequest), - /// 文件操作 - File(FileRequest), - /// 系统控制 - System(SystemRequest), - /// 自定义命令 - Custom(CustomRequest), - /// Ping（用于测量延迟） - Ping { timestamp: u64 }, } /// RPC 结果 - 统一的响应类型 #[derive(Serialize, Deserialize, Debug, Clone)] pub enum CommandResult { - /// Shell 命令结果 - Shell(ShellResponse), - /// 设备信息 - Info(DeviceInfo), - /// 音量设置结果 - Volume(SetVolumeResponse), - /// 文件操作结果 - File(FileResponse), - /// 系统控制结果 - System(SystemResponse), - /// 自定义命令结果 - Custom(CustomResponse), - /// Pong 响应 - Pong { timestamp: u64, server_time: u64 }, /// 错误 Error(CommandError), + /// Shell 命令结果 + Shell(ShellResponse), } /// 命令错误 @@ -261,48 +100,7 @@ impl std::error::Error for CommandError {} // ==================== 便捷构造方法 ==================== -impl Command { - /// 创建 Shell 命令 - pub fn shell(cmd: impl Into) -> Self { - Self::Shell(ShellRequest { - command: cmd.into(), - cwd: None, - env: None, - timeout_secs: None, - }) - } - - /// 创建带超时的 Shell 命令 - pub fn shell_with_timeout(cmd: impl Into, timeout: u32) -> Self { - Self::Shell(ShellRequest { - command: cmd.into(), - cwd: None, - env: None, - timeout_secs: Some(timeout), - }) - } - - /// 创建 Ping 命令 - pub fn ping() -> Self { - Self::Ping { - timestamp: std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_millis() as u64, - } - } -} - impl CommandResult { - /// 创建成功的 Shell 响应 - pub fn shell_ok(stdout: String) -> Self { - Self::Shell(ShellResponse { - stdout, - stderr: String::new(), - exit_code: 0, - }) - } - /// 创建错误响应 pub fn error(err: CommandError) -> Self { Self::Error(err) diff --git a/packages/client-v2/src/net/protocol.rs b/packages/client-v2/src/net/protocol.rs index 749818d..4d76d5e 100644 --- a/packages/client-v2/src/net/protocol.rs +++ b/packages/client-v2/src/net/protocol.rs @@ -65,6 +65,11 @@ pub enum ControlPacket { // ========== RPC ========== RpcRequest { id: u32, + /// 是否异步运行 + run_async: bool, + /// 超时（毫秒） + timeout: Option, + /// 请求命令 command: Command, }, RpcResponse { diff --git a/packages/client-v2/src/net/rpc.rs b/packages/client-v2/src/net/rpc.rs index 0411be5..593a4c3 100644 --- a/packages/client-v2/src/net/rpc.rs +++ b/packages/client-v2/src/net/rpc.rs @@ -4,138 +4,226 @@ //! - 请求 ID 生成 //! - 请求/响应匹配 //! - 超时处理 +//! - 异步任务管理 -use crate::net::command::CommandResult; -use parking_lot::Mutex; -use std::collections::HashMap; +use crate::net::command::{Command, CommandError, CommandResult}; +use crate::net::protocol::ControlPacket; +use crate::utils::shell::Shell; +use anyhow::Result; +use dashmap::DashMap; +use std::future::Future; +use std::sync::Arc; use std::sync::atomic::{AtomicU32, Ordering}; use std::time::Duration; +use thiserror::Error; use tokio::sync::oneshot; -/// RPC 调用错误 -#[derive(Debug)] +#[derive(Debug, Error)] pub enum RpcError { - /// 超时 + #[error("RPC timeout")] Timeout, - /// 通道关闭 + #[error("RPC cancelled")] Cancelled, - /// 未连接 - NotConnected, + #[error("RPC internal error: {0}")] + Internal(String), } -impl std::fmt::Display for RpcError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - RpcError::Timeout => write!(f, "RPC timeout"), - RpcError::Cancelled => write!(f, "RPC cancelled"), - RpcError::NotConnected => write!(f, "Not connected"), - } - } -} - -impl std::error::Error for RpcError {} - -/// 等待中的 RPC 请求 -struct PendingRequest { - tx: oneshot::Sender, - created_at: std::time::Instant, -} - -/// RPC 管理器 pub struct RpcManager { + handler: RpcHandler, next_id: AtomicU32, - pending: Mutex>, - default_timeout: Duration, + pending: Arc>>, } impl RpcManager { - /// 创建新的 RPC 管理器 pub fn new() -> Self { - Self::with_timeout(Duration::from_secs(30)) - } - - /// 创建带自定义超时的 RPC 管理器 - pub fn with_timeout(timeout: Duration) -> Self { Self { + handler: RpcHandler::new(), next_id: AtomicU32::new(1), - pending: Mutex::new(HashMap::new()), - default_timeout: timeout, + pending: Arc::new(DashMap::new()), } } - /// 注册一个新的 RPC 请求 - /// - /// 返回 (请求ID, 响应接收器) - pub fn register(&self) -> (u32, oneshot::Receiver) { - let id = self.next_id.fetch_add(1, Ordering::SeqCst); - let (tx, rx) = oneshot::channel(); - self.pending.lock().insert( - id, - PendingRequest { - tx, - created_at: std::time::Instant::now(), - }, - ); - (id, rx) - } - - /// 解析 RPC 响应 pub fn resolve(&self, id: u32, result: CommandResult) { - if let Some(req) = self.pending.lock().remove(&id) { - let _ = req.tx.send(result); + if let Some((_, tx)) = self.pending.remove(&id) { + let _ = tx.send(result); } } - /// 取消指定的 RPC 请求 pub fn cancel(&self, id: u32) { - self.pending.lock().remove(&id); + self.pending.remove(&id); } - /// 获取待处理请求数量 - pub fn pending_count(&self) -> usize { - self.pending.lock().len() + /// 发起异步 RPC 调用 + /// + /// F: 异步闭包，输入 ID，返回 Future + pub async fn call( + &self, + builder: &RpcBuilder, + send_fn: F, + ) -> Result + where + F: FnOnce(ControlPacket) -> Fut, + Fut: Future>, + { + let command = builder.command.clone(); + let timeout = builder.timeout; + let run_async = builder.run_async; + let id = self.next_id.fetch_add(1, Ordering::Relaxed); + let (tx, rx) = oneshot::channel(); + + // 1. 注册请求 + self.pending.insert(id, tx); + + // 2. RAII 守卫：确保无论因超时、报错还是外部 Drop，ID 都会被清理 + let _guard = DropGuard { + id, + pending: Arc::clone(&self.pending), + }; + + // 3. 执行异步发送逻辑 + if let Err(e) = send_fn(ControlPacket::RpcRequest { + id, + command, + timeout, + run_async, + }) + .await + { + return Err(RpcError::Internal(e.to_string())); + } + + // 4. 等待响应或超时 + match timeout { + Some(ms) => match tokio::time::timeout(Duration::from_millis(ms), rx).await { + Ok(Ok(res)) => Ok(res), + Ok(Err(_)) => Err(RpcError::Cancelled), // tx 关闭，任务被取消 + Err(_) => Err(RpcError::Timeout), + }, + None => rx.await.map_err(|_| RpcError::Cancelled), // tx 关闭，任务被取消 + } } - /// 清理超时的请求 - pub fn cleanup_expired(&self) { - let mut pending = self.pending.lock(); - let now = std::time::Instant::now(); + /// 处理 RPC 调用 + /// + /// F: 异步闭包，输入 ID，返回 Future + pub async fn handle_rpc_request( + &self, + id: u32, + run_async: bool, + timeout_ms: Option, + command: Command, + send_fn: F, + ) -> Result<()> + where + F: Fn(ControlPacket) -> Fut + Send + Sync + 'static, // 异步执行需要 'static + Fut: Future> + Send, + { + let sender = Arc::new(send_fn); + let handler = self.handler.clone(); - pending.retain(|_, req| { - if now.duration_since(req.created_at) > self.default_timeout { - false - } else { - true + let task = async move { + let execute_future = handler.handle(command, timeout_ms); + + match timeout_ms { + None => { + // 无超时逻辑 + let result = execute_future.await; + sender(ControlPacket::RpcResponse { id, result }).await + } + Some(ms) => { + let d = Duration::from_millis(ms); + match tokio::time::timeout(d, execute_future).await { + Ok(result) => { + // 1. 正常完成 + sender(ControlPacket::RpcResponse { id, result }).await + } + Err(_) => { + // 2. 触发超时：发送超时错误消息 + let err_result = CommandResult::error(CommandError::timeout(format!( + "Task {} timed out after {}ms", + id, ms + ))); + sender(ControlPacket::RpcResponse { + id, + result: err_result, + }) + .await + } + } + } } - }); - } + }; - /// 获取默认超时时间 - pub fn default_timeout(&self) -> Duration { - self.default_timeout + if run_async { + tokio::spawn(task); + Ok(()) + } else { + task.await + } } } -impl Default for RpcManager { - fn default() -> Self { - Self::new() +#[derive(Debug, Clone)] +pub struct RpcBuilder { + pub command: Command, + pub run_async: bool, + pub timeout: Option, +} + +impl RpcBuilder { + /// 基础构造器 + pub fn new(command: Command) -> Self { + Self { + command, + run_async: false, + timeout: None, + } + } + + /// 便捷方法：默认异步配置 + pub fn default(command: Command) -> Self { + Self::new(command).set_async(true).set_timeout(10_000) + } + + /// 设置为异步运行 + pub fn set_async(mut self, is_async: bool) -> Self { + self.run_async = is_async; + self + } + + /// 设置超时时长（毫秒） + pub fn set_timeout(mut self, ms: u64) -> Self { + self.timeout = Some(ms); + self } } -/// RPC 调用辅助函数 -pub async fn call_with_timeout( - timeout: Duration, - register_fn: F, -) -> Result -where - F: FnOnce() -> Fut, - Fut: std::future::Future, RpcError>>, -{ - let rx = register_fn().await?; +#[derive(Debug, Clone)] +pub struct RpcHandler; - match tokio::time::timeout(timeout, rx).await { - Ok(Ok(result)) => Ok(result), - Ok(Err(_)) => Err(RpcError::Cancelled), - Err(_) => Err(RpcError::Timeout), +impl RpcHandler { + fn new() -> Self { + Self {} + } + + async fn handle(&self, command: Command, timeout_ms: Option) -> CommandResult { + match command { + Command::Shell(req) => match Shell::run(req, timeout_ms).await { + Ok(resp) => CommandResult::Shell(resp), + Err(e) => CommandResult::Error(CommandError::internal(e.to_string())), + }, + _ => CommandResult::Error(CommandError::not_implemented()), + } + } +} + +struct DropGuard { + id: u32, + pending: Arc>>, +} + +impl Drop for DropGuard { + fn drop(&mut self) { + self.pending.remove(&self.id); } } diff --git a/packages/client-v2/src/utils/mod.rs b/packages/client-v2/src/utils/mod.rs new file mode 100644 index 0000000..327cf1b --- /dev/null +++ b/packages/client-v2/src/utils/mod.rs @@ -0,0 +1 @@ +pub mod shell; diff --git a/packages/client-v2/src/utils/shell.rs b/packages/client-v2/src/utils/shell.rs new file mode 100644 index 0000000..b4950ac --- /dev/null +++ b/packages/client-v2/src/utils/shell.rs @@ -0,0 +1,68 @@ +use anyhow::{Context, Result, anyhow}; +use core::str; +use serde::{Deserialize, Serialize}; +use std::process::Stdio; +use std::time::Duration; +use tokio::io::AsyncReadExt; +use tokio::process::Command; + +pub type ShellRequest = String; + +#[derive(Serialize, Deserialize, Debug, Clone, Default)] +pub struct ShellResponse { + pub stdout: String, + pub stderr: String, + pub exit_code: i32, +} + +pub struct Shell; + +impl Shell { + pub async fn run(command: String, timeout_ms: Option) -> Result { + // 1. 配置 Command + let mut child = Command::new("sh") + .arg("-c") + .arg(&command) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + // 当 Child 结构体被 Drop 时，自动杀死子进程（SIGKILL） + .kill_on_drop(true) + .spawn() + .context("Failed to spawn shell command")?; + + let mut stdout_reader = child.stdout.take().context("Failed to open stdout")?; + let mut stderr_reader = child.stderr.take().context("Failed to open stderr")?; + + // 2. 定义执行逻辑 + let run_logic = async move { + let mut stdout_buf = Vec::new(); + let mut stderr_buf = Vec::new(); + + // 使用 join! 并发读取流并等待进程结束 + let (res_out, res_err, res_status) = tokio::join!( + stdout_reader.read_to_end(&mut stdout_buf), + stderr_reader.read_to_end(&mut stderr_buf), + child.wait() + ); + + let status = res_status.context("Wait for child failed")?; + res_out.context("Read stdout failed")?; + res_err.context("Read stderr failed")?; + + Ok(ShellResponse { + stdout: String::from_utf8_lossy(&stdout_buf).to_string(), + stderr: String::from_utf8_lossy(&stderr_buf).to_string(), + exit_code: status.code().unwrap_or(-1), + }) + }; + + // 3. 应用超时 + if let Some(ms) = timeout_ms { + tokio::time::timeout(Duration::from_millis(ms), run_logic) + .await + .map_err(|_| anyhow!("Shell execution timed out"))? + } else { + run_logic.await + } + } +}