Skip to main content

docling_core/
env.rs

1//! Environment-variable helpers shared by every crate in the workspace.
2//!
3//! The knobs documented in the README (`DOCLING_RS_*`, `DOCLING_*`) grew one
4//! hand-rolled `std::env::var` dance per call site — presence checks that
5//! treated `FOO=0` as *on*, three copies of the same truthiness predicate,
6//! a dozen `.ok().and_then(|v| v.parse().ok())` chains. This module is the
7//! single vocabulary for all of them:
8//!
9//! - [`flag`] — boolean knobs (`DOCLING_RS_FP32=1`);
10//! - [`nonempty`] — string knobs where a blank value means "unset";
11//! - [`parse`] — numeric knobs with a coded default;
12//! - [`debug_enabled`] / [`crate::debug_log!`] — the `DOCLING_RS_DEBUG`
13//!   diagnostics channel.
14//!
15//! On targets without an environment (wasm32-unknown-unknown) `std::env::var`
16//! reports "not present", so every helper falls back to its default — the
17//! wasm builds keep compiling with no cfg noise at the call sites.
18
19/// True when `key` is set to a truthy value. Truthy is anything except the
20/// explicit "off" spellings — empty, `0`, `false`, `no`, `off` (trimmed,
21/// ASCII case-insensitive) — so both `FOO=1` and `FOO=yes` enable, and
22/// `FOO=0` actually disables instead of counting as "present, therefore on"
23/// (the trap the old `env::var(..).is_ok()` checks all shared).
24pub fn flag(key: &str) -> bool {
25    match std::env::var(key) {
26        Ok(v) => {
27            let v = v.trim();
28            !(v.is_empty()
29                || v == "0"
30                || v.eq_ignore_ascii_case("false")
31                || v.eq_ignore_ascii_case("no")
32                || v.eq_ignore_ascii_case("off"))
33        }
34        Err(_) => false,
35    }
36}
37
38/// The trimmed value of `key`, if set and non-blank. The `Option` shape makes
39/// "env override, else default" read as `nonempty(K).unwrap_or_else(..)` and
40/// composes with `.or_else` chains for multi-variable fallbacks.
41pub fn nonempty(key: &str) -> Option<String> {
42    match std::env::var(key) {
43        Ok(v) => {
44            let v = v.trim();
45            (!v.is_empty()).then(|| v.to_string())
46        }
47        Err(_) => None,
48    }
49}
50
51/// `key` parsed as `T`, if set and parseable. Unparseable values fall back to
52/// the coded default silently — tuning knobs degrade, they don't error.
53pub fn parse<T: std::str::FromStr>(key: &str) -> Option<T> {
54    std::env::var(key).ok().and_then(|v| v.trim().parse().ok())
55}
56
57/// Whether `DOCLING_RS_DEBUG` diagnostics are on. Cached on first use: the
58/// callers sit inside per-page pipeline loops, and a process does not
59/// meaningfully flip its own debug env mid-run.
60pub fn debug_enabled() -> bool {
61    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
62    *ON.get_or_init(|| flag("DOCLING_RS_DEBUG"))
63}
64
65/// `eprintln!` gated on [`env::debug_enabled`](debug_enabled) — the quiet
66/// diagnostics channel (`DOCLING_RS_DEBUG=1`). Callers keep their own
67/// `docling-<crate>:` message prefixes; the macro only owns the gate.
68#[macro_export]
69macro_rules! debug_log {
70    ($($arg:tt)*) => {
71        if $crate::env::debug_enabled() {
72            eprintln!($($arg)*);
73        }
74    };
75}
76
77/// The CPU budget thread-pool sizing should derive from: host parallelism
78/// clamped by the container's cgroup CPU quota (#262).
79/// `available_parallelism` is quota-aware on common setups, but container
80/// runtimes exist where it still reports the host cores (docling.rs#262's
81/// 8 threads under a 4-CPU limit), so the quota files are read directly as an
82/// extra clamp — a limited container must never size pools past its throttle
83/// ceiling. On non-Linux (and wasm) the quota reads simply fail and the host
84/// count stands.
85pub fn cpu_budget() -> usize {
86    let host = std::thread::available_parallelism()
87        .map(|n| n.get())
88        .unwrap_or(1);
89    match cgroup_cpu_quota() {
90        Some(q) => host.min(q).max(1),
91        None => host,
92    }
93}
94
95/// CPUs allowed by the cgroup CPU quota, rounded up (a 2.5-CPU limit gets 3
96/// threads); `None` when unlimited or undeterminable. Reads cgroup v2
97/// (`cpu.max`) and cgroup v1 (`cpu.cfs_quota_us` / `cpu.cfs_period_us`).
98fn cgroup_cpu_quota() -> Option<usize> {
99    if let Ok(s) = std::fs::read_to_string("/sys/fs/cgroup/cpu.max") {
100        return parse_cpu_max(&s);
101    }
102    for dir in ["/sys/fs/cgroup/cpu", "/sys/fs/cgroup/cpu,cpuacct"] {
103        if let (Ok(quota), Ok(period)) = (
104            std::fs::read_to_string(format!("{dir}/cpu.cfs_quota_us")),
105            std::fs::read_to_string(format!("{dir}/cpu.cfs_period_us")),
106        ) {
107            return parse_cfs(&quota, &period);
108        }
109    }
110    None
111}
112
113/// cgroup v2 `cpu.max` ("max 100000" = unlimited, "400000 100000" = 4 CPUs).
114fn parse_cpu_max(s: &str) -> Option<usize> {
115    let mut it = s.split_whitespace();
116    let quota = it.next()?;
117    if quota == "max" {
118        return None;
119    }
120    let quota: u64 = quota.parse().ok()?;
121    let period: u64 = it.next()?.parse().ok()?;
122    if period == 0 || quota == 0 {
123        return None;
124    }
125    Some(quota.div_ceil(period) as usize)
126}
127
128/// cgroup v1 CFS quota/period (quota -1 = unlimited).
129fn parse_cfs(quota: &str, period: &str) -> Option<usize> {
130    let quota: i64 = quota.trim().parse().ok()?;
131    if quota <= 0 {
132        return None;
133    }
134    let period: i64 = period.trim().parse().ok()?;
135    if period <= 0 {
136        return None;
137    }
138    Some((quota as u64).div_ceil(period as u64) as usize)
139}
140
141/// The container's memory limit in MB from the cgroup files (v2 `memory.max`,
142/// v1 `memory.limit_in_bytes`), `None` when unlimited — both spell "no limit"
143/// as either the literal `max` or an enormous sentinel (>= 2^60 bytes).
144pub fn cgroup_memory_limit_mb() -> Option<u64> {
145    for path in [
146        "/sys/fs/cgroup/memory.max",
147        "/sys/fs/cgroup/memory/memory.limit_in_bytes",
148    ] {
149        if let Ok(s) = std::fs::read_to_string(path) {
150            let s = s.trim();
151            if s == "max" {
152                return None;
153            }
154            let bytes: u64 = s.parse().ok()?;
155            if bytes >= 1 << 60 {
156                return None;
157            }
158            return Some(bytes / (1024 * 1024));
159        }
160    }
161    None
162}
163
164/// This process's resident set size in MB (`/proc/self/status` `VmRSS`);
165/// `None` off Linux. The number admission control (#263) compares against the
166/// memory ceiling.
167pub fn rss_mb() -> Option<u64> {
168    let status = std::fs::read_to_string("/proc/self/status").ok()?;
169    let line = status.lines().find(|l| l.starts_with("VmRSS:"))?;
170    let kb: u64 = line.split_whitespace().nth(1)?.parse().ok()?;
171    Some(kb / 1024)
172}
173
174#[cfg(test)]
175mod tests {
176    // Env mutation is process-global and the test harness is parallel, so
177    // every test owns uniquely-named variables and nothing else reads them.
178    use super::*;
179
180    #[test]
181    fn flag_spellings() {
182        for on in ["1", "true", "yes", "on", "anything", " 1 ", "TRUE"] {
183            std::env::set_var("DOCLING_TEST_FLAG_ON", on);
184            assert!(flag("DOCLING_TEST_FLAG_ON"), "{on:?} should enable");
185        }
186        for off in ["", "0", "false", "no", "off", " OFF ", "No"] {
187            std::env::set_var("DOCLING_TEST_FLAG_OFF", off);
188            assert!(!flag("DOCLING_TEST_FLAG_OFF"), "{off:?} should disable");
189        }
190        assert!(!flag("DOCLING_TEST_FLAG_UNSET"));
191    }
192
193    #[test]
194    fn nonempty_trims_and_drops_blank() {
195        std::env::set_var("DOCLING_TEST_NONEMPTY", "  x  ");
196        assert_eq!(nonempty("DOCLING_TEST_NONEMPTY").as_deref(), Some("x"));
197        std::env::set_var("DOCLING_TEST_NONEMPTY_BLANK", "   ");
198        assert_eq!(nonempty("DOCLING_TEST_NONEMPTY_BLANK"), None);
199        assert_eq!(nonempty("DOCLING_TEST_NONEMPTY_UNSET"), None);
200    }
201
202    #[test]
203    fn cpu_quota_parsers_cover_both_cgroup_versions() {
204        // v2: unlimited, exact, and fractional (rounds up).
205        assert_eq!(super::parse_cpu_max("max 100000\n"), None);
206        assert_eq!(super::parse_cpu_max("400000 100000"), Some(4));
207        assert_eq!(super::parse_cpu_max("250000 100000"), Some(3));
208        assert_eq!(super::parse_cpu_max("garbage"), None);
209        // v1: -1 = unlimited; fractional rounds up.
210        assert_eq!(super::parse_cfs("-1\n", "100000\n"), None);
211        assert_eq!(super::parse_cfs("400000", "100000"), Some(4));
212        assert_eq!(super::parse_cfs("150000", "100000"), Some(2));
213        assert_eq!(super::parse_cfs("x", "100000"), None);
214    }
215
216    #[test]
217    fn rss_reads_on_linux() {
218        // A running test process certainly has a nonzero RSS on Linux.
219        #[cfg(target_os = "linux")]
220        assert!(super::rss_mb().unwrap() > 0);
221    }
222
223    #[test]
224    fn parse_trims_and_ignores_garbage() {
225        std::env::set_var("DOCLING_TEST_PARSE", " 42 ");
226        assert_eq!(parse::<usize>("DOCLING_TEST_PARSE"), Some(42));
227        std::env::set_var("DOCLING_TEST_PARSE_BAD", "many");
228        assert_eq!(parse::<usize>("DOCLING_TEST_PARSE_BAD"), None);
229        assert_eq!(parse::<usize>("DOCLING_TEST_PARSE_UNSET"), None);
230    }
231}