Is there a way to create multiple pools for tokio::spawn_blocking so some tasks don't starve others?

user_1420279 0
Viewed 187

I am using Tokio for some asynchronous Rust code, and have run into a problem. I have some tasks which require access to a connection pool, and the nature of the connection pool means that only a fixed number (NUMCPUS) can run at a time - all other requests will block until there is a free connection.

Currently, I'm just using task::spawn_blocking, which kind of works. However, this has the downside that once 512 requests are blocking on the connection pool, Tokio's entire blocking pool is exhausted and all blocking tasks are just queued up. This prevents any spawn_blocking calls from elsewhere in the code that don't rely on the connection pool from running as well.

Is there any way to tell Tokio to keep a certain set of blocking tasks separate and only spawn N of them at a time, while still allowing unrelated blocking tasks to run without queueing up?

The spawn_blocking documentation suggests using Rayon for CPU intensive tasks, but a) it is not clear how to integrate Rayon with Tokio and b) my tasks are not CPU intensive anyway.

rust-tokio async-await rust
1 Answers

You can use a Semaphore: initialize it with the number of concurrently allowed tasks and have each task acquire the semaphore before processing and release it when done. Something like (untested):

use tokio::sync::Semaphore;

struct Pool {
    sem: Semaphore,
}

impl Pool {
    fn new (size: usize) -> Self {
        Pool { sem: Semaphore::new (size), }
    }

    async fn spawn<T> (&self, f: T) -> T::Output
    where
        T: Future + Send + 'static,
        T::Output: Send + 'static,
    {
        let handle = self.sem.acquire().await;
        f.await
    }
}
user_53970090
Related