word_tally/options/
threads.rs

1//! Thread count configuration for parallel processing.
2
3use core::{
4    fmt::{self, Display, Formatter},
5    sync::atomic::{AtomicBool, Ordering},
6};
7
8use rayon::ThreadPoolBuilder;
9use serde::{Deserialize, Serialize};
10
11use crate::WordTallyError;
12
13/// Thread count configuration for parallel processing.
14#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
15#[serde(rename_all = "camelCase")]
16pub enum Threads {
17    /// Use all available cores.
18    All,
19
20    /// Use a specific number of threads.
21    Count(u16),
22}
23
24impl Threads {
25    /// Get the actual number of threads that will be used.
26    #[must_use]
27    pub fn count(self) -> usize {
28        match self {
29            Self::All => rayon::current_num_threads(),
30            Self::Count(n) => n as usize,
31        }
32    }
33
34    /// Initialize the Rayon thread pool.
35    ///
36    /// # Errors
37    ///
38    /// Returns an error if the thread pool cannot be initialized or if the
39    /// number of threads specified is invalid.
40    pub fn init_pool(self) -> Result<(), WordTallyError> {
41        static INIT_ATTEMPTED: AtomicBool = AtomicBool::new(false);
42
43        // Only attempt initialization once using a thread-safe check
44        if INIT_ATTEMPTED.swap(true, Ordering::SeqCst) {
45            return Ok(());
46        }
47
48        // Configure thread pool based on the threads setting
49        match self {
50            Self::Count(count) => {
51                ThreadPoolBuilder::new()
52                    .num_threads(count as usize)
53                    .build_global()
54                    .map_err(|_| {
55                        WordTallyError::Config(format!(
56                            "failed to configure thread pool with {count} threads"
57                        ))
58                    })?;
59            }
60            Self::All => {
61                // Default Rayon behavior, no need to configure
62            }
63        }
64
65        Ok(())
66    }
67}
68
69impl Default for Threads {
70    fn default() -> Self {
71        Self::All
72    }
73}
74
75impl Display for Threads {
76    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
77        write!(f, "{}", self.count())
78    }
79}
80
81impl From<u16> for Threads {
82    fn from(count: u16) -> Self {
83        Self::Count(count)
84    }
85}