tract_data/
knobs.rs

1//! Runtime configuration knobs.
2//!
3//! A knob is a named, typed runtime setting — a kernel-selection toggle or a
4//! heuristic threshold. Its value is resolved, in priority order, from an
5//! explicit programmatic override, then the process environment, then a
6//! compiled-in default.
7//!
8//! Knobs are declared with [`declare_knob!`] and register themselves into a
9//! stack-wide inventory (this crate is the lowest in the stack, so `linalg` and
10//! everything above can declare). The inventory lets every knob be listed and
11//! documented from one place, and set through the API on targets where
12//! environment variables are unavailable (e.g. wasm).
13//!
14//! A knob is the lowest, deployer-facing layer of configuration; it is not a
15//! substitute for settings that must travel with a model.
16
17use std::collections::HashMap;
18use std::str::FromStr;
19
20use parking_lot::RwLock;
21
22pub use inventory;
23
24/// A type usable as a knob value: parseable from the string sources (env / API)
25/// and renderable for listing.
26pub trait KnobValue: Sized + Clone + Send + Sync + 'static {
27    const TYPE_NAME: &'static str;
28    /// Parse from a source string. `None` means "unset/invalid": the resolver
29    /// falls through to the next source rather than overriding with garbage.
30    fn parse_knob(s: &str) -> Option<Self>;
31    fn render_knob(&self) -> String;
32}
33
34impl KnobValue for bool {
35    const TYPE_NAME: &'static str = "bool";
36    fn parse_knob(s: &str) -> Option<Self> {
37        match s.trim().to_ascii_lowercase().as_str() {
38            "1" | "true" | "on" | "yes" => Some(true),
39            "0" | "false" | "off" | "no" => Some(false),
40            _ => None,
41        }
42    }
43    fn render_knob(&self) -> String {
44        if *self { "true" } else { "false" }.to_string()
45    }
46}
47
48macro_rules! impl_knob_value_fromstr {
49    ($ty:ty, $name:expr) => {
50        impl KnobValue for $ty {
51            const TYPE_NAME: &'static str = $name;
52            fn parse_knob(s: &str) -> Option<Self> {
53                <$ty>::from_str(s.trim()).ok()
54            }
55            fn render_knob(&self) -> String {
56                self.to_string()
57            }
58        }
59    };
60}
61impl_knob_value_fromstr!(i64, "i64");
62impl_knob_value_fromstr!(usize, "usize");
63impl_knob_value_fromstr!(f32, "f32");
64
65impl KnobValue for String {
66    const TYPE_NAME: &'static str = "string";
67    fn parse_knob(s: &str) -> Option<Self> {
68        Some(s.to_string())
69    }
70    fn render_knob(&self) -> String {
71        self.clone()
72    }
73}
74
75/// Optional knob: `None` is the unset/default state (e.g. "autodetect"), `Some`
76/// is an explicit override. A source string that does not parse as `T` is
77/// treated as unset, so the resolver falls through to the next source.
78impl<T: KnobValue> KnobValue for Option<T> {
79    const TYPE_NAME: &'static str = "option";
80    fn parse_knob(s: &str) -> Option<Self> {
81        T::parse_knob(s).map(Some)
82    }
83    fn render_knob(&self) -> String {
84        match self {
85            Some(v) => v.render_knob(),
86            None => "(unset)".to_string(),
87        }
88    }
89}
90
91static OVERRIDES: RwLock<Option<HashMap<&'static str, String>>> = RwLock::new(None);
92
93fn with_override<R>(name: &str, f: impl FnOnce(Option<&String>) -> R) -> R {
94    let guard = OVERRIDES.read();
95    f(guard.as_ref().and_then(|m| m.get(name)))
96}
97
98/// A declared knob. Construct via [`declare_knob!`] rather than directly so it
99/// is registered for listing.
100pub struct Knob<T: KnobValue> {
101    pub name: &'static str,
102    pub doc: &'static str,
103    default: fn() -> T,
104}
105
106impl<T: KnobValue> Knob<T> {
107    pub const fn new(name: &'static str, default: fn() -> T, doc: &'static str) -> Self {
108        Knob { name, doc, default }
109    }
110
111    /// The compiled-in default, ignoring override and environment.
112    pub fn default_value(&self) -> T {
113        (self.default)()
114    }
115
116    /// Resolve the current value: programmatic override, then environment, then
117    /// default.
118    pub fn get(&self) -> T {
119        if let Some(v) = with_override(self.name, |o| o.and_then(|s| T::parse_knob(s))) {
120            return v;
121        }
122        if let Ok(s) = std::env::var(self.name)
123            && let Some(v) = T::parse_knob(&s)
124        {
125            return v;
126        }
127        self.default_value()
128    }
129
130    /// Set a programmatic override (highest priority), effective for subsequent
131    /// [`get`](Self::get) calls. Prefer threading config through build options
132    /// where possible; this is global mutable state.
133    pub fn set(&self, value: T) {
134        OVERRIDES.write().get_or_insert_with(HashMap::new).insert(self.name, value.render_knob());
135    }
136
137    /// Drop the programmatic override, reverting to environment/default.
138    pub fn clear(&self) {
139        if let Some(m) = OVERRIDES.write().as_mut() {
140            m.remove(self.name);
141        }
142    }
143}
144
145/// Type-erased description of a declared knob, collected stack-wide for listing
146/// and documentation.
147pub struct KnobInfo {
148    pub name: &'static str,
149    pub doc: &'static str,
150    pub type_name: &'static str,
151    /// The default value, rendered.
152    pub default: fn() -> String,
153    /// The currently-resolved value, rendered.
154    pub current: fn() -> String,
155}
156
157inventory::collect!(KnobInfo);
158
159/// All declared knobs, sorted by name.
160pub fn all() -> Vec<&'static KnobInfo> {
161    let mut v: Vec<&'static KnobInfo> = inventory::iter::<KnobInfo>.into_iter().collect();
162    v.sort_by_key(|k| k.name);
163    v
164}
165
166/// Set a knob by name from a raw string — the environment-style source, for
167/// hosts where environment variables are unavailable (e.g. wasm). Returns
168/// `false` if no knob by that name is declared.
169pub fn set_str(name: &str, value: &str) -> bool {
170    match all().into_iter().find(|k| k.name == name) {
171        Some(info) => {
172            OVERRIDES.write().get_or_insert_with(HashMap::new).insert(info.name, value.to_string());
173            true
174        }
175        None => false,
176    }
177}
178
179/// Drop a knob's programmatic override by name. Returns `false` if not declared.
180pub fn clear_str(name: &str) -> bool {
181    match all().into_iter().find(|k| k.name == name) {
182        Some(info) => {
183            if let Some(m) = OVERRIDES.write().as_mut() {
184                m.remove(info.name);
185            }
186            true
187        }
188        None => false,
189    }
190}
191
192/// Render all declared knobs as a human-readable list.
193pub fn list() -> String {
194    use std::fmt::Write;
195    let shown = |s: String| if s.is_empty() { "(unset)".to_string() } else { s };
196    let mut out = String::new();
197    for k in all() {
198        let current = shown((k.current)());
199        let default = shown((k.default)());
200        write!(out, "{} = {}  [{}", k.name, current, k.type_name).unwrap();
201        if current != default {
202            write!(out, ", default {default}").unwrap();
203        }
204        writeln!(out, "]\n    {}\n", k.doc).unwrap();
205    }
206    out
207}
208
209/// Declare a knob: a typed runtime setting resolved from override / environment
210/// / default, registered for listing.
211///
212/// The static is named exactly like the environment variable (the `TRACT_`
213/// prefix included), so one grep finds both the declaration and any use, and
214/// the env-var name is derived from the identifier — there is no separate
215/// string to drift.
216///
217/// ```ignore
218/// declare_knob!(TRACT_MY_FLAG, bool, false, "Enable the thing.");
219/// // ... later ...
220/// if TRACT_MY_FLAG.get() { /* ... */ }
221/// ```
222#[macro_export]
223macro_rules! declare_knob {
224    ($ident:ident, $ty:ty, $default:expr, $doc:literal) => {
225        #[doc = $doc]
226        pub static $ident: $crate::knobs::Knob<$ty> =
227            $crate::knobs::Knob::new(stringify!($ident), || $default, $doc);
228
229        $crate::knobs::inventory::submit! {
230            $crate::knobs::KnobInfo {
231                name: stringify!($ident),
232                doc: $doc,
233                type_name: <$ty as $crate::knobs::KnobValue>::TYPE_NAME,
234                default: || $crate::knobs::KnobValue::render_knob(&$ident.default_value()),
235                current: || $crate::knobs::KnobValue::render_knob(&$ident.get()),
236            }
237        }
238    };
239}
240
241#[cfg(test)]
242mod tests {
243    use super::*;
244
245    declare_knob!(TRACT_TEST_KNOB_BOOL, bool, false, "Test bool knob.");
246    declare_knob!(TRACT_TEST_KNOB_INT, usize, 7, "Test int knob.");
247    declare_knob!(TRACT_TEST_KNOB_SET_INT, usize, 7, "Test int knob for set-by-name.");
248
249    #[test]
250    fn default_then_override_then_clear() {
251        assert!(!TRACT_TEST_KNOB_BOOL.get());
252        assert_eq!(TRACT_TEST_KNOB_INT.get(), 7);
253
254        TRACT_TEST_KNOB_BOOL.set(true);
255        assert!(TRACT_TEST_KNOB_BOOL.get());
256        TRACT_TEST_KNOB_BOOL.clear();
257        assert!(!TRACT_TEST_KNOB_BOOL.get());
258    }
259
260    #[test]
261    fn set_by_name() {
262        assert!(set_str("TRACT_TEST_KNOB_SET_INT", "42"));
263        assert_eq!(TRACT_TEST_KNOB_SET_INT.get(), 42);
264        assert!(clear_str("TRACT_TEST_KNOB_SET_INT"));
265        assert_eq!(TRACT_TEST_KNOB_SET_INT.get(), 7);
266        assert!(!set_str("TRACT_TEST_KNOB_DOES_NOT_EXIST", "1"));
267    }
268
269    #[test]
270    fn registered_in_inventory() {
271        let names: Vec<_> = all().iter().map(|k| k.name).collect();
272        assert!(names.contains(&"TRACT_TEST_KNOB_BOOL"));
273        assert!(names.contains(&"TRACT_TEST_KNOB_INT"));
274    }
275
276    #[test]
277    fn bool_parsing() {
278        assert_eq!(bool::parse_knob("1"), Some(true));
279        assert_eq!(bool::parse_knob("Off"), Some(false));
280        assert_eq!(bool::parse_knob("maybe"), None);
281    }
282}
tract_data/knobs.rs

tract_data/
knobs.rs
