1pub 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
38pub 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
51pub 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
57pub 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#[macro_export]
69macro_rules! debug_log {
70 ($($arg:tt)*) => {
71 if $crate::env::debug_enabled() {
72 eprintln!($($arg)*);
73 }
74 };
75}
76
77pub 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
95fn 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("a, &period);
108 }
109 }
110 None
111}
112
113fn 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
128fn 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
141pub 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
164pub 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 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 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 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 #[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}