loom_rs/
builder.rs

1//! Builder pattern for constructing Loom runtimes.
2//!
3//! The builder supports multiple configuration sources using figment:
4//! - Default values
5//! - Config files (TOML, YAML, JSON)
6//! - Environment variables
7//! - Programmatic overrides
8//! - CLI arguments via clap
9
10use crate::config::LoomConfig;
11use crate::error::Result;
12use crate::mab::{CalibrationConfig, MabKnobs};
13use crate::runtime::LoomRuntime;
14
15use figment::providers::{Env, Format, Json, Serialized, Toml, Yaml};
16use figment::Figment;
17use prometheus::Registry;
18use std::path::Path;
19
20#[cfg(feature = "cuda")]
21use crate::cuda::CudaDeviceSelector;
22
23/// Builder for constructing a `LoomRuntime`.
24///
25/// Configuration sources are merged in the following order (later sources override earlier):
26/// 1. Default values
27/// 2. Config files (in order added)
28/// 3. Environment variables
29/// 4. Programmatic overrides
30///
31/// # Examples
32///
33/// ```ignore
34/// use loom_rs::LoomBuilder;
35///
36/// let runtime = LoomBuilder::new()
37///     .file("loom.toml")
38///     .env_prefix("LOOM")
39///     .prefix("myapp")
40///     .tokio_threads(2)
41///     .build()?;
42/// ```
43pub struct LoomBuilder {
44    figment: Figment,
45    prometheus_registry: Option<Registry>,
46}
47
48impl Default for LoomBuilder {
49    fn default() -> Self {
50        Self::new()
51    }
52}
53
54impl std::fmt::Debug for LoomBuilder {
55    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
56        f.debug_struct("LoomBuilder")
57            .field("figment", &self.figment)
58            .field(
59                "prometheus_registry",
60                &self.prometheus_registry.as_ref().map(|_| "<Registry>"),
61            )
62            .finish()
63    }
64}
65
66impl LoomBuilder {
67    /// Create a new builder with default configuration.
68    pub fn new() -> Self {
69        Self {
70            figment: Figment::from(Serialized::defaults(LoomConfig::default())),
71            prometheus_registry: None,
72        }
73    }
74
75    /// Add a configuration file.
76    ///
77    /// Supports TOML, YAML, and JSON formats (detected by extension).
78    /// Files are merged in the order they are added.
79    ///
80    /// # Arguments
81    ///
82    /// * `path` - Path to the configuration file
83    ///
84    /// # Examples
85    ///
86    /// ```ignore
87    /// let builder = LoomBuilder::new()
88    ///     .file("loom.toml")
89    ///     .file("loom.local.toml"); // Overrides values from loom.toml
90    /// ```
91    pub fn file<P: AsRef<Path>>(mut self, path: P) -> Self {
92        let path = path.as_ref();
93        let extension = path.extension().and_then(|e| e.to_str()).unwrap_or("");
94
95        self.figment = match extension.to_lowercase().as_str() {
96            "toml" => self.figment.merge(Toml::file(path)),
97            "yaml" | "yml" => self.figment.merge(Yaml::file(path)),
98            "json" => self.figment.merge(Json::file(path)),
99            _ => {
100                // Default to TOML
101                self.figment.merge(Toml::file(path))
102            }
103        };
104        self
105    }
106
107    /// Add environment variables with a prefix.
108    ///
109    /// Environment variables are expected in the format `{PREFIX}_{KEY}`,
110    /// e.g., `LOOM_CPUSET`, `LOOM_TOKIO_THREADS`.
111    ///
112    /// # Arguments
113    ///
114    /// * `prefix` - The environment variable prefix (without trailing underscore)
115    ///
116    /// # Examples
117    ///
118    /// ```ignore
119    /// // Will read MYAPP_CPUSET, MYAPP_TOKIO_THREADS, etc.
120    /// let builder = LoomBuilder::new().env_prefix("MYAPP");
121    /// ```
122    pub fn env_prefix(mut self, prefix: &str) -> Self {
123        self.figment = self.figment.merge(Env::prefixed(prefix).split("_"));
124        self
125    }
126
127    /// Set the thread name prefix.
128    ///
129    /// Thread names will be formatted as `{prefix}-tokio-{NNNN}` and
130    /// `{prefix}-rayon-{NNNN}`.
131    pub fn prefix(mut self, prefix: impl Into<String>) -> Self {
132        self.figment = self
133            .figment
134            .merge(Serialized::default("prefix", prefix.into()));
135        self
136    }
137
138    /// Set the CPU set string.
139    ///
140    /// Format: `"0-7,16-23"` for ranges, `"0,2,4,6"` for individual CPUs.
141    pub fn cpuset(mut self, cpuset: impl Into<String>) -> Self {
142        self.figment = self
143            .figment
144            .merge(Serialized::default("cpuset", cpuset.into()));
145        self
146    }
147
148    /// Set the number of tokio worker threads.
149    ///
150    /// Default is 1 thread.
151    pub fn tokio_threads(mut self, n: usize) -> Self {
152        self.figment = self.figment.merge(Serialized::default("tokio_threads", n));
153        self
154    }
155
156    /// Set the number of rayon threads.
157    ///
158    /// Default is the remaining CPUs after tokio threads are allocated.
159    pub fn rayon_threads(mut self, n: usize) -> Self {
160        self.figment = self.figment.merge(Serialized::default("rayon_threads", n));
161        self
162    }
163
164    /// Set the compute pool size per result type.
165    ///
166    /// Each unique result type `R` used with `spawn_compute::<F, R>()` gets its own
167    /// pool of this size. Default is 64.
168    ///
169    /// # Guidelines
170    ///
171    /// - Set to your maximum expected concurrency per type
172    /// - Higher values use more memory but reduce allocation
173    /// - Undersized pools fall back to allocation (still correct)
174    ///
175    /// # Example
176    ///
177    /// ```ignore
178    /// let runtime = LoomBuilder::new()
179    ///     .compute_pool_size(128)  // For high-concurrency workloads
180    ///     .build()?;
181    /// ```
182    pub fn compute_pool_size(mut self, size: usize) -> Self {
183        self.figment = self
184            .figment
185            .merge(Serialized::default("compute_pool_size", size));
186        self
187    }
188
189    /// Set the MAB scheduler knobs.
190    ///
191    /// These control the adaptive scheduling decisions. Most users don't need
192    /// to modify these. See [`MabKnobs`] for details.
193    ///
194    /// # Example
195    ///
196    /// ```ignore
197    /// use loom_rs::{LoomBuilder, MabKnobs};
198    ///
199    /// let runtime = LoomBuilder::new()
200    ///     .mab_knobs(MabKnobs::default().with_k_starve(0.2))
201    ///     .build()?;
202    /// ```
203    pub fn mab_knobs(mut self, knobs: MabKnobs) -> Self {
204        self.figment = self.figment.merge(Serialized::default("mab_knobs", knobs));
205        self
206    }
207
208    /// Enable calibration at runtime startup.
209    ///
210    /// Calibration measures the overhead of offloading work to rayon,
211    /// which helps the MAB make better decisions for borderline workloads.
212    ///
213    /// Default: disabled (for fast unit test startup).
214    ///
215    /// # Example
216    ///
217    /// ```ignore
218    /// let runtime = LoomBuilder::new()
219    ///     .calibrate(true)
220    ///     .build()?;
221    /// ```
222    pub fn calibrate(mut self, enabled: bool) -> Self {
223        let config = CalibrationConfig {
224            enabled,
225            ..Default::default()
226        };
227        self.figment = self
228            .figment
229            .merge(Serialized::default("calibration", config));
230        self
231    }
232
233    /// Set calibration configuration.
234    ///
235    /// Allows full control over calibration parameters.
236    ///
237    /// # Example
238    ///
239    /// ```ignore
240    /// use loom_rs::mab::CalibrationConfig;
241    ///
242    /// let runtime = LoomBuilder::new()
243    ///     .calibration_config(
244    ///         CalibrationConfig::new()
245    ///             .enabled()
246    ///             .sample_count(500)
247    ///     )
248    ///     .build()?;
249    /// ```
250    pub fn calibration_config(mut self, config: CalibrationConfig) -> Self {
251        self.figment = self
252            .figment
253            .merge(Serialized::default("calibration", config));
254        self
255    }
256
257    /// Provide an external Prometheus registry for metrics exposition.
258    ///
259    /// When a registry is provided, loom runtime metrics will be registered
260    /// and available for Prometheus scraping.
261    ///
262    /// # Example
263    ///
264    /// ```ignore
265    /// use prometheus::Registry;
266    ///
267    /// let registry = Registry::new();
268    /// let runtime = LoomBuilder::new()
269    ///     .prometheus_registry(registry.clone())
270    ///     .build()?;
271    ///
272    /// // Later: expose via HTTP endpoint
273    /// let encoder = prometheus::TextEncoder::new();
274    /// let metric_families = registry.gather();
275    /// // encoder.encode(&metric_families, &mut buffer)?;
276    /// ```
277    pub fn prometheus_registry(mut self, registry: Registry) -> Self {
278        self.prometheus_registry = Some(registry);
279        self
280    }
281
282    /// Set the CUDA device by ID.
283    ///
284    /// This will configure the runtime to use CPUs local to the specified
285    /// CUDA device's NUMA node.
286    #[cfg(feature = "cuda")]
287    pub fn cuda_device_id(mut self, id: u32) -> Self {
288        self.figment = self.figment.merge(Serialized::default(
289            "cuda_device",
290            CudaDeviceSelector::DeviceId(id),
291        ));
292        self
293    }
294
295    /// Set the CUDA device by UUID.
296    ///
297    /// This will configure the runtime to use CPUs local to the specified
298    /// CUDA device's NUMA node.
299    #[cfg(feature = "cuda")]
300    pub fn cuda_device_uuid(mut self, uuid: impl Into<String>) -> Self {
301        self.figment = self.figment.merge(Serialized::default(
302            "cuda_device",
303            CudaDeviceSelector::Uuid(uuid.into()),
304        ));
305        self
306    }
307
308    /// Apply CLI argument overrides.
309    ///
310    /// This method applies any non-None values from the `LoomArgs` struct.
311    pub fn with_cli_args(mut self, args: &LoomArgs) -> Self {
312        if let Some(ref prefix) = args.loom_prefix {
313            self.figment = self
314                .figment
315                .merge(Serialized::default("prefix", prefix.clone()));
316        }
317        if let Some(ref cpuset) = args.loom_cpuset {
318            self.figment = self
319                .figment
320                .merge(Serialized::default("cpuset", cpuset.clone()));
321        }
322        if let Some(threads) = args.loom_tokio_threads {
323            self.figment = self
324                .figment
325                .merge(Serialized::default("tokio_threads", threads));
326        }
327        if let Some(threads) = args.loom_rayon_threads {
328            self.figment = self
329                .figment
330                .merge(Serialized::default("rayon_threads", threads));
331        }
332        #[cfg(feature = "cuda")]
333        if let Some(ref device) = args.loom_cuda_device {
334            // Parse device string - could be a number or UUID
335            if let Ok(id) = device.parse::<u32>() {
336                self.figment = self.figment.merge(Serialized::default(
337                    "cuda_device",
338                    CudaDeviceSelector::DeviceId(id),
339                ));
340            } else {
341                self.figment = self.figment.merge(Serialized::default(
342                    "cuda_device",
343                    CudaDeviceSelector::Uuid(device.clone()),
344                ));
345            }
346        }
347        self
348    }
349
350    /// Build the runtime.
351    ///
352    /// This extracts the configuration and constructs the tokio and rayon
353    /// runtimes with CPU pinning.
354    ///
355    /// # Errors
356    ///
357    /// Returns an error if:
358    /// - Configuration extraction fails
359    /// - CPU set is invalid or contains unavailable CPUs
360    /// - Runtime construction fails
361    pub fn build(self) -> Result<LoomRuntime> {
362        let mut config: LoomConfig = self.figment.extract().map_err(Box::new)?;
363        config.prometheus_registry = self.prometheus_registry;
364        LoomRuntime::from_config(config)
365    }
366}
367
368/// CLI arguments for Loom configuration.
369///
370/// Use with clap's `Parser` derive macro. These arguments can be applied
371/// to a `LoomBuilder` using `with_cli_args`.
372///
373/// # Examples
374///
375/// ```ignore
376/// use clap::Parser;
377/// use loom_rs::{LoomBuilder, LoomArgs};
378///
379/// #[derive(Parser)]
380/// struct MyArgs {
381///     #[command(flatten)]
382///     loom: LoomArgs,
383///     // ... other args
384/// }
385///
386/// let args = MyArgs::parse();
387/// let runtime = LoomBuilder::new()
388///     .with_cli_args(&args.loom)
389///     .build()?;
390/// ```
391#[derive(Debug, Default, Clone, clap::Args)]
392pub struct LoomArgs {
393    /// Thread name prefix
394    #[arg(long)]
395    pub loom_prefix: Option<String>,
396
397    /// CPU set (e.g., "0-7,16-23")
398    #[arg(long)]
399    pub loom_cpuset: Option<String>,
400
401    /// Number of tokio worker threads
402    #[arg(long)]
403    pub loom_tokio_threads: Option<usize>,
404
405    /// Number of rayon threads
406    #[arg(long)]
407    pub loom_rayon_threads: Option<usize>,
408
409    /// CUDA device ID or UUID
410    #[cfg(feature = "cuda")]
411    #[arg(long)]
412    pub loom_cuda_device: Option<String>,
413}
414
415#[cfg(test)]
416mod tests {
417    use super::*;
418
419    #[test]
420    fn test_builder_defaults() {
421        let config: LoomConfig = LoomBuilder::new().figment.extract().unwrap();
422        assert_eq!(config.prefix, "loom");
423        assert!(config.cpuset.is_none());
424        assert!(config.tokio_threads.is_none());
425        assert!(config.rayon_threads.is_none());
426    }
427
428    #[test]
429    fn test_builder_programmatic_override() {
430        let config: LoomConfig = LoomBuilder::new()
431            .prefix("myapp")
432            .cpuset("0-3")
433            .tokio_threads(2)
434            .rayon_threads(6)
435            .figment
436            .extract()
437            .unwrap();
438
439        assert_eq!(config.prefix, "myapp");
440        assert_eq!(config.cpuset, Some("0-3".to_string()));
441        assert_eq!(config.tokio_threads, Some(2));
442        assert_eq!(config.rayon_threads, Some(6));
443    }
444
445    #[test]
446    fn test_builder_cli_args() {
447        let args = LoomArgs {
448            loom_prefix: Some("cliapp".to_string()),
449            loom_cpuset: Some("4-7".to_string()),
450            loom_tokio_threads: Some(1),
451            loom_rayon_threads: Some(3),
452            #[cfg(feature = "cuda")]
453            loom_cuda_device: None,
454        };
455
456        let config: LoomConfig = LoomBuilder::new()
457            .prefix("original")
458            .with_cli_args(&args)
459            .figment
460            .extract()
461            .unwrap();
462
463        // CLI args should override programmatic values
464        assert_eq!(config.prefix, "cliapp");
465        assert_eq!(config.cpuset, Some("4-7".to_string()));
466        assert_eq!(config.tokio_threads, Some(1));
467        assert_eq!(config.rayon_threads, Some(3));
468    }
469
470    #[test]
471    fn test_builder_partial_cli_args() {
472        let args = LoomArgs {
473            loom_prefix: Some("cliapp".to_string()),
474            loom_cpuset: None,
475            loom_tokio_threads: None,
476            loom_rayon_threads: None,
477            #[cfg(feature = "cuda")]
478            loom_cuda_device: None,
479        };
480
481        let config: LoomConfig = LoomBuilder::new()
482            .prefix("original")
483            .cpuset("0-3")
484            .with_cli_args(&args)
485            .figment
486            .extract()
487            .unwrap();
488
489        // Only prefix should be overridden
490        assert_eq!(config.prefix, "cliapp");
491        assert_eq!(config.cpuset, Some("0-3".to_string()));
492    }
493}
loom_rs/builder.rs

loom_rs/
builder.rs
