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

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

作者：京东科技贾世闻

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"，但实际输出的结果大概这样

tokio_multi_thread

run 0

run 1

run 2

.......

run 16

spawn 0

run 17

......

run 99

spawn 1

spawn 2

......

spawn 29

......

spawn 58

spawn 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 test

tokio_multi_thread

run 0

run 1

spawn 0

run 2

spawn 1

......

run 31

spawn 30

run 32

spawn 31

run 33

......

run 96

spawn 95

run 97

spawn 96

run 98

spawn 97

run 99

spawn 98

spawn 99

复制代码

符合预期，代码不多，有兴趣的同学可以动手尝试一下。

发布于: 刚刚阅读数: 3

原文链接:【http://xie.infoq.cn/article/c0320fc2bfe01dd52c507b211】。文章转载请联系作者。

京东科技开发者

关注

拥抱技术，与开发者携手创造未来！ 2018-11-20 加入

我们将持续为人工智能、大数据、云计算、物联网等相关领域的开发者，提供技术干货、行业技术内容、技术落地实践等文章内容。京东云开发者社区官方网站【https://developer.jdcloud.com/】，欢迎大家来玩

发布

暂无评论

创作场景

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

京东科技开发者

评论