文盘 Rust -- 用 Tokio 实现简易任务池

2023-04-07
北京
本文字数：1317 字
阅读完需：约 4 分钟

作者： jiashiwen 原文来源：https://tidb.net/blog/a3f2d5d4

notice”Rust is a trademark of the Mozilla Foundation in the US and other countries.”

Tokio 无疑是 Rust 世界中最优秀的异步 Runtime 实现。非阻塞的特性带来了优异的性能，但是在实际的开发中我们往往需要在某些情况下阻塞任务来实现某些功能。我们看看下面的例子

fn main(){        let max_task = 1;        let rt = runtime::Builder::new_multi_thread()            .worker_threads(max_task)                        .build()            .unwrap();     
        rt.block_on(async {            println!("tokio_multi_thread ");            for i in 0..100 {                println!("run {}", i);                     tokio::spawn(async move {                    println!("spawn {}", i);                    thread::sleep(Duration::from_secs(2));                });            }        });    }

复制代码

我们期待的运行结构是通过异步任务打印出 99 个 “spawn i”，但实际输出的结果大概这样。原因是在 async move 以后，主线程未阻塞。

tokio_multi_threadrun 0run 1run 2.......run 16spawn 0run 17......run 99spawn 1spawn 2......spawn 29......spawn 58spawn 59

复制代码

59 执行完后面就没有输出了，如果把 max_task 设置为 2，情况会好一点，但是也没有执行完所有的异步操作，也就是说在资源不足的情况下，Tokio 会抛弃某些任务，这不符合我们的预期。那么能不能再达到了某一阀值的情况下阻塞一下，不再给 Tokio 新的任务呢。这有点类似线程池，当达达最大线程数的时候阻塞后面的任务待有释放的线程后再继续。我们看看下面的代码。

fn main(){        let max_task = 2;        let rt = runtime::Builder::new_multi_thread()            .worker_threads(max_task)            .enable_time()            .build()            .unwrap();             let mut set = JoinSet::new();        rt.block_on(async {            for i in 0..100 {                println!("run {}", i);                while set.len() >= max_task {                    set.join_next().await;                }                set.spawn(async move {                    sleep().await;                    println!("spawn {}", i);                });            }            while set.len() > 0 {                set.join_next().await;            }        });    }

复制代码

我们使用 JoinSet 来管理派生出来的任务。set.join_next().await; 保证至少一个任务被执行完成。结合 set 的 len，我们可以在任务达到上限时阻塞任务派生。当循环结束，可能还有未完成的任务，所以只要 set.len() 大于 0 就等待任务结束。

输出大概长这样

running 1 testtokio_multi_threadrun 0run 1spawn 0run 2spawn 1......run 31spawn 30run 32spawn 31run 33......run 96spawn 95run 97spawn 96run 98spawn 97run 99spawn 98spawn 99