ad_plugins_rs/
par_util.rs

1#[cfg(feature = "parallel")]
2use std::sync::OnceLock;
3#[cfg(feature = "parallel")]
4use std::sync::atomic::{AtomicUsize, Ordering};
5
6/// Minimum element count to justify rayon overhead.
7pub const PAR_THRESHOLD: usize = 4096;
8
9/// Number of CPU cores reserved for driver threads, tokio runtime, etc.
10/// The rayon pool will use `available_cores - RESERVED_CORES` threads (minimum 1).
11#[cfg(feature = "parallel")]
12const RESERVED_CORES: usize = 2;
13
14/// Returns true if the data size warrants parallel processing.
15pub fn should_parallelize(num_elements: usize) -> bool {
16    num_elements >= PAR_THRESHOLD
17}
18
19/// Shared rayon ThreadPool.
20///
21/// Plugins in non-blocking mode each have their own data thread, so multiple
22/// plugins may submit rayon work concurrently. A single shared pool ensures
23/// work-stealing without over-subscription.
24///
25/// The pool is sized to `available_cores - RESERVED_CORES` to leave headroom
26/// for port driver data threads, autoconnect tasks, and the tokio runtime.
27/// Call [`set_num_threads`] before the first `thread_pool()` access to override.
28#[cfg(feature = "parallel")]
29static POOL: OnceLock<rayon::ThreadPool> = OnceLock::new();
30
31/// User-specified thread count override. 0 means "use default formula".
32#[cfg(feature = "parallel")]
33static NUM_THREADS_OVERRIDE: AtomicUsize = AtomicUsize::new(0);
34
35/// Set the number of rayon worker threads before the pool is first used.
36///
37/// Must be called before any plugin processes an array. Has no effect if the
38/// pool has already been initialized.
39#[cfg(feature = "parallel")]
40pub fn set_num_threads(n: usize) {
41    NUM_THREADS_OVERRIDE.store(n, Ordering::Relaxed);
42}
43
44#[cfg(feature = "parallel")]
45pub fn thread_pool() -> &'static rayon::ThreadPool {
46    POOL.get_or_init(|| {
47        let user = NUM_THREADS_OVERRIDE.load(Ordering::Relaxed);
48        let num_threads = if user > 0 {
49            user
50        } else {
51            let available = std::thread::available_parallelism()
52                .map(|n| n.get())
53                .unwrap_or(1);
54            available.saturating_sub(RESERVED_CORES).max(1)
55        };
56        rayon::ThreadPoolBuilder::new()
57            .num_threads(num_threads)
58            .build()
59            .expect("failed to create rayon thread pool")
60    })
61}
ad_plugins_rs/par_util.rs

ad_plugins_rs/
par_util.rs
