Skip to main content

asc_engine/
budget.rs

1//! What the machine offers ([`Device`]) and how a resident host's memory
2//! cache sizes itself from it ([`MemoryPolicy`]).
3
4use asc_cache::ProcessUsage;
5use std::path::Path;
6
7/// The machine's memory and CPUs, as the automatic budgets see them.
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9#[non_exhaustive]
10pub struct Device {
11    /// Physical memory (lowered to a cgroup v2 `memory.max`), when known.
12    pub total_memory: Option<u64>,
13    /// Memory available right now ([`available_memory`]), when known.
14    pub available_memory: Option<u64>,
15    /// CPUs this process may run on.
16    pub cpus: usize,
17}
18
19impl Device {
20    /// Read the machine's memory and CPUs now.
21    pub fn probe() -> Device {
22        Device { total_memory: total_memory(), available_memory: available_memory(), cpus: asc_cache::cpus() }
23    }
24}
25
26/// How a memory cache follows the machine: a share of the memory the
27/// user's resident hosts could hold together, split between them.
28#[derive(Debug, Clone, Copy, PartialEq, Eq)]
29#[non_exhaustive]
30pub struct MemoryPolicy {
31    /// Share of (available memory + what the caches hold), in percent.
32    pub percent: u8,
33    /// Most one process takes, in bytes.
34    pub ceiling: usize,
35    /// Split the share between the user's registered processes.
36    pub split: bool,
37}
38
39impl MemoryPolicy {
40    /// Default share of the available memory.
41    pub const PERCENT: u8 = 25;
42    /// Least one process takes: below this the LRU would thrash on every
43    /// query (it always keeps at least one entry anyway).
44    pub const FLOOR: usize = 64 << 20;
45    /// Bounds of the automatic ceiling ([`MemoryPolicy::auto_ceiling`]).
46    pub const CEILING_MIN: usize = 2 << 30;
47    /// Upper bound of the automatic ceiling, in bytes.
48    pub const CEILING_MAX: usize = 8 << 30;
49
50    /// The default policy for `device`.
51    pub fn for_device(device: &Device) -> MemoryPolicy {
52        MemoryPolicy { ceiling: MemoryPolicy::auto_ceiling(device.total_memory), ..MemoryPolicy::default() }
53    }
54
55    /// A sixteenth of physical memory, within
56    /// [`MemoryPolicy::CEILING_MIN`]..=[`MemoryPolicy::CEILING_MAX`]: 2 GiB
57    /// up to 32 GB of memory (where a quarter of the available memory is
58    /// the tighter bound anyway), 6 GiB on a 96 GB machine. One server then
59    /// keeps a phone's worth of hot DEX on a large machine, while sixteen
60    /// still split a quarter of it.
61    pub fn auto_ceiling(total_memory: Option<u64>) -> usize {
62        total_memory.map_or(MemoryPolicy::CEILING_MIN, |t| {
63            usize::try_from(t / 16).unwrap_or(usize::MAX).clamp(MemoryPolicy::CEILING_MIN, MemoryPolicy::CEILING_MAX)
64        })
65    }
66
67    /// The budget of a process alone with `available` bytes free.
68    pub fn alone(&self, available: Option<u64>) -> usize {
69        self.share(available, 0, &[])
70    }
71
72    /// This process's share: [`MemoryPolicy::percent`] of what the user's
73    /// registered processes could hold together (the available memory plus
74    /// what their caches hold already), split evenly when
75    /// [`MemoryPolicy::split`], within [`MemoryPolicy::FLOOR`]..=ceiling.
76    /// The ceiling when the available memory is unknown.
77    pub fn share(&self, available: Option<u64>, held_here: u64, others: &[ProcessUsage]) -> usize {
78        let Some(available) = available else { return self.ceiling.max(MemoryPolicy::FLOOR) };
79        let others = if self.split { others } else { &[] };
80        let held: u64 = held_here + others.iter().map(|o| o.bytes).sum::<u64>();
81        let share = u128::from(available.saturating_add(held)) * u128::from(self.percent) / 100 / (others.len() as u128 + 1);
82        usize::try_from(share).unwrap_or(usize::MAX).clamp(MemoryPolicy::FLOOR, self.ceiling.max(MemoryPolicy::FLOOR))
83    }
84}
85
86impl Default for MemoryPolicy {
87    fn default() -> MemoryPolicy {
88        MemoryPolicy { percent: MemoryPolicy::PERCENT, ceiling: MemoryPolicy::CEILING_MIN, split: true }
89    }
90}
91
92/// Bytes of memory available to this process: on Linux `MemAvailable`,
93/// lowered to the headroom left under a cgroup v2 `memory.max`; on Apple
94/// systems the free, speculative and inactive pages. `None` elsewhere.
95pub fn available_memory() -> Option<u64> {
96    let system = system_available();
97    match (system, cgroup_limit().map(|(max, current)| max.saturating_sub(current))) {
98        (Some(a), Some(b)) => Some(a.min(b)),
99        (a, b) => a.or(b),
100    }
101}
102
103/// Physical memory, lowered to a cgroup v2 `memory.max`. `None` where it
104/// cannot be read.
105pub fn total_memory() -> Option<u64> {
106    match (system_total(), cgroup_limit().map(|(max, _)| max)) {
107        (Some(a), Some(b)) => Some(a.min(b)),
108        (a, b) => a.or(b),
109    }
110}
111
112#[cfg(target_os = "linux")]
113fn system_available() -> Option<u64> {
114    meminfo_field(&std::fs::read_to_string("/proc/meminfo").ok()?, "MemAvailable:")
115}
116
117#[cfg(target_os = "linux")]
118fn system_total() -> Option<u64> {
119    meminfo_field(&std::fs::read_to_string("/proc/meminfo").ok()?, "MemTotal:")
120}
121
122#[cfg(target_vendor = "apple")]
123fn system_available() -> Option<u64> {
124    apple::available()
125}
126
127#[cfg(target_vendor = "apple")]
128fn system_total() -> Option<u64> {
129    apple::total()
130}
131
132#[cfg(not(any(target_os = "linux", target_vendor = "apple")))]
133fn system_available() -> Option<u64> {
134    None
135}
136
137#[cfg(not(any(target_os = "linux", target_vendor = "apple")))]
138fn system_total() -> Option<u64> {
139    None
140}
141
142/// A `/proc/meminfo` field (`MemAvailable:`, `MemTotal:`), in bytes.
143#[cfg_attr(not(target_os = "linux"), allow(dead_code))]
144fn meminfo_field(meminfo: &str, field: &str) -> Option<u64> {
145    let line = meminfo.lines().find(|l| l.starts_with(field))?;
146    let kb: u64 = line.split_whitespace().nth(1)?.parse().ok()?;
147    kb.checked_mul(1024)
148}
149
150/// `(memory.max, memory.current)` of this process's cgroup v2, if limited.
151fn cgroup_limit() -> Option<(u64, u64)> {
152    let cgroup = std::fs::read_to_string("/proc/self/cgroup").ok()?;
153    let rel = cgroup.lines().find_map(|l| l.strip_prefix("0::"))?;
154    let dir = Path::new("/sys/fs/cgroup").join(rel.trim_start_matches('/'));
155    let read = |f: &str| std::fs::read_to_string(dir.join(f)).ok();
156    cgroup_limit_from(&read("memory.max")?, &read("memory.current")?)
157}
158
159/// Limit and usage from the text of `memory.max` (`max` when unlimited)
160/// and `memory.current`.
161fn cgroup_limit_from(max: &str, current: &str) -> Option<(u64, u64)> {
162    Some((max.trim().parse().ok()?, current.trim().parse().ok()?))
163}
164
165#[cfg(target_vendor = "apple")]
166mod apple {
167    use std::os::raw::c_uint;
168
169    unsafe extern "C" {
170        // Declared here: libc marks it deprecated in favour of the mach2 crate.
171        fn mach_host_self() -> c_uint;
172    }
173
174    pub(super) fn total() -> Option<u64> {
175        let mut bytes: u64 = 0;
176        let mut len = std::mem::size_of::<u64>();
177        // SAFETY: the name is NUL-terminated and `bytes` / `len` describe a u64 out buffer.
178        let rc = unsafe { libc::sysctlbyname(c"hw.memsize".as_ptr(), (&mut bytes as *mut u64).cast(), &mut len, std::ptr::null_mut(), 0) };
179        (rc == 0 && len == std::mem::size_of::<u64>()).then_some(bytes)
180    }
181
182    pub(super) fn available() -> Option<u64> {
183        let mut info = std::mem::MaybeUninit::<libc::vm_statistics64>::zeroed();
184        let mut count = libc::HOST_VM_INFO64_COUNT;
185        // SAFETY: `info` is a vm_statistics64 out buffer of `count` integers,
186        // as HOST_VM_INFO64 expects; the host port needs no release.
187        let rc = unsafe { libc::host_statistics64(mach_host_self(), libc::HOST_VM_INFO64, info.as_mut_ptr().cast(), &mut count) };
188        if rc != 0 {
189            return None;
190        }
191        // SAFETY: host_statistics64 succeeded, so it filled `info`.
192        let info = unsafe { info.assume_init() };
193        // SAFETY: plain query of a system constant.
194        let page = u64::try_from(unsafe { libc::sysconf(libc::_SC_PAGESIZE) }).ok()?;
195        let pages = u64::from(info.free_count) + u64::from(info.speculative_count) + u64::from(info.inactive_count);
196        Some(pages.saturating_mul(page))
197    }
198}
199
200#[cfg(test)]
201mod tests {
202    use super::*;
203    use std::time::Duration;
204
205    const GIB: u64 = 1 << 30;
206
207    #[test]
208    fn alone_is_a_clamped_share_of_available_memory() {
209        let p = MemoryPolicy::default();
210        assert_eq!(p.alone(None), MemoryPolicy::CEILING_MIN);
211        assert_eq!(p.alone(Some(4 * GIB)), GIB as usize);
212        assert_eq!(p.alone(Some(64 * GIB)), MemoryPolicy::CEILING_MIN);
213        assert_eq!(p.alone(Some(100 << 20)), MemoryPolicy::FLOOR);
214        assert_eq!(p.alone(Some(u64::MAX)), MemoryPolicy::CEILING_MIN);
215        let half = MemoryPolicy { percent: 50, ceiling: 16 << 30, split: true };
216        assert_eq!(half.alone(Some(8 * GIB)), (4 * GIB) as usize);
217    }
218
219    #[test]
220    fn share_splits_between_processes_unless_told_not_to() {
221        let other = |bytes| ProcessUsage::new(1, bytes, Duration::ZERO, None);
222        let p = MemoryPolicy { ceiling: 8 << 30, ..MemoryPolicy::default() };
223        assert_eq!(p.share(Some(3 * GIB), GIB, &[]), GIB as usize);
224        let others = vec![other(GIB / 8); 7];
225        assert_eq!(p.share(Some(8 * GIB), GIB / 8, &others), (9 * GIB / 4 / 8) as usize);
226        assert_eq!(p.share(Some(GIB), 0, &vec![other(0); 63]), MemoryPolicy::FLOOR);
227        assert_eq!(p.share(None, 0, &others), 8 << 30);
228        let alone = MemoryPolicy { split: false, ..p };
229        assert_eq!(alone.share(Some(8 * GIB), 0, &others), (2 * GIB) as usize);
230    }
231
232    #[test]
233    fn the_ceiling_follows_physical_memory() {
234        assert_eq!(MemoryPolicy::auto_ceiling(None), 2 << 30);
235        assert_eq!(MemoryPolicy::auto_ceiling(Some(12 * GIB)), 2 << 30);
236        assert_eq!(MemoryPolicy::auto_ceiling(Some(96 * GIB)), 6 << 30);
237        assert_eq!(MemoryPolicy::auto_ceiling(Some(512 * GIB)), 8 << 30);
238        let d = Device { total_memory: Some(64 * GIB), available_memory: None, cpus: 8 };
239        assert_eq!(MemoryPolicy::for_device(&d).ceiling, 4 << 30);
240    }
241
242    #[test]
243    fn memory_probes_parse() {
244        let meminfo = "MemTotal:       16000000 kB\nMemFree:         1000000 kB\nMemAvailable:    8000000 kB\n";
245        assert_eq!(meminfo_field(meminfo, "MemAvailable:"), Some(8_000_000 * 1024));
246        assert_eq!(meminfo_field(meminfo, "MemTotal:"), Some(16_000_000 * 1024));
247        assert_eq!(meminfo_field("MemFree: 1 kB\n", "MemAvailable:"), None);
248        assert_eq!(cgroup_limit_from("1073741824\n", "268435456\n"), Some((1 << 30, 256 << 20)));
249        assert_eq!(cgroup_limit_from("max\n", "268435456\n"), None);
250        if cfg!(any(target_os = "linux", target_vendor = "apple")) {
251            let d = Device::probe();
252            assert!(d.total_memory.is_some() && d.available_memory.is_some(), "{d:?}");
253        }
254    }
255}