寻求Tokio `select!`宏的动态替代方案——多动态Channel监听

问题 #

我开发的应用需要同时监听多个消息通道。若通道数量是硬编码固定的，使用Tokio官方示例中的select!宏即可轻松实现：

use tokio::sync::mpsc;

#[tokio::main]
async fn main() {
    let (tx1, mut rx1) = mpsc::channel(128);
    let (tx2, mut rx2) = mpsc::channel(128);
    let (tx3, mut rx3) = mpsc::channel(128);

    loop {
        let msg = tokio::select! {
            Some(msg) = rx1.recv() => msg,
            Some(msg) = rx2.recv() => msg,
            Some(msg) = rx3.recv() => msg,
            else => { break }
        };

        println!("Got {:?}", msg);
    }

    println!("All channels have been closed.");
}

但我的场景中，通道以动态数量存储在Vec中，无法使用上述硬编码方式。我需要类似如下的实现：

use tokio::sync::mpsc;

#[tokio::main]
async fn main() {
    let channels = get_channels(); // Returns Vec<mpsc::Receiver<_>>

    while let Some(msg) = magic_crate::select_dynamic(&channels.iter()).await {
        println!("Got {:?}", msg);
    }

    println!("All channels closed");
}

我认为futures::select_all无法满足需求，因为该函数除了返回首个完成的Future结果外，还需要像tokio::select!那样取消其他Future，以确保其他通道的消息不会丢失。请问是否存在可行方案？我对futures::select_all的理解是否有误？

解决方案 #

纠正对futures::select_all的误解 #

你对futures::select_all的理解有误：该函数在某个Future完成时，会自动取消所有未完成的Future，同时返回完成的结果以及剩余的Future集合。你可以将剩余Future重新传入select_all循环监听，完全不会丢失其他通道的消息。

基于futures::select_all的实现 #

通过循环配合select_all，可以实现动态数量通道的监听，示例代码如下：

use futures::future::select_all;
use tokio::sync::mpsc;

#[tokio::main]
async fn main() {
    // 将每个通道的recv() Future包装为统一类型，放入集合
    let mut futures = get_channels()
        .into_iter()
        .map(|mut rx| Box::pin(rx.recv()) as _)
        .collect::<Vec<_>>();

    while !futures.is_empty() {
        let (result, _, remaining) = select_all(futures).await;
        
        match result {
            Some(msg) => println!("Got {:?}", msg),
            None => println!("A channel closed"),
        }

        // 剩余Future重新赋值，继续监听未关闭的通道
        futures = remaining;
    }

    println!("All channels closed");
}

更简洁的Stream合并方案 #

如果觉得手动管理Future集合麻烦，可以使用tokio-util的MergeQueue，将多个mpsc::Receiver合并为一个Stream，直接迭代处理消息：

use tokio::sync::mpsc;
use tokio_util::sync::MergeQueue;
use futures::StreamExt;

#[tokio::main]
async fn main() {
    let channels = get_channels();
    
    let mut merge_queue = MergeQueue::new();
    // 将每个通道的消息流加入合并队列
    for mut rx in channels {
        merge_queue.push(async move {
            while let Some(msg) = rx.recv().await {
                msg
            }
        });
    }

    // 迭代合并后的Stream，处理所有通道的消息
    while let Some(msg) = merge_queue.next().await {
        println!("Got {:?}", msg);
    }

    println!("All channels closed");
}

MergeQueue会自动管理所有输入的消息流，当任意通道有消息时立即返回，处理完成后继续监听剩余通道，完全符合动态监听的需求。

总结 #

• futures::select_all完全可以满足你的需求，它的行为和tokio::select!一致，会取消未完成的Future，且支持循环复用剩余Future实现动态监听。
• 追求代码简洁性的话，tokio-util的MergeQueue是更优选择，无需手动维护Future集合。

内容的提问来源于stack exchange，提问作者Jake