rust tokio select!宏详解

rust tokio select!宏详解

简介

本文介绍Tokio中select!的用法，重点是使用过程中可能遇到的问题，比如阻塞问题、优先级问题、cancel safe问题。在Tokio 中，select! 是一个宏，用于同时等待多个异步任务，并在其中任意一个任务完成时执行相应的逻辑。

基本用法

如下代码演示了如何使用 Tokio 库实现一个异步的消息传递系统，其中包括三个无限通道和一个关闭通道。程序使用了 select! 宏来等待通道和关闭通道的事件，并在事件发生时执行相应的操作。

程序的主要步骤如下：

1. 创建三个无限通道和一个用于传递关闭信号的通道。
2. 向三个通道中发送一些数据。
3. 开启一个异步任务并在两秒后发送关闭信号。
4. 在主循环中使用 select! 宏等待通道和关闭通道的事件。
5. 当一个通道接收到数据时，打印出数据。
6. 当关闭通道接收到信号时，退出循环。

程序中的 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个分支。