Wenpan Rust -- Implementing a simple task pool with Tokio

Tokio is undoubtedly the best asynchronous runtime implementation in the Rust world. The non-blocking feature brings excellent performance, but in actual development, we often need to block tasks in some cases to achieve certain functions.

Let's look at the following example

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));
                });
            }
        });
    }

The running structure we expect is to print out 99 "spawn i" through asynchronous tasks, but the actual output is probably like this

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 There will be no output after execution. If max_task is set to 2, the situation will be better, but all asynchronous operations have not been executed, that is to say, in the case of insufficient resources, Tokio will abandon some tasks. Meet our expectations. So can it be blocked again when a certain threshold is reached, and no new tasks will be given to Tokio? This is a bit similar to the thread pool. When the maximum number of threads is reached, the subsequent tasks will be blocked and the threads will be released before continuing.

Let's look at the code below.

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;
            }
        });
    }

We use JoinSet to manage derived tasks. set.join_next().await; Guarantee that at least one task is executed. Combined with the len of set, we can block task forking when the task reaches the upper limit. When the loop ends, there may be unfinished tasks, so as long as set.len() is greater than 0, wait for the task to end.

The output looks like this

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

As expected, there are not many codes, interested students can try it out.

-end-

This article is shared from the WeChat public account - JD Cloud Developers (JDT_Developers).
If there is any infringement, please contact [email protected] to delete it.
This article participates in the " OSC Source Creation Program ". You are welcome to join in and share it.

{{o.name}}

{{m.name}}