Rust中等待Tokio异步任务完成的问题求助

我正在编写一个程序，要求当MPSC通道收到消息时启动新的Tokio异步任务。目前程序能成功启动任务，但无法保持运行直到所有任务完成。

我的代码如下：

#[tokio::main]
async fn main() {
    let (sender, mut receiver) = mpsc::channel(100);

    let handles = Arc::new(tokio::sync::Mutex::new(vec![]));
    let handles_clone = handles.clone();
    let consumer_task = tokio::spawn(async move {
        while let Some(item) = receiver.recv().await {
            println!("Consumer received: {}", item);
            let h = tokio::spawn(async move {
                println!("Thread {} working", item);
                sleep(Duration::from_secs(2)).await;
                println!("Thread {} done", item);
            });
            {
                let mut handles = handles_clone.lock().await;
                handles.push(h);
            }
        }
    });

    let producer_task = tokio::spawn(async move {
        for i in 0..5 {
            let _ = sender.send(i).await;
        }
    });

    let mut handles = handles.lock().await;
    handles.push(consumer_task);
    handles.push(producer_task);

    let a = handles.deref();

    join_all(a);
}

我的思路是需要对启动的所有任务调用join_all，但无法同时在接收线程中向向量添加任务句柄，以及在程序末尾调用join_all阻塞至所有任务完成。我使用Arc<Mutex<_>>包装任务句柄向量，以便在接收线程中引用，但调用join_all时出现如下错误：

`&tokio::task::JoinHandle<()>` is not a future
the trait `futures::Future` is not implemented for `&tokio::task::JoinHandle<()>`
&tokio::task::JoinHandle<()> must be a future or must implement `IntoFuture` to be awaited
the trait `futures::Future` is implemented for `tokio::task::JoinHandle<T>`
`futures::Future` is implemented for `&mut tokio::task::JoinHandle<()>`, but not for `&tokio::task::JoinHandle<()>`

问题根源 #

1. join_all需要拥有所有权的JoinHandle，而你传递的是共享引用（&JoinHandle），仅JoinHandle本身或可变引用实现了Future trait，共享引用不满足要求。
2. 直接持有Arc<Mutex<Vec<JoinHandle>>>无法直接提取所有权，锁持有期间向量无法被移动。

修正步骤 #

1. 先等待生产者和消费者完成：生产者发送完所有消息后关闭通道，消费者会在通道关闭后退出循环，确保所有子任务都已被添加到句柄向量。
2. 提取任务句柄所有权：等消费者任务结束后，锁定Arc<Mutex<Vec<JoinHandle>>>，将向量内容转移到本地变量，释放锁后再调用join_all。

修正后的代码 #

use tokio::{sync::mpsc, task::JoinHandle, time::{sleep, Duration}};
use std::sync::Arc;
use futures::future::join_all;

#[tokio::main]
async fn main() {
    let (sender, mut receiver) = mpsc::channel(100);

    let handles = Arc::new(tokio::sync::Mutex::new(vec![]));
    let handles_clone = handles.clone();
    
    // 消费者任务：接收消息并启动子任务
    let consumer_task = tokio::spawn(async move {
        while let Some(item) = receiver.recv().await {
            println!("Consumer received: {}", item);
            let h = tokio::spawn(async move {
                println!("Task {} working", item);
                sleep(Duration::from_secs(2)).await;
                println!("Task {} done", item);
            });
            handles_clone.lock().await.push(h);
        }
        println!("Consumer task finished");
    });

    // 生产者任务：发送消息后关闭通道
    let producer_task = tokio::spawn(async move {
        for i in 0..5 {
            let _ = sender.send(i).await;
            println!("Producer sent: {}", i);
        }
        // 关闭通道，让消费者知道没有新消息了
        drop(sender);
        println!("Producer task finished");
    });

    // 先等待生产者和消费者完成，确保所有子任务都已创建并添加到句柄列表
    producer_task.await.unwrap();
    consumer_task.await.unwrap();

    // 提取所有任务句柄的所有权，避免持有锁等待任务完成
    let mut handles_lock = handles.lock().await;
    let task_handles = std::mem::take(&mut *handles_lock);
    drop(handles_lock);

    // 等待所有子任务完成
    join_all(task_handles).await;
    println!("All tasks finished");
}

关键说明 #

• 关闭通道：生产者任务结束时drop(sender)，触发消费者的recv()返回None，退出循环，保证所有消息都被处理、所有子任务都已创建。
• 转移所有权：用std::mem::take将锁内的向量内容转移到本地变量，获得所有JoinHandle的所有权，满足join_all的要求。
• 锁的生命周期：提前释放锁，避免在等待任务完成期间长时间持有锁，提升并发效率。

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