1use asc_cache::ProcessUsage;
5use std::path::Path;
6
7#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9#[non_exhaustive]
10pub struct Device {
11 pub total_memory: Option<u64>,
13 pub available_memory: Option<u64>,
15 pub cpus: usize,
17}
18
19impl Device {
20 pub fn probe() -> Device {
22 Device { total_memory: total_memory(), available_memory: available_memory(), cpus: asc_cache::cpus() }
23 }
24}
25
26#[derive(Debug, Clone, Copy, PartialEq, Eq)]
29#[non_exhaustive]
30pub struct MemoryPolicy {
31 pub percent: u8,
33 pub ceiling: usize,
35 pub split: bool,
37}
38
39impl MemoryPolicy {
40 pub const PERCENT: u8 = 25;
42 pub const FLOOR: usize = 64 << 20;
45 pub const CEILING_MIN: usize = 2 << 30;
47 pub const CEILING_MAX: usize = 8 << 30;
49
50 pub fn for_device(device: &Device) -> MemoryPolicy {
52 MemoryPolicy { ceiling: MemoryPolicy::auto_ceiling(device.total_memory), ..MemoryPolicy::default() }
53 }
54
55 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 pub fn alone(&self, available: Option<u64>) -> usize {
69 self.share(available, 0, &[])
70 }
71
72 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
92pub 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
103pub 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#[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
150fn 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
159fn 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 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 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 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 let info = unsafe { info.assume_init() };
193 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}