loom_rs/
config.rs

1//! Configuration types for loom-rs runtime.
2
3use serde::{Deserialize, Serialize};
4
5#[cfg(feature = "cuda")]
6use crate::cuda::CudaDeviceSelector;
7
8use crate::mab::{CalibrationConfig, MabKnobs};
9use crate::pool::DEFAULT_POOL_SIZE;
10use prometheus::Registry;
11
12/// Configuration for the Loom runtime.
13///
14/// This struct can be deserialized from TOML, YAML, JSON, or environment variables
15/// using figment.
16#[derive(Debug, Clone, Deserialize, Serialize)]
17pub struct LoomConfig {
18    /// Thread name prefix (default: "loom")
19    #[serde(default = "default_prefix")]
20    pub prefix: String,
21
22    /// CPU set string (e.g., "0-7,16-23") or None for all CPUs
23    #[serde(default)]
24    pub cpuset: Option<String>,
25
26    /// Number of tokio worker threads (default: 1)
27    #[serde(default)]
28    pub tokio_threads: Option<usize>,
29
30    /// Number of rayon threads (default: remaining CPUs after tokio threads)
31    #[serde(default)]
32    pub rayon_threads: Option<usize>,
33
34    /// Size of compute task pool per result type (default: 64)
35    #[serde(default = "default_compute_pool_size")]
36    pub compute_pool_size: usize,
37
38    /// CUDA device selection (feature-gated)
39    #[cfg(feature = "cuda")]
40    #[serde(default)]
41    pub cuda_device: Option<CudaDeviceSelector>,
42
43    /// MAB scheduler configuration knobs.
44    /// If None, default knobs are used.
45    #[serde(default)]
46    pub mab_knobs: Option<MabKnobs>,
47
48    /// Calibration configuration.
49    /// If None or disabled, calibration is skipped.
50    #[serde(default)]
51    pub calibration: Option<CalibrationConfig>,
52
53    /// Prometheus registry for metrics exposition.
54    /// If provided, metrics will be registered for scraping.
55    /// Not serializable - must be set programmatically.
56    #[serde(skip)]
57    pub prometheus_registry: Option<Registry>,
58}
59
60fn default_compute_pool_size() -> usize {
61    DEFAULT_POOL_SIZE
62}
63
64fn default_prefix() -> String {
65    "loom".to_string()
66}
67
68impl Default for LoomConfig {
69    fn default() -> Self {
70        Self {
71            prefix: default_prefix(),
72            cpuset: None,
73            tokio_threads: None,
74            rayon_threads: None,
75            compute_pool_size: default_compute_pool_size(),
76            #[cfg(feature = "cuda")]
77            cuda_device: None,
78            mab_knobs: None,
79            calibration: None,
80            prometheus_registry: None,
81        }
82    }
83}
84
85impl LoomConfig {
86    /// Create a new configuration with default values.
87    pub fn new() -> Self {
88        Self::default()
89    }
90
91    /// Get the effective number of tokio threads.
92    ///
93    /// Returns the configured value or 1 as the default.
94    pub fn effective_tokio_threads(&self) -> usize {
95        self.tokio_threads.unwrap_or(1)
96    }
97
98    /// Get the effective number of rayon threads.
99    ///
100    /// Returns the configured value or calculates based on available CPUs
101    /// minus tokio threads.
102    pub fn effective_rayon_threads(&self, available_cpus: usize) -> usize {
103        self.rayon_threads
104            .unwrap_or_else(|| available_cpus.saturating_sub(self.effective_tokio_threads()))
105    }
106}
107
108#[cfg(test)]
109mod tests {
110    use super::*;
111
112    #[test]
113    fn test_default_config() {
114        let config = LoomConfig::default();
115        assert_eq!(config.prefix, "loom");
116        assert!(config.cpuset.is_none());
117        assert!(config.tokio_threads.is_none());
118        assert!(config.rayon_threads.is_none());
119        assert_eq!(config.compute_pool_size, 64);
120        assert!(config.mab_knobs.is_none());
121        assert!(config.calibration.is_none());
122    }
123
124    #[test]
125    fn test_effective_tokio_threads() {
126        let mut config = LoomConfig::default();
127        assert_eq!(config.effective_tokio_threads(), 1);
128
129        config.tokio_threads = Some(4);
130        assert_eq!(config.effective_tokio_threads(), 4);
131    }
132
133    #[test]
134    fn test_effective_rayon_threads() {
135        let mut config = LoomConfig::default();
136        // With 8 CPUs and 1 tokio thread, should get 7 rayon threads
137        assert_eq!(config.effective_rayon_threads(8), 7);
138
139        config.tokio_threads = Some(2);
140        // With 8 CPUs and 2 tokio threads, should get 6 rayon threads
141        assert_eq!(config.effective_rayon_threads(8), 6);
142
143        config.rayon_threads = Some(4);
144        // Explicit override
145        assert_eq!(config.effective_rayon_threads(8), 4);
146    }
147
148    #[test]
149    fn test_deserialize_config() {
150        let toml = r#"
151            prefix = "myapp"
152            cpuset = "0-3"
153            tokio_threads = 2
154            rayon_threads = 6
155            compute_pool_size = 128
156        "#;
157
158        let config: LoomConfig = toml::from_str(toml).unwrap();
159        assert_eq!(config.prefix, "myapp");
160        assert_eq!(config.cpuset, Some("0-3".to_string()));
161        assert_eq!(config.tokio_threads, Some(2));
162        assert_eq!(config.rayon_threads, Some(6));
163        assert_eq!(config.compute_pool_size, 128);
164    }
165}
loom_rs/config.rs

loom_rs/
config.rs
