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

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

1. fn main(){
4.         let max_task = 1;
7.         let rt = runtime::Builder::new_multi_thread()
10.             .worker_threads(max_task)
13.             
16.             .build()
19.             .unwrap();     
22.         rt.block_on(async {
25.             println!("tokio_multi_thread ");
28.             for i in 0..100 {
31.                 println!("run {}", i);     
34.                 tokio::spawn(async move {
37.                     println!("spawn {}", i);
40.                     thread::sleep(Duration::from_secs(2));
43.                 });
46.             }
49.         });
52.     }

复制代码

我们期待的运行结构是通过异步任务打印出99个 “spawn i"，但实际输出的结果大概这样

1. tokio_multi_thread
4. run 0
7. run 1
10. run 2
13. .......
16. run 16
19. spawn 0
22. run 17
25. ......
28. run 99
31. spawn 1
34. spawn 2
37. ......
40. spawn 29
43. ......
46. spawn 58
49. spawn 59

复制代码

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

1. fn main(){
4.         let max_task = 2;
7.         let rt = runtime::Builder::new_multi_thread()
10.             .worker_threads(max_task)
13.             .enable_time()
16.             .build()
19.             .unwrap();     
22.         let mut set = JoinSet::new();
25.         rt.block_on(async {
28.             for i in 0..100 {
31.                 println!("run {}", i);
34.                 while set.len() >= max_task {
37.                     set.join_next().await;
40.                 }
43.                 set.spawn(async move {
46.                     sleep().await;
49.                     println!("spawn {}", i);
52.                 });
55.             }
58.             while set.len() > 0 {
61.                 set.join_next().await;
64.             }
67.         });
70.     }

复制代码

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

1. running 1 test
4. tokio_multi_thread
7. run 0
10. run 1
13. spawn 0
16. run 2
19. spawn 1
22. ......
25. run 31
28. spawn 30
31. run 32
34. spawn 31
37. run 33
40. ......
43. run 96
46. spawn 95
49. run 97
52. spawn 96
55. run 98
58. spawn 97
61. run 99
64. spawn 98
67. spawn 99

复制代码

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

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！