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

refactor: handlers

Browse Source
This commit is contained in:
Del Wang
2026-01-02 15:16:31 +08:00
parent 4ed8d3dca9
commit 6f55021651
13 changed files with 458 additions and 252 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/client-v2/Makefile
+5 -2
View File
@@ -1,13 +1,16 @@
build-server:
clear && cargo build --release --bin server --features app
cargo build --release --bin server --features app
build-client:
clear && docker run --rm -v $(shell pwd):/app idootop/open-xiaoai-runtime:oh2p \
docker run --rm -v $(shell pwd):/app idootop/open-xiaoai-runtime:oh2p \
cargo build --target armv7-unknown-linux-gnueabihf --release --bin client --features app
run-server:
target/release/server
test:
clear && make build-client && make deploy && make build-server && make run-server
# 部署到小爱音箱（调试自用）
deploy:
dd if=target/armv7-unknown-linux-gnueabihf/release/client \
packages/client-v2/src/app/client/core.rs
+63 -172
View File
@@ -1,12 +1,11 @@
#![cfg(target_os = "linux")]
use crate::audio::codec::OpusCodec;
use crate::app::client::handlers;
use crate::app::client::session::ClientSession;
use crate::audio::config::AudioConfig;
use crate::audio::player::AudioPlayer;
use crate::audio::recorder::AudioRecorder;
use crate::net::discovery::Discovery;
use crate::net::network::{AudioSocket, ClientNetwork, ControlConnection};
use crate::net::protocol::{AudioPacket, ControlPacket, DeviceInfo, RpcResult};
use crate::net::network::{AudioSocket, ClientNetwork};
use crate::net::protocol::{ControlPacket, DeviceInfo, RpcResult};
use anyhow::{Context, Result};
use std::net::SocketAddr;
use std::sync::Arc;
@@ -14,6 +13,7 @@ use tokio::sync::broadcast;
pub struct Client {
info: DeviceInfo,
session: Arc<tokio::sync::Mutex<Option<Arc<ClientSession>>>>,
}
impl Client {
@@ -24,10 +24,11 @@ impl Client {
mac: mac.to_string(),
version,
},
session: Arc::new(tokio::sync::Mutex::new(None)),
}
}
pub async fn run(self) -> Result<()> {
pub async fn run(self: Arc<Self>) -> Result<()> {
println!("正在寻找服务端...");
let (server_ip, server_port) = Discovery::discover_server().await?;
let server_addr = SocketAddr::new(server_ip, server_port);
@@ -48,177 +49,67 @@ impl Client {
p => return Err(anyhow::anyhow!("认证失败: {:?}", p)),
}
let (reader, writer) = control.split();
let session = Arc::new(ClientSession::new(self.info.clone(), writer));
*self.session.lock().await = Some(session.clone());
let (stop_tx, _) = broadcast::channel::<()>(1);
let audio_socket = Arc::new(AudioSocket::bind().await?);
// 启动测试 RPC 调用的任务
let client_clone = self.clone();
tokio::spawn(async move {
tokio::time::sleep(std::time::Duration::from_secs(2)).await;
println!("测试向服务端发起 RPC: status");
match client_clone.call_server("status", vec![]).await {
Ok(res) => println!("收到服务端响应: {:?}", res),
Err(e) => eprintln!("向服务端发起 RPC 失败: {}", e),
}
});
let mut reader = reader;
loop {
tokio::select! {
packet = control.recv_packet() => {
let packet = packet?;
match packet {
ControlPacket::StartRecording { config } => {
println!("开始录音: {:?}", config);
let socket = audio_socket.clone();
let mut stop_rx = stop_tx.subscribe();
let server_addr = control.peer_addr()?;
tokio::spawn(async move {
if let Err(e) = handle_recording(config, socket, server_addr, stop_rx).await {
eprintln!("录音出错: {}", e);
}
});
}
ControlPacket::StopRecording => {
println!("停止录音");
let _ = stop_tx.send(());
}
ControlPacket::StartPlayback { config } => {
println!("开始播放: {:?}", config);
let socket = audio_socket.clone();
let mut stop_rx = stop_tx.subscribe();
tokio::spawn(async move {
if let Err(e) = handle_playback(config, socket, stop_rx).await {
eprintln!("播放出错: {}", e);
}
});
}
ControlPacket::StopPlayback => {
println!("停止播放");
let _ = stop_tx.send(());
}
ControlPacket::RpcRequest { id, method, args } => {
println!("收到 RPC 请求: {} {:?}", method, args);
let result = handle_rpc(&method, args).await;
control.send_packet(&ControlPacket::RpcResponse { id, result }).await?;
}
ControlPacket::Ping => {
control.send_packet(&ControlPacket::Pong).await?;
}
_ => {}
}
let packet = reader.recv_packet().await?;
let self_clone = self.clone();
let session_clone = session.clone();
let audio_socket = audio_socket.clone();
let stop_tx = stop_tx.clone();
tokio::spawn(async move {
if let Err(e) = handlers::handle_packet(
self_clone,
session_clone,
packet,
audio_socket,
stop_tx,
server_addr,
)
.await
{
eprintln!("处理控制包出错: {}", e);
}
}
}
}
}
async fn handle_recording(
config: AudioConfig,
socket: Arc<AudioSocket>,
server_addr: SocketAddr,
mut stop_rx: broadcast::Receiver<()>,
) -> Result<()> {
#[cfg(target_os = "linux")]
{
let recorder = AudioRecorder::new(&config)?;
let mut codec = OpusCodec::new(&config)?;
let mut pcm_buf = vec![0i16; config.frame_size];
let mut opus_buf = vec![0u8; 4096];
loop {
if stop_rx.try_recv().is_ok() {
break;
}
let n = recorder.read(&mut pcm_buf)?;
if n > 0 {
let opus_len = codec.encode(&pcm_buf[..n], &mut opus_buf)?;
let packet = AudioPacket {
data: opus_buf[..opus_len].to_vec(),
};
socket.send_packet(&packet, server_addr).await?;
}
}
}
#[cfg(not(target_os = "linux"))]
{
println!("当前系统不支持 ALSA 录音，模拟发送音频数据...");
loop {
if stop_rx.try_recv().is_ok() {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(20)).await;
let packet = AudioPacket {
data: vec![0u8; 10],
};
socket.send_packet(&packet, server_addr).await?;
}
}
Ok(())
}
async fn handle_playback(
config: AudioConfig,
socket: Arc<AudioSocket>,
mut stop_rx: broadcast::Receiver<()>,
) -> Result<()> {
#[cfg(target_os = "linux")]
{
let player = AudioPlayer::new(&config)?;
let mut codec = OpusCodec::new(&config)?;
let mut pcm_buf = vec![0i16; config.frame_size];
let mut udp_buf = vec![0u8; 4096];
loop {
if stop_rx.try_recv().is_ok() {
break;
}
// 这里简单处理，UDP 接收可能阻塞。实际建议加超时或 select
let (packet, _) = socket.recv_packet(&mut udp_buf).await?;
let pcm_len = codec.decode(&packet.data, &mut pcm_buf)?;
player.write(&pcm_buf[..pcm_len])?;
}
}
#[cfg(not(target_os = "linux"))]
{
println!("当前系统不支持 ALSA 播放，模拟接收音频数据...");
let mut udp_buf = vec![0u8; 4096];
loop {
tokio::select! {
_ = stop_rx.recv() => break,
res = socket.recv_packet(&mut udp_buf) => {
let _ = res?;
}
}
}
}
Ok(())
}
async fn handle_rpc(method: &str, args: Vec<String>) -> RpcResult {
if method == "shell" && !args.is_empty() {
let cmd = &args[0];
println!("执行 Shell: {}", cmd);
// 模拟执行
#[cfg(target_os = "linux")]
{
use std::process::Command;
let output = Command::new("sh").arg("-c").arg(cmd).output();
match output {
Ok(out) => RpcResult {
stdout: String::from_utf8_lossy(&out.stdout).to_string(),
stderr: String::from_utf8_lossy(&out.stderr).to_string(),
code: out.status.code().unwrap_or(0),
},
Err(e) => RpcResult {
stdout: "".to_string(),
stderr: e.to_string(),
code: -1,
},
}
}
#[cfg(not(target_os = "linux"))]
{
RpcResult {
stdout: format!("模拟执行 {} 成功", cmd),
stderr: "".to_string(),
code: 0,
}
}
} else {
RpcResult {
stdout: "".to_string(),
stderr: "未知方法或参数错误".to_string(),
code: -1,
});
}
}
pub async fn call_server(&self, method: &str, args: Vec<String>) -> Result<RpcResult> {
let session = self
.session
.lock()
.await
.clone()
.context("Session not established")?;
let (id, rx) = session.rpc.alloc_id();
session
.writer
.lock()
.await
.send_packet(&ControlPacket::RpcRequest {
id,
method: method.to_string(),
args,
})
.await?;
Ok(rx.await?)
}
}
packages/client-v2/src/app/client/entry.rs
+1 -1
View File
@@ -9,6 +9,6 @@ pub async fn run_client() -> Result<()> {
let mac = "00:11:22:33:44:55";
let version = 1;
let client = Client::new(model, mac, version);
let client = std::sync::Arc::new(Client::new(model, mac, version));
client.run().await
}
packages/client-v2/src/app/client/handlers.rs
+202
View File
@@ -0,0 +1,202 @@
#![cfg(target_os = "linux")]
use crate::app::client::core::Client;
use crate::app::client::session::ClientSession;
use crate::audio::codec::OpusCodec;
use crate::audio::config::AudioConfig;
use crate::audio::player::AudioPlayer;
use crate::audio::recorder::AudioRecorder;
use crate::net::network::AudioSocket;
use crate::net::protocol::{AudioPacket, ControlPacket, RpcResult};
use anyhow::Result;
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::sync::broadcast;
pub async fn handle_packet(
client: Arc<Client>,
session: Arc<ClientSession>,
packet: ControlPacket,
audio_socket: Arc<AudioSocket>,
stop_tx: broadcast::Sender<()>,
server_addr: SocketAddr,
) -> Result<()> {
match packet {
ControlPacket::RpcResponse { id, result } => {
session.rpc.fulfill(id, result);
}
ControlPacket::RpcRequest { id, method, args } => {
println!("收到 RPC 请求: {} {:?}", method, args);
let result = handle_rpc(&method, args).await;
let response = ControlPacket::RpcResponse { id, result };
session.writer.lock().await.send_packet(&response).await?;
}
ControlPacket::StartRecording { config } => {
// ...
println!("开始录音: {:?}", config);
let mut stop_rx = stop_tx.subscribe();
tokio::spawn(async move {
if let Err(e) = handle_recording(config, audio_socket, server_addr, stop_rx).await {
eprintln!("录音出错: {}", e);
}
});
}
ControlPacket::StopRecording => {
println!("停止录音");
let _ = stop_tx.send(());
}
ControlPacket::StartPlayback { config } => {
println!("开始播放: {:?}", config);
let mut stop_rx = stop_tx.subscribe();
tokio::spawn(async move {
if let Err(e) = handle_playback(config, audio_socket, stop_rx).await {
eprintln!("播放出错: {}", e);
}
});
}
ControlPacket::StopPlayback => {
println!("停止播放");
let _ = stop_tx.send(());
}
ControlPacket::Ping => {
let _ = session
.writer
.lock()
.await
.send_packet(&ControlPacket::Pong)
.await;
}
_ => {}
}
Ok(())
}
pub async fn handle_rpc(method: &str, args: Vec<String>) -> RpcResult {
match method {
"shell" => handle_shell(args).await,
_ => RpcResult {
stdout: "".to_string(),
stderr: format!("Unknown method: {}", method),
code: -1,
},
}
}
async fn handle_shell(args: Vec<String>) -> RpcResult {
if args.is_empty() {
return RpcResult {
stdout: "".to_string(),
stderr: "Missing command argument".to_string(),
code: -1,
};
}
let cmd = &args[0];
println!("Executing Shell: {}", cmd);
#[cfg(target_os = "linux")]
{
use std::process::Command;
let output = Command::new("sh").arg("-c").arg(cmd).output();
match output {
Ok(out) => RpcResult {
stdout: String::from_utf8_lossy(&out.stdout).to_string(),
stderr: String::from_utf8_lossy(&out.stderr).to_string(),
code: out.status.code().unwrap_or(0),
},
Err(e) => RpcResult {
stdout: "".to_string(),
stderr: e.to_string(),
code: -1,
},
}
}
#[cfg(not(target_os = "linux"))]
{
RpcResult {
stdout: format!("Mock execution of {} successful", cmd),
stderr: "".to_string(),
code: 0,
}
}
}
async fn handle_recording(
config: AudioConfig,
socket: Arc<AudioSocket>,
server_addr: SocketAddr,
mut stop_rx: broadcast::Receiver<()>,
) -> Result<()> {
#[cfg(target_os = "linux")]
{
let recorder = AudioRecorder::new(&config)?;
let mut codec = OpusCodec::new(&config)?;
let mut pcm_buf = vec![0i16; config.frame_size];
let mut opus_buf = vec![0u8; 4096];
loop {
if stop_rx.try_recv().is_ok() {
break;
}
let n = recorder.read(&mut pcm_buf)?;
if n > 0 {
let opus_len = codec.encode(&pcm_buf[..n], &mut opus_buf)?;
let packet = AudioPacket {
data: opus_buf[..opus_len].to_vec(),
};
socket.send_packet(&packet, server_addr).await?;
}
}
}
#[cfg(not(target_os = "linux"))]
{
println!("当前系统不支持 ALSA 录音，模拟发送音频数据...");
loop {
if stop_rx.try_recv().is_ok() {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(20)).await;
let packet = AudioPacket {
data: vec![0u8; 10],
};
socket.send_packet(&packet, server_addr).await?;
}
}
Ok(())
}
async fn handle_playback(
config: AudioConfig,
socket: Arc<AudioSocket>,
mut stop_rx: broadcast::Receiver<()>,
) -> Result<()> {
#[cfg(target_os = "linux")]
{
let player = AudioPlayer::new(&config)?;
let mut codec = OpusCodec::new(&config)?;
let mut pcm_buf = vec![0i16; config.frame_size];
let mut udp_buf = vec![0u8; 4096];
loop {
if stop_rx.try_recv().is_ok() {
break;
}
let (packet, _) = socket.recv_packet(&mut udp_buf).await?;
let pcm_len = codec.decode(&packet.data, &mut pcm_buf)?;
player.write(&pcm_buf[..pcm_len])?;
}
}
#[cfg(not(target_os = "linux"))]
{
println!("当前系统不支持 ALSA 播放，模拟接收音频数据...");
let mut udp_buf = vec![0u8; 4096];
loop {
tokio::select! {
_ = stop_rx.recv() => break,
res = socket.recv_packet(&mut udp_buf) => {
let _ = res?;
}
}
}
}
Ok(())
}
packages/client-v2/src/app/client/mod.rs
+2
View File
@@ -1,2 +1,4 @@
pub mod core;
pub mod entry;
pub mod handlers;
pub mod session;
packages/client-v2/src/app/client/session.rs
+20
View File
@@ -0,0 +1,20 @@
use crate::net::network::ControlWriter;
use crate::net::protocol::DeviceInfo;
use crate::net::rpc::RpcManager;
use std::sync::Arc;
pub struct ClientSession {
pub info: DeviceInfo,
pub writer: Arc<tokio::sync::Mutex<ControlWriter>>,
pub rpc: Arc<RpcManager>,
}
impl ClientSession {
pub fn new(info: DeviceInfo, writer: ControlWriter) -> Self {
Self {
info,
writer: Arc::new(tokio::sync::Mutex::new(writer)),
rpc: Arc::new(RpcManager::new()),
}
}
}
packages/client-v2/src/app/server/core.rs
+17 -70
View File
@@ -1,53 +1,19 @@
use crate::app::server::handlers;
use crate::app::server::session::ServerSession;
use crate::audio::codec::OpusCodec;
use crate::audio::config::AudioConfig;
use crate::audio::wav::{WavReader, WavWriter};
use crate::net::discovery::Discovery;
use crate::net::network::{AudioSocket, ControlConnection, ServerNetwork};
use crate::net::protocol::{AudioPacket, ControlPacket, DeviceInfo, RpcResult};
use crate::net::protocol::{AudioPacket, ControlPacket, RpcResult};
use anyhow::{Context, Result};
use parking_lot::Mutex;
use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::Arc;
use std::sync::atomic::{AtomicU32, Ordering};
use tokio::sync::{broadcast, oneshot};
pub struct RpcManager {
next_id: AtomicU32,
pending: Mutex<HashMap<u32, oneshot::Sender<RpcResult>>>,
}
impl RpcManager {
pub fn new() -> Self {
Self {
next_id: AtomicU32::new(1),
pending: Mutex::new(HashMap::new()),
}
}
pub fn alloc_id(&self) -> (u32, oneshot::Receiver<RpcResult>) {
let id = self.next_id.fetch_add(1, Ordering::SeqCst);
let (tx, rx) = oneshot::channel();
self.pending.lock().insert(id, tx);
(id, rx)
}
pub fn fulfill(&self, id: u32, result: RpcResult) {
if let Some(tx) = self.pending.lock().remove(&id) {
let _ = tx.send(result);
}
}
}
pub struct Session {
pub info: DeviceInfo,
pub control: Arc<tokio::sync::Mutex<ControlConnection>>,
pub addr: SocketAddr,
pub rpc: Arc<RpcManager>,
}
pub struct Server {
sessions: Arc<Mutex<HashMap<SocketAddr, Arc<Session>>>>,
sessions: Arc<Mutex<HashMap<SocketAddr, Arc<ServerSession>>>>,
audio_socket: Arc<AudioSocket>,
}
@@ -98,36 +64,20 @@ impl Server {
);
control.send_packet(&ControlPacket::IdentifyOk).await?;
let rpc = Arc::new(RpcManager::new());
let control = Arc::new(tokio::sync::Mutex::new(control));
let session = Arc::new(Session {
info,
control: control.clone(),
addr,
rpc: rpc.clone(),
});
let (mut reader, writer) = control.split();
let session = Arc::new(ServerSession::new(info, writer, addr));
self.sessions.lock().insert(addr, session.clone());
// 处理控制消息循环
loop {
let mut ctrl = control.lock().await;
let packet = ctrl.recv_packet().await?;
drop(ctrl); // 释放锁以允许发送
match packet {
ControlPacket::RpcResponse { id, result } => {
rpc.fulfill(id, result);
let packet = reader.recv_packet().await?;
let session = session.clone();
tokio::spawn(async move {
if let Err(e) = handlers::handle_packet(session, packet).await {
eprintln!("处理来自 {} 的包时出错: {}", addr, e);
}
ControlPacket::Ping => {
control
.lock()
.await
.send_packet(&ControlPacket::Pong)
.await?;
}
_ => {}
}
});
}
}
@@ -141,7 +91,7 @@ impl Server {
.context("未找到 Session")?;
let (id, rx) = session.rpc.alloc_id();
session
.control
.writer
.lock()
.await
.send_packet(&ControlPacket::RpcRequest {
@@ -163,7 +113,7 @@ impl Server {
// 发送开始录音指令
session
.control
.writer
.lock()
.await
.send_packet(&ControlPacket::StartRecording {
@@ -189,7 +139,7 @@ impl Server {
.cloned()
.context("未找到 Session")?;
session
.control
.writer
.lock()
.await
.send_packet(&ControlPacket::StopRecording)
@@ -206,7 +156,7 @@ impl Server {
.context("未找到 Session")?;
session
.control
.writer
.lock()
.await
.send_packet(&ControlPacket::StartPlayback {
@@ -233,7 +183,7 @@ impl Server {
.cloned()
.context("未找到 Session")?;
session
.control
.writer
.lock()
.await
.send_packet(&ControlPacket::StopPlayback)
@@ -259,8 +209,6 @@ async fn save_audio_to_wav(
println!("正在录制到 {}...", path);
// 这里需要一个停止机制，目前简单起见，如果一段时间没收到包就停止，或者通过全局状态
// 为简单演示，我们录制 100 个包
for _ in 0..100 {
let (packet, _) = socket.recv_packet(&mut udp_buf).await?;
let pcm_len = codec.decode(&packet.data, &mut pcm_buf)?;
@@ -296,7 +244,6 @@ async fn stream_wav_to_client(
};
socket.send_packet(&packet, target).await?;
// 控制发送频率，约 20ms 一帧
tokio::time::sleep(std::time::Duration::from_millis(20)).await;
}
packages/client-v2/src/app/server/handlers.rs
+50
View File
@@ -0,0 +1,50 @@
use crate::app::server::session::ServerSession;
use crate::net::protocol::{ControlPacket, RpcResult};
use anyhow::Result;
use std::sync::Arc;
pub async fn handle_packet(session: Arc<ServerSession>, packet: ControlPacket) -> Result<()> {
match packet {
ControlPacket::RpcResponse { id, result } => {
session.rpc.fulfill(id, result);
}
ControlPacket::RpcRequest { id, method, args } => {
println!(
"收到来自客户端 {} 的 RPC 请求: {} {:?}",
session.addr, method, args
);
let result = handle_server_rpc(&method, args).await;
session
.writer
.lock()
.await
.send_packet(&ControlPacket::RpcResponse { id, result })
.await?;
}
ControlPacket::Ping => {
session
.writer
.lock()
.await
.send_packet(&ControlPacket::Pong)
.await?;
}
_ => {}
}
Ok(())
}
async fn handle_server_rpc(method: &str, _args: Vec<String>) -> RpcResult {
match method {
"status" => RpcResult {
stdout: "Server is running normally".to_string(),
stderr: "".to_string(),
code: 0,
},
_ => RpcResult {
stdout: "".to_string(),
stderr: format!("Server does not support method: {}", method),
code: -1,
},
}
}
packages/client-v2/src/app/server/mod.rs
+2
View File
@@ -1,2 +1,4 @@
pub mod core;
pub mod entry;
pub mod handlers;
pub mod session;
packages/client-v2/src/app/server/session.rs
+23
View File
@@ -0,0 +1,23 @@
use crate::net::network::ControlWriter;
use crate::net::protocol::DeviceInfo;
use crate::net::rpc::RpcManager;
use std::net::SocketAddr;
use std::sync::Arc;
pub struct ServerSession {
pub info: DeviceInfo,
pub writer: Arc<tokio::sync::Mutex<ControlWriter>>,
pub addr: SocketAddr,
pub rpc: Arc<RpcManager>,
}
impl ServerSession {
pub fn new(info: DeviceInfo, writer: ControlWriter, addr: SocketAddr) -> Self {
Self {
info,
writer: Arc::new(tokio::sync::Mutex::new(writer)),
addr,
rpc: Arc::new(RpcManager::new()),
}
}
}
packages/client-v2/src/net/mod.rs
+1
View File
@@ -1,3 +1,4 @@
pub mod discovery;
pub mod network;
pub mod protocol;
pub mod rpc;
packages/client-v2/src/net/network.rs
+40 -7
View File
@@ -3,9 +3,11 @@ use anyhow::{Context, Result};
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf};
use tokio::net::{TcpListener, TcpStream, UdpSocket};
/// UDP 音频传输
// ... (rest of AudioSocket is same)
pub struct AudioSocket {
socket: Arc<UdpSocket>,
}
@@ -39,6 +41,40 @@ impl AudioSocket {
}
}
/// TCP 控制连接读取端
pub struct ControlReader {
reader: OwnedReadHalf,
}
impl ControlReader {
pub async fn recv_packet(&mut self) -> Result<ControlPacket> {
let len = self.reader.read_u32().await? as usize;
if len > 10 * 1024 * 1024 {
return Err(anyhow::anyhow!("Packet too large: {}", len));
}
let mut buf = vec![0u8; len];
self.reader.read_exact(&mut buf).await?;
let packet = postcard::from_bytes(&buf)?;
Ok(packet)
}
}
/// TCP 控制连接写入端
pub struct ControlWriter {
writer: OwnedWriteHalf,
}
impl ControlWriter {
pub async fn send_packet(&mut self, packet: &ControlPacket) -> Result<()> {
let bytes = postcard::to_allocvec(packet)?;
let len = bytes.len() as u32;
self.writer.write_u32(len).await?;
self.writer.write_all(&bytes).await?;
self.writer.flush().await?;
Ok(())
}
}
/// TCP 控制连接
pub struct ControlConnection {
stream: TcpStream,
@@ -54,6 +90,7 @@ impl ControlConnection {
let len = bytes.len() as u32;
self.stream.write_u32(len).await?;
self.stream.write_all(&bytes).await?;
self.stream.flush().await?;
Ok(())
}
@@ -68,13 +105,9 @@ impl ControlConnection {
Ok(packet)
}
pub fn split(
self,
) -> (
tokio::net::tcp::OwnedReadHalf,
tokio::net::tcp::OwnedWriteHalf,
) {
self.stream.into_split()
pub fn split(self) -> (ControlReader, ControlWriter) {
let (r, w) = self.stream.into_split();
(ControlReader { reader: r }, ControlWriter { writer: w })
}
pub fn peer_addr(&self) -> Result<SocketAddr> {
packages/client-v2/src/net/rpc.rs
+32
View File
@@ -0,0 +1,32 @@
use crate::net::protocol::RpcResult;
use parking_lot::Mutex;
use std::collections::HashMap;
use std::sync::atomic::{AtomicU32, Ordering};
use tokio::sync::oneshot;
pub struct RpcManager {
next_id: AtomicU32,
pending: Mutex<HashMap<u32, oneshot::Sender<RpcResult>>>,
}
impl RpcManager {
pub fn new() -> Self {
Self {
next_id: AtomicU32::new(1),
pending: Mutex::new(HashMap::new()),
}
}
pub fn alloc_id(&self) -> (u32, oneshot::Receiver<RpcResult>) {
let id = self.next_id.fetch_add(1, Ordering::SeqCst);
let (tx, rx) = oneshot::channel();
self.pending.lock().insert(id, tx);
(id, rx)
}
pub fn fulfill(&self, id: u32, result: RpcResult) {
if let Some(tx) = self.pending.lock().remove(&id) {
let _ = tx.send(result);
}
}
}