3

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.

kmdreko
  • 42,554
  • 6
  • 57
  • 106
Antimony
  • 37,781
  • 10
  • 100
  • 107

1 Answers1

3

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
    }
}
Jmb
  • 18,893
  • 2
  • 28
  • 55
  • 2
    You need to await `self.sem.acquire()`, otherwise the semaphore is not actually acquired, and this will compile, but result in the same starvation as the OP's original code. Also, while at it, it might be a good idea to accept a closure, and submit it to `spawn_blocking()` in order to avoid repeating that part: https://play.rust-lang.org/?version=stable&mode=debug&edition=2018&gist=20be58c9d2556d36ce4710a819ec5331 – user4815162342 Oct 13 '21 at 12:59
  • Good point about awaiting the `acquire`. However the way I understood the question, the `spawn_blocking` was just a hack to use the tokio pool limit, so it should not be required any longer (although a regular `tokio::spawn` may make sense depending on the actual use case). – Jmb Oct 13 '21 at 13:20
  • Interesting observation. The way I understood the OP's question is that `spawn_blocking()` is needed because the API they're invoking is inherently blocking. It also runs only a fixed number of instances at a time, but I understood the signal instance to be blocking as well, so failing to use `spawn_blocking()` would block the thread that calls it. The OP's problem was that unchecked use of `spawn_blocking()` swamped the thread pool internally used by `spawn_blocking()` - and your answer addresses that problem perfectly. – user4815162342 Oct 13 '21 at 13:56