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操作才将所有缓冲区写入。