Hass/open-xiaoai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: 异步 RPC

Browse Source
This commit is contained in:
Del Wang
2026-01-08 09:53:59 +08:00
parent fb6c1fd9f8
commit 3cf6771cc9
16 changed files with 364 additions and 529 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/client-v2/Cargo.lock
Generated
+1
View File
@@ -847,6 +847,7 @@ dependencies = [
"postcard",
"serde",
"symphonia",
"thiserror",
"tokio",
"tokio-util",
]
packages/client-v2/Cargo.toml
+1
View File
@@ -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"
packages/client-v2/src/app/client/mod.rs
+19 -86
View File
@@ -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<Session>, 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<Session>) -> 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<CommandResult> {
pub async fn rpc(&self, builder: &RpcBuilder) -> Result<CommandResult> {
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<ShellResponse> {
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;
packages/client-v2/src/app/client/pipeline.rs
-1
View File
@@ -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;
packages/client-v2/src/app/client/session.rs
+13 -11
View File
@@ -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<RpcManager>,
pub rpc_manager: Arc<RpcManager>,
/// 会话取消令牌
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<CommandResult> {
let (id, rx) = self.rpc.register();
self.conn
.send(&ControlPacket::RpcRequest { id, command })
/// 发起 RPC 调用（支持超时和异步控制）
pub async fn rpc(&self, request: &RpcBuilder) -> Result<CommandResult> {
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);
}
/// 开始录音管道
packages/client-v2/src/app/server/audio_bus.rs
-19
View File
@@ -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);
}
}
packages/client-v2/src/app/server/mod.rs
+19 -66
View File
@@ -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<Session>, 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<CommandResult> {
/// 向客户端发起 RPC
pub async fn rpc(&self, addr: SocketAddr, builder: &RpcBuilder) -> Result<CommandResult> {
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<ShellResponse> {
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
}
/// 发送事件
packages/client-v2/src/app/server/session.rs
+13 -11
View File
@@ -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<Connection>,
/// RPC 管理器
pub rpc: Arc<RpcManager>,
pub rpc_manager: Arc<RpcManager>,
/// 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<CommandResult> {
let (id, rx) = self.rpc.register();
self.conn
.send(&ControlPacket::RpcRequest { id, command })
/// 发起 RPC 调用（支持超时和异步控制）
pub async fn rpc(&self, request: &RpcBuilder) -> Result<CommandResult> {
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);
}
/// 开始录音流
packages/client-v2/src/bin/client.rs
+13 -14
View File
@@ -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(),
packages/client-v2/src/bin/server.rs
+25 -22
View File
@@ -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,
packages/client-v2/src/lib.rs
+1
View File
@@ -1,3 +1,4 @@
pub mod app;
pub mod audio;
pub mod net;
pub mod utils;
packages/client-v2/src/net/command.rs
+4 -206
View File
@@ -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<String>,
/// 环境变量（可选）
pub env: Option<Vec<(String, String)>>,
/// 超时秒数（可选）
pub timeout_secs: Option<u32>,
}
/// 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<u8> },
/// 删除文件
Delete { path: String },
/// 列出目录
List { path: String },
/// 获取文件信息
Stat { path: String },
}
/// 文件操作响应
#[derive(Serialize, Deserialize, Debug, Clone)]
pub enum FileResponse {
/// 读取结果
Data(Vec<u8>),
/// 写入成功
Written { bytes: usize },
/// 删除成功
Deleted,
/// 目录列表
Entries(Vec<FileEntry>),
/// 文件信息
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<u8>,
}
/// 自定义命令响应
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct CustomResponse {
pub payload: Vec<u8>,
}
// ==================== 统一命令枚举 ====================
/// 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<String>) -> Self {
Self::Shell(ShellRequest {
command: cmd.into(),
cwd: None,
env: None,
timeout_secs: None,
})
}
/// 创建带超时的 Shell 命令
pub fn shell_with_timeout(cmd: impl Into<String>, 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)
packages/client-v2/src/net/protocol.rs
+5
View File
@@ -65,6 +65,11 @@ pub enum ControlPacket {
// ========== RPC ==========
RpcRequest {
id: u32,
/// 是否异步运行
run_async: bool,
/// 超时（毫秒）
timeout: Option<u64>,
/// 请求命令
command: Command,
},
RpcResponse {
packages/client-v2/src/net/rpc.rs
+181 -93
View File
@@ -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<CommandResult>,
created_at: std::time::Instant,
}
/// RPC 管理器
pub struct RpcManager {
handler: RpcHandler,
next_id: AtomicU32,
pending: Mutex<HashMap<u32, PendingRequest>>,
default_timeout: Duration,
pending: Arc<DashMap<u32, oneshot::Sender<CommandResult>>>,
}
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<CommandResult>) {
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<F, Fut>(
&self,
builder: &RpcBuilder,
send_fn: F,
) -> Result<CommandResult, RpcError>
where
F: FnOnce(ControlPacket) -> Fut,
Fut: Future<Output = Result<()>>,
{
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<F, Fut>(
&self,
id: u32,
run_async: bool,
timeout_ms: Option<u64>,
command: Command,
send_fn: F,
) -> Result<()>
where
F: Fn(ControlPacket) -> Fut + Send + Sync + 'static, // 异步执行需要 'static
Fut: Future<Output = Result<()>> + 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<u64>,
}
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<F, Fut>(
timeout: Duration,
register_fn: F,
) -> Result<CommandResult, RpcError>
where
F: FnOnce() -> Fut,
Fut: std::future::Future<Output = Result<oneshot::Receiver<CommandResult>, 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<u64>) -> 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<DashMap<u32, oneshot::Sender<CommandResult>>>,
}
impl Drop for DropGuard {
fn drop(&mut self) {
self.pending.remove(&self.id);
}
}
packages/client-v2/src/utils/mod.rs
+1
View File
@@ -0,0 +1 @@
pub mod shell;
packages/client-v2/src/utils/shell.rs
+68
View File
@@ -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<u64>) -> Result<ShellResponse> {
// 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
}
}
}