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