rust tokio select!宏详解
简介
本文介绍Tokio
中select!
的用法,重点是使用过程中可能遇到的问题,比如阻塞问题、优先级问题、cancel safe
问题。在Tokio
中,select!
是一个宏,用于同时等待多个异步任务,并在其中任意一个任务完成时执行相应的逻辑。
基本用法
如下代码演示了如何使用 Tokio
库实现一个异步的消息传递系统,其中包括三个无限通道和一个关闭通道。程序使用了 select!
宏来等待通道和关闭通道的事件,并在事件发生时执行相应的操作。
程序的主要步骤如下:
- 创建三个无限通道和一个用于传递关闭信号的通道。
- 向三个通道中发送一些数据。
- 开启一个异步任务并在两秒后发送关闭信号。
- 在主循环中使用
select!
宏等待通道和关闭通道的事件。 - 当一个通道接收到数据时,打印出数据。
- 当关闭通道接收到信号时,退出循环。
程序中的 select!
宏使用了类似于 match
的语法,但是它可以同时等待多个异步事件。当其中一个事件发生时,宏将执行相应的代码块,并跳出循环。在本例中,当一个通道接收到数据时,打印出数据;当关闭通道接收到信号时,退出循环。
select!
经常与loop
搭配使用,循环地从多个通道中接收事件并处理。
use std::time::Duration;
use tokio::select;
#[tokio::main]
async fn main() {
let (sender1, mut receiver1) = tokio::sync::mpsc::unbounded_channel::<String>();
let (sender2, mut receiver2) = tokio::sync::mpsc::unbounded_channel::<String>();
let (sender3, mut receiver3) = tokio::sync::mpsc::unbounded_channel::<String>();
let (shutdown_sender, mut shutdown_receiver) = tokio::sync::watch::channel(());
for i in 0..3 {
sender1.send(i.to_string()).unwrap();
sender2.send(i.to_string()).unwrap();
sender3.send(i.to_string()).unwrap();
}
tokio::spawn(async move {
tokio::time::sleep(Duration::from_secs(2)).await;
shutdown_sender.send(()).unwrap(); //两秒后关闭
});
loop {
select! {
ret = receiver1.recv() => {
println!("channel 1 received: {:?}", ret);
},
ret = receiver2.recv() => {
println!("channel 2 received: {:?}", ret);
},
ret = receiver3.recv() => {
println!("channel 3 received: {:?}", ret);
},
_ = shutdown_receiver.changed() => {
println!("shutdown received");
break;
}
};
}
}
可能遇到的坑
阻塞
select
中的各个分支是并行执行的,这里的并行是指分支中的各个future
在并行执行。不过一旦某个分支的future
完成并进入了分支代码块,如果在分支代码中有一些阻塞的操作,则其他分支是没有机会执行的。
比如下面代码,在receiver1.recv()
完成时,sleep
了10s,sleep
期间其他的分支是不会执行的。即使在2s
后发送了shutdown
信号,select!
因为无法及时处理此信号,实际上循环也无法退出。
loop {
select! {
ret = receiver1.recv() => {
println!("channel 1 received: {:?}", ret);
tokio::time::sleep(Duration::from_secs(10)).await;//这里等待期间,其他的分支是无法被执行的
},
ret = receiver2.recv() => {
println!("channel 2 received: {:?}", ret);
},
ret = receiver3.recv() => {
println!("channel 3 received: {:?}", ret);
},
_ = shutdown_receiver.changed() => {
println!("shutdown received");
break;
}
};
}
这个坑在网络编程中比较容易踩到,比如select
这里是从channel
中取出上层应用传来的数据,并将其写入到socket
中,而写socket
的操作是有可能阻塞的,阻塞期间其他的分支是无法执行的。
顺序
1、默认情况下select
中的各个分支执行顺序是随机的,比如上面例子中三个channel
都有消息的情况下,具体去执行哪个分支是随机的。执行结果如下:
2、如果想要区分优先级,可以加标志biased
,这样每次select
将会按照从上到下的顺序去poll
每个future
,也就是说优先级顺序是从上往下的。比如某些场景下需要按优先级处理各个channel
中的数据时这个特性就很有用。代码如下:
loop {
select! {
biased;//按顺序优先执行
ret = receiver1.recv() => {
println!("channel 1 received: {:?}", ret);
},
ret = receiver2.recv() => {
println!("channel 2 received: {:?}", ret);
},
ret = receiver3.recv() => {
println!("channel 3 received: {:?}", ret);
},
_ = shutdown_receiver.changed() => {
println!("shutdown received");
break;
}
};
}
运行结果如下:
3、顺序执行时注意饿死问题
添加了biased
标志后,顺序靠前的future
总是先被执行,在上述例子中,极端情况下如果靠前的channel
总是有数据,那后面的channel
就没有机会被执行。比如例子中如果前三个channel
中一直有数据,那shutdown_receiver
就无法收到shutdown
信号,导致程序功能不符合预期。
解决这个问题很简单,就是把更关键的控制性的future
放在最前方。
关于cancel safe
select!
中如果某个分支future completed
了,会将其他分支的future cancel
掉,这个cancel
操作要格外小心,因为如果future
不是cancel safe
的可能会丢数据。tokio的官方文档中给出了常见的cancel safe
和不safe
的future
。
那么如何判断自己实现的future
是否是cancel safe
的呢? 很简单、只需要思考如果future
中的代码执行到.await
时被cancel
了,是否是安全的。我们来看下cancel unsafe
的代码长啥样:
pub async fn read_and_write(mut message_recevier: UnboundedReceiver<Bytes>, mut file: File) {
let message = message_recevier.recv().await.unwrap();
file.write(&message).await.unwrap();
}
该方法从一个channel
中读取消息,并将此消息写入到文件中,这个future
就明显不是cancel safe
的。为啥呢?试想一下,此future
从channel
中读到消息之后,在写文件时被cancel
掉了,那message
岂不是就丢了。
实际项目中一定要格外小心这个cancel safe
问题,很容易造成丢数据或者数据重复等不良反应,而且一旦出现了还很难复现、不太容易想到是这里的问题。网络编程中尤其要注意tokio::io::AsyncWriteExt::write_all
不是cancel safe
的,因为它内部可能是多次调用write
操作才将所有缓冲区写入。
数量
1、首先select!
中的分支仅支持显式地用代码书写,无法动态增减。就是说在写代码时select
中的futures
数量就固定了,程序运行过程中无法动态删减。
2、目前最多支持64个分支。