cu29_runtime/
thread_pool.rs

1//! Construction of configured worker thread pools (std-only).
2//!
3//! A pool is built from a [`ThreadPoolConfig`] into a rayon [`ThreadPool`] whose
4//! worker threads optionally get:
5//! - a fixed CPU affinity, applied **Spread**: worker `i` is pinned to
6//!   `affinity[i % affinity.len()]`, so `threads == affinity.len()` yields one
7//!   worker pinned per dedicated core; and
8//! - a scheduling policy/priority (normal/niceness or `SCHED_FIFO`/`SCHED_RR`).
9//!
10//! Affinity and scheduling are applied **once per worker thread at startup**
11//! (never on the per-CopperList hot path), and only when the `rt-scheduling`
12//! feature is enabled on a supported platform (Linux for real-time policies; CPU
13//! affinity is cross-platform). When the feature is off, the requested affinity
14//! and policy are ignored and a warning is emitted.
15//!
16//! Per-pool [`OnError`](crate::config::OnError) controls what happens when a request cannot be applied
17//! (for example, setting a real-time priority without `CAP_SYS_NICE`):
18//! [`OnError::Warn`](crate::config::OnError::Warn) logs and continues with default scheduling, while
19//! [`OnError::Strict`](crate::config::OnError::Strict) fails the build.
20
21#[cfg(feature = "rt-scheduling")]
22use crate::config::OnError;
23use crate::config::{SchedulingPolicy, ThreadPoolConfig};
24#[allow(unused_imports)] // pulls the `warning!` macro and its support symbols into scope
25use crate::log::*;
26use cu29_traits::{CuError, CuResult};
27use rayon::ThreadPool;
28
29/// Builds a rayon thread pool from a declarative [`ThreadPoolConfig`], applying
30/// the configured CPU affinity and scheduling policy/priority to each worker.
31pub fn build_pool(spec: &ThreadPoolConfig) -> CuResult<ThreadPool> {
32    let id = spec.id.clone();
33    let pool = rayon::ThreadPoolBuilder::new()
34        .num_threads(spec.threads)
35        .thread_name({
36            let id = id.clone();
37            move |i| format!("cu-pool-{id}-{i}")
38        })
39        .build()
40        .map_err(|e| CuError::from(format!("Failed to build thread pool '{id}': {e}")))?;
41
42    apply_scheduling(&pool, spec)?;
43    Ok(pool)
44}
45
46#[cfg(feature = "rt-scheduling")]
47fn apply_scheduling(pool: &ThreadPool, spec: &ThreadPoolConfig) -> CuResult<()> {
48    // Nothing to pin or reschedule: leave workers on the OS default.
49    if spec.affinity.is_none() && spec.policy == SchedulingPolicy::Fair {
50        return Ok(());
51    }
52
53    let affinity = spec.affinity.as_deref();
54    let policy = spec.policy;
55
56    // `broadcast` runs the closure on every worker thread and waits, so each
57    // setting is applied from within the thread it targets, with deterministic,
58    // ordered results. The pool has no jobs in flight yet at this point.
59    let results: Vec<CuResult<()>> =
60        pool.broadcast(|ctx| apply_to_current_thread(affinity, policy, ctx.index()));
61
62    for (worker, result) in results.into_iter().enumerate() {
63        if let Err(e) = result {
64            match spec.on_error {
65                OnError::Strict => {
66                    return Err(CuError::from(format!(
67                        "Thread pool '{}' worker {worker} could not apply scheduling: {e}",
68                        spec.id
69                    )));
70                }
71                OnError::Warn => {
72                    let pool_id = spec.id.as_str();
73                    let reason = e.to_string();
74                    warning!(
75                        "Thread pool {} worker {} could not apply scheduling ({}); using default scheduling",
76                        pool_id,
77                        worker,
78                        reason.as_str()
79                    );
80                }
81            }
82        }
83    }
84
85    Ok(())
86}
87
88#[cfg(not(feature = "rt-scheduling"))]
89fn apply_scheduling(_pool: &ThreadPool, spec: &ThreadPoolConfig) -> CuResult<()> {
90    if spec.affinity.is_some() || spec.policy != SchedulingPolicy::Fair {
91        let pool_id = spec.id.as_str();
92        warning!(
93            "Thread pool {} requests CPU affinity/scheduling but the 'rt-scheduling' feature is disabled; using default scheduling",
94            pool_id
95        );
96    }
97    Ok(())
98}
99
100/// Applies a pool's CPU affinity and scheduling policy to the **current** thread,
101/// as worker `index` (Spread: pinned to `affinity[index % affinity.len()]`).
102///
103/// This is for worker threads that are not part of a rayon pool — notably the
104/// `parallel-rt` stage workers, which are plain `std::thread::scope` threads.
105/// Behavior mirrors [`build_pool`]: [`OnError::Warn`](crate::config::OnError::Warn)
106/// logs and returns `Ok`, [`OnError::Strict`](crate::config::OnError::Strict)
107/// returns `Err`. When the `rt-scheduling` feature is off the request is ignored
108/// (with a warning).
109#[cfg(feature = "rt-scheduling")]
110pub fn apply_current_thread_scheduling(spec: &ThreadPoolConfig, index: usize) -> CuResult<()> {
111    if spec.affinity.is_none() && spec.policy == SchedulingPolicy::Fair {
112        return Ok(());
113    }
114    match apply_to_current_thread(spec.affinity.as_deref(), spec.policy, index) {
115        Ok(()) => Ok(()),
116        Err(e) => {
117            let pool_id = spec.id.as_str();
118            let reason = e.to_string();
119            match spec.on_error {
120                OnError::Strict => {
121                    error!(
122                        "Thread pool {} worker {} could not apply scheduling: {}",
123                        pool_id,
124                        index,
125                        reason.as_str()
126                    );
127                    Err(CuError::from(format!(
128                        "Thread pool '{}' worker {index} could not apply scheduling: {reason}",
129                        spec.id
130                    )))
131                }
132                OnError::Warn => {
133                    warning!(
134                        "Thread pool {} worker {} could not apply scheduling ({}); using default scheduling",
135                        pool_id,
136                        index,
137                        reason.as_str()
138                    );
139                    Ok(())
140                }
141            }
142        }
143    }
144}
145
146#[cfg(not(feature = "rt-scheduling"))]
147pub fn apply_current_thread_scheduling(spec: &ThreadPoolConfig, _index: usize) -> CuResult<()> {
148    if spec.affinity.is_some() || spec.policy != SchedulingPolicy::Fair {
149        let pool_id = spec.id.as_str();
150        warning!(
151            "Thread pool {} requests CPU affinity/scheduling but the 'rt-scheduling' feature is disabled; using default scheduling",
152            pool_id
153        );
154    }
155    Ok(())
156}
157
158#[cfg(feature = "rt-scheduling")]
159fn apply_to_current_thread(
160    affinity: Option<&[usize]>,
161    policy: SchedulingPolicy,
162    index: usize,
163) -> CuResult<()> {
164    if let Some(cores) = affinity
165        && !cores.is_empty()
166    {
167        let core = cores[index % cores.len()];
168        set_affinity(core)?;
169    }
170    set_policy(policy)
171}
172
173#[cfg(feature = "rt-scheduling")]
174fn set_affinity(core: usize) -> CuResult<()> {
175    if core_affinity::set_for_current(core_affinity::CoreId { id: core }) {
176        Ok(())
177    } else {
178        Err(CuError::from(format!(
179            "failed to pin worker to CPU core {core}"
180        )))
181    }
182}
183
184#[cfg(all(feature = "rt-scheduling", target_os = "linux"))]
185fn set_policy(policy: SchedulingPolicy) -> CuResult<()> {
186    match policy {
187        SchedulingPolicy::Fair => Ok(()),
188        SchedulingPolicy::Nice(nice) => set_nice(nice),
189        SchedulingPolicy::Fifo { priority } => set_rt_policy(libc::SCHED_FIFO, priority),
190        SchedulingPolicy::RoundRobin { priority } => set_rt_policy(libc::SCHED_RR, priority),
191    }
192}
193
194#[cfg(all(feature = "rt-scheduling", not(target_os = "linux")))]
195fn set_policy(policy: SchedulingPolicy) -> CuResult<()> {
196    match policy {
197        SchedulingPolicy::Fair => Ok(()),
198        _ => Err(CuError::from(
199            "real-time scheduling policies are only supported on Linux",
200        )),
201    }
202}
203
204#[cfg(all(feature = "rt-scheduling", target_os = "linux"))]
205fn set_rt_policy(policy: libc::c_int, priority: u8) -> CuResult<()> {
206    let param = libc::sched_param {
207        sched_priority: libc::c_int::from(priority),
208    };
209    // SAFETY: `pthread_self()` always returns a valid handle for the current
210    // thread, and `param` is a fully initialized `sched_param`.
211    let ret = unsafe { libc::pthread_setschedparam(libc::pthread_self(), policy, &param) };
212    if ret != 0 {
213        // `pthread_setschedparam` returns the error number directly.
214        return Err(CuError::from(format!(
215            "pthread_setschedparam failed ({}); setting a real-time priority typically requires CAP_SYS_NICE",
216            std::io::Error::from_raw_os_error(ret)
217        )));
218    }
219    Ok(())
220}
221
222#[cfg(all(feature = "rt-scheduling", target_os = "linux"))]
223fn set_nice(nice: i8) -> CuResult<()> {
224    // On Linux niceness is per-thread; address this thread by its tid.
225    // SAFETY: `gettid` takes no arguments and cannot fail.
226    let tid = unsafe { libc::syscall(libc::SYS_gettid) } as libc::id_t;
227
228    // `setpriority` returns -1 both on error and for the legitimate nice value
229    // -1, so disambiguate via errno.
230    // SAFETY: `__errno_location` returns a valid per-thread pointer.
231    unsafe { *libc::__errno_location() = 0 };
232    // SAFETY: simple syscall wrapper with scalar arguments.
233    let ret = unsafe { libc::setpriority(libc::PRIO_PROCESS, tid, libc::c_int::from(nice)) };
234    if ret == -1 {
235        let err = std::io::Error::last_os_error();
236        if err.raw_os_error().unwrap_or(0) != 0 {
237            return Err(CuError::from(format!("setpriority failed: {err}")));
238        }
239    }
240    Ok(())
241}
242
243#[cfg(test)]
244mod tests {
245    use super::*;
246    use crate::config::{OnError, SchedulingPolicy, ThreadPoolConfig};
247
248    fn spec(id: &str, threads: usize) -> ThreadPoolConfig {
249        ThreadPoolConfig {
250            id: id.to_string(),
251            threads,
252            affinity: None,
253            policy: SchedulingPolicy::Fair,
254            on_error: OnError::Warn,
255        }
256    }
257
258    #[test]
259    fn builds_plain_pool_with_requested_thread_count() {
260        let pool = build_pool(&spec("plain", 3)).unwrap();
261        assert_eq!(pool.current_num_threads(), 3);
262    }
263
264    #[test]
265    fn warn_mode_tolerates_unappliable_request() {
266        // A real-time priority request that may or may not succeed depending on
267        // platform/privilege; Warn mode must build the pool regardless.
268        let mut s = spec("warn", 2);
269        s.policy = SchedulingPolicy::Fifo { priority: 50 };
270        let pool = build_pool(&s).unwrap();
271        assert_eq!(pool.current_num_threads(), 2);
272    }
273
274    // Pinning to a non-existent (but in-range) CPU reliably fails regardless of
275    // privilege, so Strict mode must surface the error. Derive the bogus core id
276    // from the host's actual core count so the test stays valid on very large
277    // machines and unusual cpuset configurations.
278    #[cfg(all(feature = "rt-scheduling", target_os = "linux"))]
279    #[test]
280    fn strict_mode_fails_on_invalid_affinity() {
281        let Some(cores) = core_affinity::get_core_ids() else {
282            return; // affinity unsupported on this platform; nothing to assert
283        };
284        let max_core = cores.iter().map(|c| c.id).max().unwrap_or(0);
285        let invalid_core = max_core.saturating_add(1);
286        let mut s = spec("strict", 1);
287        s.affinity = Some(vec![invalid_core]);
288        s.on_error = OnError::Strict;
289        assert!(build_pool(&s).is_err());
290    }
291
292    // When affinity targets valid cores, pinning succeeds.
293    #[cfg(feature = "rt-scheduling")]
294    #[test]
295    fn affinity_to_valid_cores_succeeds() {
296        let Some(cores) = core_affinity::get_core_ids() else {
297            return; // affinity unsupported on this platform; nothing to assert
298        };
299        if cores.is_empty() {
300            return;
301        }
302        let mut s = spec("affinity", cores.len());
303        s.affinity = Some(cores.iter().map(|c| c.id).collect());
304        s.on_error = OnError::Strict;
305        let pool = build_pool(&s).unwrap();
306        assert_eq!(pool.current_num_threads(), cores.len());
307    }
308}
cu29_runtime/thread_pool.rs

cu29_runtime/
thread_pool.rs
