1
use log::debug;
2
use rayon::ThreadPool;
3
use rayon::ThreadPoolBuilder;
4

            
5
use merc_utilities::MercError;
6

            
7
/// Builds a rayon thread pool with `num_threads` workers.
8
///
9
/// When `pinned` is set, workers are pinned round-robin to the available CPU
10
/// cores. If no core information is available on this platform, the pool is
11
/// still created but worker pinning is skipped.
12
pub fn configure_rayon_thread_pool(num_threads: usize, pinned: bool) -> Result<ThreadPool, MercError> {
13
    let worker_count = num_threads.max(1);
14
    let cores = if pinned {
15
        core_affinity2::get_core_ids().unwrap_or_default()
16
    } else {
17
        Vec::new()
18
    };
19

            
20
    ThreadPoolBuilder::new()
21
        .num_threads(worker_count)
22
        .spawn_handler(move |thread| {
23
            let thread_index = thread.index();
24
            let core = if cores.is_empty() {
25
                None
26
            } else {
27
                Some(cores[thread_index % cores.len()])
28
            };
29

            
30
            let mut builder = std::thread::Builder::new();
31
            if let Some(name) = thread.name() {
32
                builder = builder.name(name.to_owned());
33
            }
34
            if let Some(stack_size) = thread.stack_size() {
35
                builder = builder.stack_size(stack_size);
36
            }
37

            
38
            builder.spawn(move || {
39
                if let Some(core) = core {
40
                    match core.set_affinity() {
41
                        Ok(()) => {
42
                            debug!("Pinned rayon worker thread {} to core {}", thread_index, core.0);
43
                        }
44
                        Err(error) => {
45
                            debug!(
46
                                "Failed to pin rayon worker thread {} to core {}: {error}",
47
                                thread_index, core.0
48
                            );
49
                        }
50
                    }
51
                }
52

            
53
                thread.run();
54
            })?;
55

            
56
            Ok(())
57
        })
58
        .build()
59
        .map_err(|error| MercError::from(format!("Failed to build thread pool: {error}")))
60
}