rete_core/
sysmem.rs

1//! How much memory this process can still use, and the default query budget
2//! that follows from it.
3//!
4//! A lazily opened file caps its dictionary, tile and block caches with a
5//! memory budget ([`crate::Rete::set_memory_budget`]). A fixed default of
6//! [`DEFAULT_QUERY_BUDGET_MB`] (4096 MiB) is more than a 4 GiB container holds
7//! before the query's own working state is counted, so the default is sized
8//! from the memory actually available instead: half of it, at least
9//! [`MIN_DEFAULT_QUERY_BUDGET_MB`], at most the old fixed figure.
10//! `RETE_QUERY_BUDGET_MB` overrides the default, and an explicit budget (the
11//! CLI's `--memory-budget-mb`, a client's `set_memory_budget`) wins over both.
12//!
13//! The probe is the one the CLI's memory-aware local open uses: the smaller of
14//! the kernel's `MemAvailable` and the headroom under every cgroup memory
15//! limit above the process.
16
17use crate::file::DEFAULT_QUERY_BUDGET_MB;
18
19/// The lowest default query budget, whatever the memory probe says: below it
20/// the caches thrash rather than save memory (one decoded dictionary chunk or
21/// index tile is 64 KiB to a few MiB).
22pub const MIN_DEFAULT_QUERY_BUDGET_MB: u64 = 256;
23
24/// The default budget takes `1 / QUERY_BUDGET_SHARE` of the available memory,
25/// leaving the rest to what the budget does not cover: hash-join build sides,
26/// `GROUP BY` groups, `DISTINCT` sets, the result itself.
27const QUERY_BUDGET_SHARE: u64 = 2;
28
29/// The environment variable that overrides the default query budget, in MiB
30/// (`0` = unlimited).
31pub const QUERY_BUDGET_ENV: &str = "RETE_QUERY_BUDGET_MB";
32
33/// The default query budget in MiB (`0` = unlimited) for this process:
34/// `RETE_QUERY_BUDGET_MB` when set, else sized from [`available_memory`] (see
35/// [`default_query_budget_mb_from`]).
36pub fn default_query_budget_mb() -> u64 {
37    let env = std::env::var(QUERY_BUDGET_ENV)
38        .ok()
39        .and_then(|v| v.trim().parse::<u64>().ok());
40    // Only probe the memory when it is going to be used.
41    let available = if env.is_some() {
42        None
43    } else {
44        available_memory()
45    };
46    default_query_budget_mb_from(env, available)
47}
48
49/// [`default_query_budget_mb`] with its inputs explicit: `env_mb` is the
50/// override (it wins outright, `0` = unlimited), `available` the bytes the
51/// process can still allocate (`None` = unknown, which keeps the fixed
52/// [`DEFAULT_QUERY_BUDGET_MB`]). Otherwise half the available memory, clamped
53/// to `[MIN_DEFAULT_QUERY_BUDGET_MB, DEFAULT_QUERY_BUDGET_MB]`.
54pub fn default_query_budget_mb_from(env_mb: Option<u64>, available: Option<u64>) -> u64 {
55    if let Some(mb) = env_mb {
56        return mb;
57    }
58    match available {
59        Some(bytes) => ((bytes >> 20) / QUERY_BUDGET_SHARE)
60            .clamp(MIN_DEFAULT_QUERY_BUDGET_MB, DEFAULT_QUERY_BUDGET_MB),
61        None => DEFAULT_QUERY_BUDGET_MB,
62    }
63}
64
65/// Bytes this process can still allocate before something kills it: the
66/// smaller of the kernel's `MemAvailable` and the headroom under every cgroup
67/// memory limit above the process (v2 `memory.max`, or v1
68/// `memory.limit_in_bytes`). Inside a container `MemAvailable` describes the
69/// whole VM, so the cgroup limit is the one that matters there. Reclaimable
70/// page cache (`inactive_file`) does not count as used. `None` when neither
71/// source is readable (not Linux).
72pub fn available_memory() -> Option<u64> {
73    let meminfo = std::fs::read_to_string("/proc/meminfo")
74        .ok()
75        .and_then(|s| meminfo_available(&s));
76    let cgroup = cgroup_headroom(std::path::Path::new("/sys/fs/cgroup"));
77    match (meminfo, cgroup) {
78        (Some(a), Some(b)) => Some(a.min(b)),
79        (a, b) => a.or(b),
80    }
81}
82
83/// `MemAvailable` from a `/proc/meminfo` text, in bytes.
84fn meminfo_available(meminfo: &str) -> Option<u64> {
85    meminfo.lines().find_map(|line| {
86        let rest = line.strip_prefix("MemAvailable:")?;
87        let kb: u64 = rest.trim().trim_end_matches("kB").trim().parse().ok()?;
88        Some(kb.saturating_mul(1024))
89    })
90}
91
92/// One `key value` field of a cgroup `memory.stat` file.
93fn stat_field(stat: &str, key: &str) -> Option<u64> {
94    stat.lines().find_map(|line| {
95        let (k, v) = line.split_once(' ')?;
96        (k == key).then(|| v.trim().parse().ok())?
97    })
98}
99
100fn read_u64(path: &std::path::Path) -> Option<u64> {
101    std::fs::read_to_string(path).ok()?.trim().parse().ok()
102}
103
104/// Headroom under one cgroup v2 directory's `memory.max` (`None` when it has
105/// no limit: `max`, or no such file).
106fn cgroup_v2_headroom(dir: &std::path::Path) -> Option<u64> {
107    let max = read_u64(&dir.join("memory.max"))?;
108    let current = read_u64(&dir.join("memory.current")).unwrap_or(0);
109    let inactive = std::fs::read_to_string(dir.join("memory.stat"))
110        .ok()
111        .and_then(|s| stat_field(&s, "inactive_file"))
112        .unwrap_or(0);
113    Some(max.saturating_sub(current.saturating_sub(inactive)))
114}
115
116/// Headroom under a cgroup v1 memory controller directory. v1 reports "no
117/// limit" as a page-rounded `i64::MAX`, so anything past 2^60 is no limit.
118fn cgroup_v1_headroom(dir: &std::path::Path) -> Option<u64> {
119    let limit = read_u64(&dir.join("memory.limit_in_bytes")).filter(|&l| l < 1 << 60)?;
120    let usage = read_u64(&dir.join("memory.usage_in_bytes")).unwrap_or(0);
121    let inactive = std::fs::read_to_string(dir.join("memory.stat"))
122        .ok()
123        .and_then(|s| stat_field(&s, "total_inactive_file"))
124        .unwrap_or(0);
125    Some(limit.saturating_sub(usage.saturating_sub(inactive)))
126}
127
128/// The tightest headroom of the cgroups this process sits in, under the
129/// cgroup filesystem mounted at `root`. For v2 that walks from the process's
130/// own group (`0::<path>` in `/proc/self/cgroup`) up to `root`, since a limit
131/// set on any ancestor binds; inside a container with its own cgroup
132/// namespace the path is `/` and `root` itself carries the container's limit.
133fn cgroup_headroom(root: &std::path::Path) -> Option<u64> {
134    let own = std::fs::read_to_string("/proc/self/cgroup").ok();
135    let v2 = own
136        .as_deref()
137        .and_then(|s| s.lines().find_map(|l| l.strip_prefix("0::")))
138        .map(|p| p.trim().trim_start_matches('/').to_string());
139    let mut best: Option<u64> = None;
140    if let Some(rel) = v2 {
141        let mut dir = root.join(rel);
142        loop {
143            if let Some(h) = cgroup_v2_headroom(&dir) {
144                best = Some(best.map_or(h, |b| b.min(h)));
145            }
146            if dir.as_path() == root || !dir.pop() || !dir.starts_with(root) {
147                break;
148            }
149        }
150    }
151    best.or_else(|| cgroup_v1_headroom(&root.join("memory")))
152}
153
154#[cfg(test)]
155mod tests {
156    use super::*;
157
158    const MIB: u64 = 1 << 20;
159    const GIB: u64 = 1 << 30;
160
161    /// A fresh, empty directory under the system temp dir (no `tempfile` in
162    /// this crate's dev-dependencies), removed when dropped.
163    struct TempDir(std::path::PathBuf);
164
165    impl TempDir {
166        fn new(tag: &str) -> Self {
167            let n = std::time::SystemTime::now()
168                .duration_since(std::time::UNIX_EPOCH)
169                .map(|d| d.as_nanos())
170                .unwrap_or(0);
171            let p =
172                std::env::temp_dir().join(format!("rete-sysmem-{tag}-{}-{n}", std::process::id()));
173            std::fs::create_dir_all(&p).unwrap();
174            TempDir(p)
175        }
176    }
177
178    impl Drop for TempDir {
179        fn drop(&mut self) {
180            let _ = std::fs::remove_dir_all(&self.0);
181        }
182    }
183
184    #[test]
185    fn the_default_budget_follows_available_memory() {
186        // A 4 GiB container with ~3.8 GiB free: half of it, under the old
187        // fixed 4096 MiB that the container could not hold.
188        assert_eq!(default_query_budget_mb_from(None, Some(3800 * MIB)), 1900);
189        // A roomy machine keeps the old default.
190        assert_eq!(
191            default_query_budget_mb_from(None, Some(48 * GIB)),
192            DEFAULT_QUERY_BUDGET_MB
193        );
194        // A tiny one gets the floor.
195        assert_eq!(
196            default_query_budget_mb_from(None, Some(300 * MIB)),
197            MIN_DEFAULT_QUERY_BUDGET_MB
198        );
199        // Unknown memory (macOS, Windows): the old fixed figure.
200        assert_eq!(
201            default_query_budget_mb_from(None, None),
202            DEFAULT_QUERY_BUDGET_MB
203        );
204    }
205
206    #[test]
207    fn the_env_override_wins_unclamped() {
208        assert_eq!(default_query_budget_mb_from(Some(0), Some(GIB)), 0);
209        assert_eq!(default_query_budget_mb_from(Some(64), Some(48 * GIB)), 64);
210        assert_eq!(default_query_budget_mb_from(Some(9000), Some(GIB)), 9000);
211    }
212
213    #[test]
214    fn parses_meminfo() {
215        let s =
216            "MemTotal:       49283460 kB\nMemFree:         1234 kB\nMemAvailable:   40000000 kB\n";
217        assert_eq!(meminfo_available(s), Some(40_000_000 * 1024));
218        assert_eq!(meminfo_available("MemTotal: 1 kB\n"), None);
219    }
220
221    /// A fake cgroup v2 tree: a limit on the group, a tighter one on its
222    /// parent; reclaimable page cache does not count as used.
223    #[test]
224    fn cgroup_v2_headroom_takes_the_tightest_limit() {
225        let root = TempDir::new("v2");
226        let leaf = root.0.join("a/b");
227        std::fs::create_dir_all(&leaf).unwrap();
228        std::fs::write(leaf.join("memory.max"), "max\n").unwrap();
229        std::fs::write(root.0.join("a/memory.max"), format!("{}\n", 4 * GIB)).unwrap();
230        std::fs::write(
231            root.0.join("a/memory.current"),
232            format!("{}\n", GIB + 256 * MIB),
233        )
234        .unwrap();
235        std::fs::write(
236            root.0.join("a/memory.stat"),
237            format!("anon 1\ninactive_file {}\nactive_file 9\n", 256 * MIB),
238        )
239        .unwrap();
240        assert_eq!(cgroup_v2_headroom(&leaf), None, "`max` is no limit");
241        assert_eq!(
242            cgroup_v2_headroom(&root.0.join("a")),
243            Some(3 * GIB),
244            "4 GiB limit, 1.25 GiB charged of which 256 MiB is page cache"
245        );
246    }
247
248    #[test]
249    fn cgroup_v1_unlimited_is_no_limit() {
250        let dir = TempDir::new("v1");
251        std::fs::write(dir.0.join("memory.limit_in_bytes"), "9223372036854771712\n").unwrap();
252        assert_eq!(cgroup_v1_headroom(&dir.0), None);
253        std::fs::write(dir.0.join("memory.limit_in_bytes"), format!("{}", 2 * GIB)).unwrap();
254        std::fs::write(dir.0.join("memory.usage_in_bytes"), format!("{GIB}")).unwrap();
255        assert_eq!(cgroup_v1_headroom(&dir.0), Some(GIB));
256    }
257
258    /// On the Linux this suite runs on (the dev container), the probe finds
259    /// something; elsewhere `None` is the documented answer.
260    #[test]
261    fn available_memory_is_readable_on_linux() {
262        if cfg!(target_os = "linux") {
263            let a = available_memory().expect("Linux exposes MemAvailable");
264            assert!(a > 0);
265        }
266    }
267}
rete_core/sysmem.rs

rete_core/
sysmem.rs
