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.
15///
16/// The whole question, not the size half of it. Without the `parallel` feature
17/// there is no pool to run on, so the answer is no whatever the size — and a
18/// caller that had to pair this with its own `#[cfg]` fallback was carrying
19/// the other half itself. One of them (`process::apply_element_ops`) then left
20/// its element count with no reader at all in a serial build.
21pub fn should_parallelize(num_elements: usize) -> bool {
22 cfg!(feature = "parallel") && num_elements >= PAR_THRESHOLD
23}
24
25/// Shared rayon ThreadPool.
26///
27/// Plugins in non-blocking mode each have their own data thread, so multiple
28/// plugins may submit rayon work concurrently. A single shared pool ensures
29/// work-stealing without over-subscription.
30///
31/// The pool is sized to `available_cores - RESERVED_CORES` to leave headroom
32/// for port driver data threads, autoconnect tasks, and the tokio runtime.
33/// Call [`set_num_threads`] before the first `thread_pool()` access to override.
34#[cfg(feature = "parallel")]
35static POOL: OnceLock<rayon::ThreadPool> = OnceLock::new();
36
37/// User-specified thread count override. 0 means "use default formula".
38#[cfg(feature = "parallel")]
39static NUM_THREADS_OVERRIDE: AtomicUsize = AtomicUsize::new(0);
40
41/// Set the number of rayon worker threads before the pool is first used.
42///
43/// Must be called before any plugin processes an array. Has no effect if the
44/// pool has already been initialized.
45#[cfg(feature = "parallel")]
46pub fn set_num_threads(n: usize) {
47 NUM_THREADS_OVERRIDE.store(n, Ordering::Relaxed);
48}
49
50#[cfg(feature = "parallel")]
51pub fn thread_pool() -> &'static rayon::ThreadPool {
52 POOL.get_or_init(|| {
53 let user = NUM_THREADS_OVERRIDE.load(Ordering::Relaxed);
54 let num_threads = if user > 0 {
55 user
56 } else {
57 let available = std::thread::available_parallelism()
58 .map(|n| n.get())
59 .unwrap_or(1);
60 available.saturating_sub(RESERVED_CORES).max(1)
61 };
62 rayon::ThreadPoolBuilder::new()
63 .num_threads(num_threads)
64 .build()
65 .expect("failed to create rayon thread pool")
66 })
67}