在现代的异步编程中，Tokio是一个强大的Rust异步运行时，它提供了构建高性能网络应用的能力。在这篇文章中，我们将探讨如何使用Tokio库中的 TcpStream 来建立TCP连接、收发数据，并实现心跳机制以保持连接的活跃状态。
建立Tcp连接
tokio 中的TcpStream的 connect提供了建立连接的功能
pub async fn connect<A: ToSocketAddrs>(addr: A) -> Result<TcpStream>
发送和接收数据
tokio 中的 TcpStream 的 into_split 接口可将TcpStream 分成读写两半，可以并发地读取和写入流：
pub fn into_split(self) -> (OwnedReadHalf, OwnedWriteHalf)Splits a TcpStream into a read half and a write half, which can be used to read and write the stream concurrently.
实现心跳机制
在长连接中，保持连接的活跃性是至关重要的。需要通过心跳机制来维持 TCP 连接的活跃性。在一定时间内没有收发消息，需要发送心跳包给服务器维持连接。
代码实现
此处只把tokio的依赖列出，Cargo.toml 文件中添加以下依赖项：
tokio = { version = "1.39.2", features = ["rt", "rt-multi-thread", "macros", "net", "time", "sync", "io-util"] }
下面代码包含TcpClient简单封装，测试代码。还包含数据收发、心跳维持的两种写法，具体如下
use std::string::ToString;
use std::sync::OnceLock;
use anyhow::anyhow;
use prost::bytes::BytesMut;
use tokio::io;
use tokio::net::TcpStream;
use tokio::runtime::{Builder, Runtime};
use tokio::sync::mpsc;
use tokio::sync::mpsc::Sender;
use tokio::time::{Duration, timeout};
// 简单封装，仅做示意
struct TcpClient {
rt: Runtime,
addr: String,
sender: OnceLock<Sender<Vec<u8>>>,
}
impl TcpClient {
fn new(addr: String) -> anyhow::Result<Self> {
let rt = Builder::new_multi_thread()
.enable_io()
.enable_time()
.worker_threads(2)
.thread_name("thread_name")
.build()?;
anyhow::Ok(TcpClient {
rt,
addr,
sender: OnceLock::new(),
})
}
// 建立连接，收发数据，维持心跳
fn connect(&mut self) -> anyhow::Result<()> {
let addr = self.addr.clone();
let (tx, mut rx) = mpsc::channel(100);
let tx_heart = tx.clone();
self.sender
.set(tx)
.map_err(|e| anyhow!("already connect"))?;
self.rt.spawn(async move {
let stream = TcpStream::connect(addr).await.map_err(|e| anyhow!(format!("err {}", e.to_string())))?;
let (read_half, write_half) = stream.into_split();
let (heart_sender, mut heart_receiver) = mpsc::channel(10);
let heart_sender_read = heart_sender.clone();
let heart_sender_write = heart_sender.clone();
tokio::select! {
_ = async {
let mut read_form_net = BytesMut::with_capacity(1024 * 4);
loop {
match read_half.readable().await {
// if the readiness event is a false positive.
Ok(_) => {
match read_half.try_read_buf(&mut read_form_net) {
Ok(0) => {
return anyhow!("no data");
}
Ok(read_bytes_size) => {
// todo tcp 拆包逻辑
// 有数据读取，心跳重新计时(此时不方便区分去读的数据是否是心跳)
if heart_sender_read.capacity() > 0 {
heart_sender_read.send(true).await;
}
}
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
continue;
}
Err(e) => {
return anyhow!(format!("error_message {:?}", e));
}
}
}
Err(e) => {
return anyhow!(format!("error_message {:?}", e));
}
}
}
} => {
println!("read task completed");
}
_ = async {
loop {
match write_half.writable().await {
Ok(_) => {
// Try to write data, this may still fail with `WouldBlock`
// if the readiness event is a false positive.
let write_data = match rx.recv().await {
Some(data) => {
println!("rx data is {:?}", data);
if data != HEART_PACK_DATA {
// 数据处理
}
data
}
None => { vec![] }
};
if write_data.is_empty() {
continue;
}
match write_half.try_write(write_data.as_slice()) {
Ok(n) => {
if write_data != HEART_PACK_DATA {
// 有非心跳数据写入，心跳重新计时
println!("hear send write data {:?}", write_data);
if heart_sender_write.capacity() > 0 {
heart_sender_write.send(true).await;
}
}
}
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
continue;
}
Err(e) => {
return Err::<(), std::io::Error>(io::Error::new(io::ErrorKind::Other, "error_message"));
}
}
}
Err(e) => {
return Err::<(), std::io::Error>(io::Error::new(io::ErrorKind::Other, "error_message"));
}
}
}
} => {
println!("write task completed");
}
_ = async {
// 0x89F
loop {
let res = timeout(Duration::from_secs(HEART_KEEPALIVE_TIME), heart_receiver.recv()).await;
match res {
Ok(_) => {
println!("socket data not ping");
continue;
},
Err(_) => {
// 发送心跳数据
println!("socket ping");
tx_heart.send(HEART_PACK_DATA.to_vec()).await;
},
}
}
} => {
println!("heart task completed");
}
}
anyhow::Ok(())
});
anyhow::Ok(())
}
// 发送数据
fn send_data(&self, data: Vec<u8>) -> anyhow::Result<()> {
let sender = self.sender.get().ok_or(anyhow!("not connect"))?;
let sender_clone = sender.clone();
self.rt.spawn(async move { sender_clone.send(data).await });
anyhow::Ok(())
}
}
const HEART_PACK_DATA: [u8; 2] = [0x10, 0x10];
const HEART_KEEPALIVE_TIME: u64 = 30;
// 测试代码
#[tokio::main]
async fn main() -> anyhow::Result<()> {
let data = vec![
0x31, 0x31, 0x3A, 0x7C, 0x0B, 0x00, 0x00, 0x0A, 0x00, 0x08, 0x00, 0x53, 0x48, 0x36, 0x30,
0x31, 0x32, 0x31, 0x31, 0x7D,
];
let mut tcp_client = TcpClient::new("127.0.0.1:3000".to_string())?;
tcp_client.connect();
loop {
tcp_client.send_data(data.clone());
tokio::time::sleep(Duration::from_secs(16)).await;
}
Ok(())
}
上述代码connect实现依赖tokio::select!，其内部启动三个async task
- 读task：从socket loop读数据
- 写task：往socket loop写数据
- 心跳task：如果超过心跳时间
timeout没有被唤醒，即heart_receiver心跳时间内没收到数据(heart_sender会在有数据读写时往heart_receiver发数据进行唤醒)，此时往socket 写入心跳包。
数据收发维持心跳的另一种写法
另外对于connect再提供另一种写法，依然依赖tokio::select!，内部启动三个task。但是loop循环稍有不同。tokio::select!在外部的loop内，每个task不再loop。此时心跳task和读写task不再交互。而是依赖一个
let mut keepalive_timeout: Option<Pin<Box<Sleep>>> = Some(Box::pin(time::sleep(Duration::from_secs(HEART_KEEPALIVE_TIME))));
keepalive_timeout sleep时间重置代码
let timeout = keepalive_timeout.as_mut().unwrap();
timeout.as_mut().reset(Instant::now() + Duration::from_secs(HEART_KEEPALIVE_TIME));
每次loop时通过上面代码把keepalive_timeout的sleep时间重置，如果心跳时间内没有数据读写，心跳task先执行，其他task取消，重新下一次loop (此写法参考rumqttc)内部的select方法实现
具体实现如下
fn connect(&mut self) -> anyhow::Result<()> {
let addr = self.addr.clone();
let (tx, mut rx) = mpsc::channel(100);
let tx_heart = tx.clone();
self.sender
.set(tx)
.map_err(|e| anyhow!("already connect"))?;
self.rt.spawn(async move {
let stream = TcpStream::connect(addr).await.map_err(|e| anyhow!(format!("err {}", e.to_string())))?;
let (read_half, write_half) = stream.into_split();
let mut keepalive_timeout: Option<Pin<Box<Sleep>>> = Some(Box::pin(time::sleep(Duration::from_secs(HEART_KEEPALIVE_TIME))));
'outer_loop: loop {
let timeout = keepalive_timeout.as_mut().unwrap();
timeout.as_mut().reset(Instant::now() + Duration::from_secs(HEART_KEEPALIVE_TIME));
tokio::select! {
o = async {
let mut read_form_net = BytesMut::with_capacity(1024 * 4);
match read_half.readable().await {
// if the readiness event is a false positive.
Ok(_) => {
match read_half.try_read_buf(&mut read_form_net) {
Ok(0) => {
return Err(anyhow!("no data"));
}
Ok(read_bytes_size) => {
// todo tcp 拆包逻辑，可以使用BytesMut拆包
}
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
}
Err(e) => {
return Err(anyhow!(format!("error_message {:?}", e)));
}
}
return anyhow::Ok(())
}
Err(e) => {
return Err(anyhow!(format!("error_message {:?}", e)));
}
};
} => {
// 读数据错误断开连接，重现开始. connect 应该传入 listener，此处断开回调listener
println!("read task completed {:?}", o);
if o.is_err() { break 'outer_loop }
}
o = async {
match write_half.writable().await {
Ok(_) => {
let write_data = match rx.recv().await {
Some(data) => {
println!("rx data is {:?}", data);
if data != HEART_PACK_DATA {
// 数据处理
}
data
}
None => { vec![] }
};
if !write_data.is_empty() {
match write_half.try_write(write_data.as_slice()) {
Ok(n) => {}
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
}
Err(e) => {
return Err(anyhow!(format!("error_message {:?}", e)));
}
}
}
return anyhow::Ok(())
}
Err(e) => {
return Err(anyhow!(format!("error_message {:?}", e)));
}
};
} => {
println!("write task completed {:?}", o);
if o.is_err() {
// 写数据错误断开连接，重现开始. connect 应该传入 listener，此处断开回调listener
break 'outer_loop
}
}
_ = keepalive_timeout.as_mut().unwrap() => {
// HEART_KEEPALIVE_TIME 时间内没有读写，发送心跳
println!("socket ping");
if let Err(r) = tx_heart.send(HEART_PACK_DATA.to_vec()).await {
break 'outer_loop
}
let timeout = keepalive_timeout.as_mut().unwrap();
timeout.as_mut().reset(Instant::now() + Duration::from_secs(HEART_KEEPALIVE_TIME));
}
}
}
anyhow::Ok(())
});
anyhow::Ok(())
}
