rayon-rs / rayon

Rayon: A data parallelism library for Rust

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Handle/guard support for current thread pool

nazar-pc opened this issue · comments

I'm doing some really not present things in the code righ now just to make rayon-based parallel code running on a specific thread pool.

Specifically I have two different thread pools and a piece of code that depending on some conditions needs to run on one of those thread pools.

The only way do do it right now is to use thread_pool.install(), but something like this would have been much more convenient:

{
    let _guard = if foo { thread_pool_a.handle().enter() } else { thread_pool_b.handle().enter() };

    // Do compute here
}

Moreover, Handle::current() should also be possible to capture current thread pool's handle, regardless of whether custom or global thread pool is used.

This is similar in behavior to tokio and tracing that both have local guards that allow you to enter/exit runtime/context/scope/span/etc.

My educated guess is that something like this is already implicitly present, just not exposed to the user.

The implicit support is just our thread-local WorkerThread pointer -- so if the current thread is already part of a pool, it will use that, otherwise we use the global pool. Then install works by sending its closure as a job to execute in the pool, so whatever runs in that closure will also see that thread-local connected to the pool.

I guess for a guard, we would need a separate TLS which is checked before anything else. I fear that will add overhead to the critical path of what would normally be a local-only operation though.