1use asc_cache::ProcessUsage;
5use std::path::Path;
6
7#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9pub struct Device {
10 pub total_memory: Option<u64>,
12 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
26pub struct MemoryPolicy {
27 pub percent: u8,
29 pub ceiling: usize,
31 pub split: bool,
33}
34
35impl MemoryPolicy {
36 pub const PERCENT: u8 = 25;
38 pub const FLOOR: usize = 64 << 20;
41 pub const CEILING_MIN: usize = 2 << 30;
43 pub const CEILING_MAX: usize = 8 << 30;
44
45 pub fn for_device(device: &Device) -> MemoryPolicy {
47 MemoryPolicy { ceiling: MemoryPolicy::auto_ceiling(device.total_memory), ..MemoryPolicy::default() }
48 }
49
50 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 pub fn alone(&self, available: Option<u64>) -> usize {
64 self.share(available, 0, &[])
65 }
66
67 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
87pub 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
98pub 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#[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
145fn 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
154fn 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 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 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 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 let info = unsafe { info.assume_init() };
188 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}