Can rayon and tokio cooperate?

[CAD97]

Rayon, Tokio, async-std (and probably other crates) all involve managing an internal thread pool.

Based on a quick scan of the docs, it looks like you can configure rayon's global pool to run on Tokio's blocking thread pool with something along the lines of

rayon::ThreadBoolBuilder::new()
    .spawn_handler(|thread| {
        tokio::task::spawn_blocking(|| thread.run());
        Ok(())
    })
    .build_global()?;

Based on Tokio's large blocking thread limit and that it expects those threads to mainly be blocking, there may not be any benefit to this—and depending on details, it may actually not be a good idea, since IIUC Tokio expects spawn_blocking threads to finish (which the global rayon threads won't), and that entails at least some amount of overhead.

Is this something that should be done, or should rayon's thread pool be kept distinct from Tokio's? The primary advantage of putting rayon on tokio would probably be that the rayon worker threads are running in a Tokio context and thus have access to the runtime context for spawning futures.

A "vibe check" of the API suggests that the most "proper" way to nest rayon CPU work in the Tokio runtime looks to be

pub async fn spawn_compute<R: 'static + Send>(
    compute: impl 'static + Send + FnOnce() -> R,
) -> Result<R, tokio::task::JoinError> {
    tokio::task::spawn_blocking(|| {
        let mut out = None;
        rayon::scope(|s| {
            s.spawn(|_| out = Some(compute()));
        });
        out.unwrap()
    })
    .await
}

// or perhaps just

pub async fn spawn_compute<'scope, R: 'static + Send>(
    op: impl 'static + Send + FnOnce(&rayon::Scope<'scope>) -> R,
) -> Result<R, tokio::task::JoinError> {
    tokio::task::spawn_blocking(|| {
        rayon::scope(op)
    })
    .await
}

[Aiden2207]

Don't give rayon tokio's blocking threads. Because rayon's threads won't stop until the whole pool is dropped, and tokio only has a fairly small number of blocking threads, you will very quickly starve out anything that needs access to the blocking threads. That is a huge problem because a large number of the asynchronous APIs in tokio require running blocking tasks in the background. Running a separate thread pool for rayon, with the driver thread being a blocking task spawned by tokio (like in your second example) is the way to go.

[CAD97]

tokio only has a fairly small number of blocking threads

512 doesn't feel like a small number to me. In fact, spawn_blocking specifically says that

The thread limit is very large by default, because spawn_blocking is often used for various kinds of IO operations that cannot be performed asynchronously. When you run CPU-bound code using spawn_blocking, you should keep this large upper limit in mind. When running many CPU-bound computations, a semaphore or some other synchronization primitive should be used to limit the number of computation executed in parallel. Specialized CPU-bound executors, such as rayon, may also be a good fit.

---

That said, I think I agree that rayon's worker threads just being managed by rayon and not in Tokio's pool is better than running the rayon workers on Tokio's pool.

I think I'd still want to ensure that the rayon workers are in a Tokio runtime context, e.g.

        .spawn_handler(|thread| {
            let rt = tokio::runtime::Handle::current();
            let mut b = std::thread::Builder::new();
            if let Some(name) = thread.name() {
                b = b.name(name.to_owned());
            }
            if let Some(stack_size) = thread.stack_size() {
                b = b.stack_size(stack_size);
            }
            b.spawn(move || {
                let _guard = rt.enter();
                thread.run()
            })?;
            Ok(())
        })

... although tbf the extent of times a blocking task will need to pingpong back to the async runtime is probably low.

[Aiden2207]

@CAD97

512 doesn't feel like a small number to me. In fact, spawn_blocking specifically says that

My bad, it's the worker threads that there isn't a lot of (default is the number of CPUs, and tokio recommends keeping the number fairly small). I got the number of worker and blocking threads mixed up in my head somehow.

[alice]

I would not use spawn_blocking to run rayon worker threads. Just entering the Tokio runtime context on the rayon thread is a better solution.