tenferro_cpu/
context.rs

1use std::env;
2use std::sync::Arc;
3
4use crate::{Error, Result};
5
6/// Reusable CPU execution context carrying CPU parallelism policy.
7///
8/// `CpuContext` stores the requested thread count as a kernel-level
9/// parallelism hint and owns the Rayon pool used by multi-threaded CPU work.
10///
11/// # Examples
12///
13/// ```
14/// use tenferro_cpu::CpuContext;
15///
16/// let ctx = CpuContext::with_threads(1).unwrap();
17/// let value = ctx.install(|| 1 + 1);
18/// assert_eq!(value, 2);
19/// assert_eq!(ctx.num_threads(), 1);
20/// ```
21#[derive(Clone, Debug)]
22pub struct CpuContext {
23    num_threads: usize,
24    pool: Option<Arc<rayon::ThreadPool>>,
25}
26
27impl CpuContext {
28    /// Create a CPU context from `RAYON_NUM_THREADS`, or fall back to a
29    /// single-threaded context with a stderr warning when validation fails.
30    ///
31    /// # Examples
32    ///
33    /// ```
34    /// use tenferro_cpu::CpuContext;
35    ///
36    /// let ctx = CpuContext::from_env();
37    /// assert!(ctx.num_threads() >= 1);
38    /// ```
39    pub fn from_env() -> Self {
40        Self::try_from_env().unwrap_or_else(|err| {
41            eprintln!(
42                "tenferro_cpu: falling back to single-threaded CPU context after configuration error: {err}"
43            );
44            Self::single_threaded()
45        })
46    }
47
48    /// Try to create a CPU context from `RAYON_NUM_THREADS`.
49    ///
50    /// # Examples
51    ///
52    /// ```
53    /// use tenferro_cpu::CpuContext;
54    ///
55    /// let ctx = CpuContext::try_from_env()
56    ///     .unwrap_or_else(|_| CpuContext::with_threads(1).unwrap());
57    /// assert!(ctx.num_threads() >= 1);
58    /// ```
59    pub fn try_from_env() -> Result<Self> {
60        match env::var("RAYON_NUM_THREADS") {
61            Ok(value) => {
62                let num_threads = value.parse::<usize>().map_err(|err| Error::InvalidConfig {
63                    op: "CpuContext::try_from_env",
64                    message: format!("invalid RAYON_NUM_THREADS value {value:?}: {err}"),
65                })?;
66                Self::with_threads(num_threads).map_err(|err| match err {
67                    Error::InvalidConfig { message, .. } => Error::InvalidConfig {
68                        op: "CpuContext::try_from_env",
69                        message: format!("invalid RAYON_NUM_THREADS value {value:?}: {message}"),
70                    },
71                    err => err,
72                })
73            }
74            Err(env::VarError::NotPresent) => {
75                Self::with_threads(super::affinity::available_parallelism())
76            }
77            Err(err) => Err(Error::InvalidConfig {
78                op: "CpuContext::try_from_env",
79                message: format!("failed to read RAYON_NUM_THREADS: {err}"),
80            }),
81        }
82    }
83
84    /// Create a CPU context with a fixed parallelism hint.
85    ///
86    /// # Examples
87    ///
88    /// ```
89    /// use tenferro_cpu::CpuContext;
90    ///
91    /// let ctx = CpuContext::with_threads(2).unwrap();
92    /// assert_eq!(ctx.num_threads(), 2);
93    /// ```
94    ///
95    /// # Errors
96    ///
97    /// Returns an error when `num_threads` is zero or Rayon rejects the pool.
98    pub fn with_threads(num_threads: usize) -> Result<Self> {
99        if num_threads == 0 {
100            return Err(Error::InvalidConfig {
101                op: "CpuContext::with_threads",
102                message: "thread count must be at least 1".into(),
103            });
104        }
105        let pool = if num_threads == 1 {
106            None
107        } else {
108            Some(Arc::new(
109                rayon::ThreadPoolBuilder::new()
110                    .num_threads(num_threads)
111                    .build()
112                    .map_err(|err| Error::InvalidConfig {
113                        op: "CpuContext::with_threads",
114                        message: format!("failed to build CPU thread pool: {err}"),
115                    })?,
116            ))
117        };
118        Ok(Self { num_threads, pool })
119    }
120
121    fn single_threaded() -> Self {
122        Self {
123            num_threads: 1,
124            pool: None,
125        }
126    }
127
128    /// Return this context's CPU parallelism hint.
129    ///
130    /// # Examples
131    ///
132    /// ```
133    /// use tenferro_cpu::CpuContext;
134    ///
135    /// let ctx = CpuContext::with_threads(2).unwrap();
136    /// assert_eq!(ctx.num_threads(), 2);
137    /// ```
138    pub fn num_threads(&self) -> usize {
139        self.num_threads
140    }
141
142    /// Run a closure inside this context's CPU execution scope.
143    ///
144    /// # Examples
145    ///
146    /// ```
147    /// use tenferro_cpu::CpuContext;
148    ///
149    /// let ctx = CpuContext::with_threads(1).unwrap();
150    /// let value = ctx.install(|| 1 + 1);
151    /// assert_eq!(value, 2);
152    /// ```
153    pub fn install<R: Send>(&self, op: impl FnOnce() -> R + Send) -> R {
154        match &self.pool {
155            Some(pool) => pool.install(op),
156            None => op(),
157        }
158    }
159
160    /// Return the faer parallelism policy for work run inside this context.
161    ///
162    /// `Par::rayon(0)` is intentional for multi-threaded contexts: faer reads
163    /// `rayon::current_num_threads()`, so calls made under [`Self::install`]
164    /// inherit this context's Rayon pool size.
165    #[cfg(feature = "cpu-faer")]
166    #[doc(hidden)]
167    pub fn faer_par(&self) -> faer::Par {
168        if self.num_threads == 1 {
169            faer::Par::Seq
170        } else {
171            faer::Par::rayon(0)
172        }
173    }
174
175    #[cfg(feature = "cpu-faer")]
176    #[doc(hidden)]
177    pub fn faer_seq(&self) -> faer::Par {
178        faer::Par::Seq
179    }
180}
181
182#[cfg(test)]
183mod tests;
tenferro_cpu/context.rs

tenferro_cpu/
context.rs
