pub const RESERVE_MB: u64 = 2048;
pub const PER_RUN_PEAK_MB: u64 = 1500;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Machine {
pub logical_cores: usize,
pub physical_cores: Option<usize>,
pub total_memory_mb: u64,
pub available_memory_mb: u64,
}
#[derive(Debug, Clone)]
pub struct Derivation {
pub value: usize,
pub reason: String,
}
#[derive(Debug, Clone)]
pub struct Derived {
pub runs_at_once: Derivation,
pub threads_per_writer: Derivation,
pub index_memory_mb: Derivation,
}
pub fn read() -> Machine {
let logical_cores = std::thread::available_parallelism()
.map(|n| n.get())
.unwrap_or(1);
platform(logical_cores)
}
#[cfg(target_os = "linux")]
fn platform(logical_cores: usize) -> Machine {
let meminfo = std::fs::read_to_string("/proc/meminfo").unwrap_or_default();
Machine {
logical_cores,
physical_cores: linux_physical_cores(),
total_memory_mb: meminfo_kb(&meminfo, "MemTotal") / 1024,
available_memory_mb: meminfo_kb(&meminfo, "MemAvailable") / 1024,
}
}
#[cfg(target_os = "linux")]
fn meminfo_kb(meminfo: &str, key: &str) -> u64 {
meminfo
.lines()
.find_map(|line| {
let rest = line.strip_prefix(key)?.strip_prefix(':')?;
rest.split_whitespace().next()?.parse::<u64>().ok()
})
.unwrap_or(0)
}
#[cfg(target_os = "linux")]
fn linux_physical_cores() -> Option<usize> {
let mut cores = std::collections::BTreeSet::new();
for entry in std::fs::read_dir("/sys/devices/system/cpu").ok()?.flatten() {
let topology = entry.path().join("topology");
let Ok(core) = std::fs::read_to_string(topology.join("core_id")) else {
continue;
};
let package = std::fs::read_to_string(topology.join("physical_package_id"))
.unwrap_or_default()
.trim()
.to_string();
cores.insert((package, core.trim().to_string()));
}
(!cores.is_empty()).then_some(cores.len())
}
#[cfg(target_os = "macos")]
fn platform(logical_cores: usize) -> Machine {
Machine {
logical_cores,
physical_cores: sysctl::<i32>(c"hw.physicalcpu")
.filter(|n| *n > 0)
.map(|n| n as usize),
total_memory_mb: sysctl::<u64>(c"hw.memsize").unwrap_or(0) / 1024 / 1024,
available_memory_mb: macos_available_mb(),
}
}
#[cfg(target_os = "macos")]
fn sysctl<T: Default>(name: &std::ffi::CStr) -> Option<T> {
let mut out = T::default();
let mut len = std::mem::size_of::<T>();
let rc = unsafe {
libc::sysctlbyname(
name.as_ptr(),
&mut out as *mut T as *mut libc::c_void,
&mut len,
std::ptr::null_mut(),
0,
)
};
(rc == 0 && len == std::mem::size_of::<T>()).then_some(out)
}
#[cfg(target_os = "macos")]
#[allow(deprecated)]
fn macos_available_mb() -> u64 {
let mut vm: libc::vm_statistics64 = unsafe { std::mem::zeroed() };
let mut count = libc::HOST_VM_INFO64_COUNT;
let rc = unsafe {
libc::host_statistics64(
libc::mach_host_self(),
libc::HOST_VM_INFO64,
&mut vm as *mut libc::vm_statistics64 as libc::host_info64_t,
&mut count,
)
};
if rc != 0 {
return 0;
}
let page = unsafe { libc::sysconf(libc::_SC_PAGESIZE) };
if page <= 0 {
return 0;
}
let pages =
u64::from(vm.free_count) + u64::from(vm.inactive_count) + u64::from(vm.purgeable_count);
pages.saturating_mul(page as u64) / 1024 / 1024
}
#[cfg(windows)]
fn platform(logical_cores: usize) -> Machine {
use windows_sys::Win32::System::SystemInformation::{GlobalMemoryStatusEx, MEMORYSTATUSEX};
let mut status = MEMORYSTATUSEX {
dwLength: std::mem::size_of::<MEMORYSTATUSEX>() as u32,
..Default::default()
};
let ok = unsafe { GlobalMemoryStatusEx(&mut status) } != 0;
Machine {
logical_cores,
physical_cores: None,
total_memory_mb: if ok {
status.ullTotalPhys / 1024 / 1024
} else {
0
},
available_memory_mb: if ok {
status.ullAvailPhys / 1024 / 1024
} else {
0
},
}
}
#[cfg(not(any(target_os = "linux", target_os = "macos", windows)))]
fn platform(logical_cores: usize) -> Machine {
Machine {
logical_cores,
physical_cores: None,
total_memory_mb: 0,
available_memory_mb: 0,
}
}
pub fn derive(machine: &Machine, per_run_peak_mb: u64) -> Derived {
let peak = per_run_peak_mb.max(1);
let core_budget = machine.logical_cores.saturating_sub(1).max(1);
let unknown_memory = machine.available_memory_mb == 0;
let headroom = machine.available_memory_mb.saturating_sub(RESERVE_MB);
let memory_runs = if unknown_memory {
1
} else {
usize::try_from(headroom / peak)
.unwrap_or(usize::MAX)
.max(1)
};
let runs = memory_runs.min(core_budget);
let runs_reason = if unknown_memory {
format!(
"the memory reading failed, so one run at a time until there is a reading, on {} logical cores",
machine.logical_cores
)
} else if headroom < peak {
format!(
"{} free is under the {} reserve plus {} for a run, so the floor of 1 applies; {} logical cores with one kept free would allow {core_budget}",
gib(machine.available_memory_mb),
gib(RESERVE_MB),
gib(peak),
machine.logical_cores,
)
} else {
format!(
"{} free minus a {} reserve, at {} a run, allows {memory_runs}, and {} logical cores with one kept free allows {core_budget}, so {runs}",
gib(machine.available_memory_mb),
gib(RESERVE_MB),
gib(peak),
machine.logical_cores,
)
};
let index_memory = index_memory_mb(headroom);
Derived {
runs_at_once: Derivation {
value: runs,
reason: runs_reason,
},
threads_per_writer: threads_for(machine, runs),
index_memory_mb: Derivation {
value: index_memory,
reason: format!(
"{index_memory} MiB a store: the {} MiB floor, doubled once past 16 GiB beyond the reserve and again past 64 GiB, and {} beyond the reserve out of the {} free now",
crate::index::INDEX_MEMORY_MB,
beyond(headroom),
size(machine.available_memory_mb),
),
},
}
}
pub fn threads_for(machine: &Machine, runs: usize) -> Derivation {
let core_budget = machine.logical_cores.saturating_sub(1).max(1);
let runs = runs.max(1);
let per_run = (core_budget / runs).max(1);
let threads = (crate::embed::embed_threads() / runs).clamp(1, per_run);
let alone = if runs > 1 {
format!(
", and all {} for a run going alone",
crate::embed::embed_threads()
)
} else {
String::new()
};
Derivation {
value: threads,
reason: format!(
"{} embedding threads split between {} and held inside {core_budget} cores, so {threads} each{alone}",
crate::embed::embed_threads(),
runs_phrase(runs),
),
}
}
#[cfg(test)]
mod threads_tests {
use super::*;
#[test]
fn the_split_names_the_runs_in_force_however_few_cores_there_are() {
let small = Machine {
logical_cores: 3,
physical_cores: None,
total_memory_mb: 8192,
available_memory_mb: 4096,
};
let three = threads_for(&small, 3);
assert!(
three.reason.contains("3 runs"),
"the sentence describes a machine that is not this one: {}",
three.reason
);
assert!(three.value >= 1, "a run always gets at least one thread");
assert!(threads_for(&small, 1).reason.contains("1 run "));
}
}
fn index_memory_mb(headroom_mb: u64) -> usize {
let base = crate::index::INDEX_MEMORY_MB;
match headroom_mb {
0..16_384 => base,
16_384..65_536 => base * 2,
_ => base * 4,
}
}
fn runs_phrase(runs: usize) -> String {
if runs == 1 {
"1 run".to_string()
} else {
format!("{runs} runs")
}
}
fn gib(mb: u64) -> String {
let text = format!("{:.1}", mb as f64 / 1024.0);
format!("{} GiB", text.trim_end_matches(".0"))
}
fn size(mb: u64) -> String {
if mb < 1024 {
format!("{mb} MiB")
} else {
gib(mb)
}
}
fn beyond(headroom_mb: u64) -> String {
if headroom_mb == 0 {
"none is free".to_string()
} else {
format!("{} is free", size(headroom_mb))
}
}
#[cfg(test)]
mod tests {
use super::*;
fn machine(logical_cores: usize, total_gib: u64, available_gib: u64) -> Machine {
Machine {
logical_cores,
physical_cores: None,
total_memory_mb: total_gib * 1024,
available_memory_mb: available_gib * 1024,
}
}
fn assert_bounds(m: &Machine, d: &Derived) {
let runs = d.runs_at_once.value;
let threads = d.threads_per_writer.value;
assert!(runs >= 1, "{m:?} proposed {runs} runs");
assert!(threads >= 1, "{m:?} proposed {threads} threads");
assert!(
runs * threads <= m.logical_cores,
"{m:?} proposed {runs} x {threads} on {} cores",
m.logical_cores
);
assert!(
!d.runs_at_once.reason.is_empty()
&& !d.threads_per_writer.reason.is_empty()
&& !d.index_memory_mb.reason.is_empty()
);
}
#[test]
fn bounds_hold_on_every_shape_of_machine() {
let machines = [
machine(1, 2, 2), machine(2, 4, 3), machine(4, 8, 6),
machine(8, 16, 12),
machine(10, 32, 24),
machine(32, 128, 100), machine(64, 512, 480),
];
for m in machines {
let d = derive(&m, PER_RUN_PEAK_MB);
assert_bounds(&m, &d);
}
}
#[test]
fn one_free_core_survives_a_machine_with_more_than_one() {
let m = machine(32, 128, 100);
let d = derive(&m, PER_RUN_PEAK_MB);
assert!(d.runs_at_once.value * d.threads_per_writer.value < m.logical_cores);
}
#[test]
fn a_failed_memory_read_still_proposes_one_run() {
let m = Machine {
logical_cores: 8,
physical_cores: None,
total_memory_mb: 0,
available_memory_mb: 0,
};
let d = derive(&m, PER_RUN_PEAK_MB);
assert_bounds(&m, &d);
assert_eq!(d.runs_at_once.value, 1);
assert!(
d.runs_at_once.reason.contains("reading failed"),
"reason must say the read failed, got: {}",
d.runs_at_once.reason
);
}
#[test]
fn a_zero_peak_does_not_divide_by_zero() {
let m = machine(8, 16, 12);
let d = derive(&m, 0);
assert_bounds(&m, &d);
}
#[test]
fn reads_this_machine() {
let m = read();
assert!(m.logical_cores >= 1);
let d = derive(&m, PER_RUN_PEAK_MB);
assert_bounds(&m, &d);
println!("machine: {m:?}");
println!(
"runs at once = {} ({})",
d.runs_at_once.value, d.runs_at_once.reason
);
println!(
"threads per writer = {} ({})",
d.threads_per_writer.value, d.threads_per_writer.reason
);
println!(
"index memory MB = {} ({})",
d.index_memory_mb.value, d.index_memory_mb.reason
);
}
#[test]
#[cfg(target_os = "macos")]
fn total_memory_matches_sysctl() {
let Ok(out) = std::process::Command::new("sysctl")
.args(["-n", "hw.memsize"])
.output()
else {
return;
};
let Ok(bytes) = String::from_utf8_lossy(&out.stdout).trim().parse::<u64>() else {
return;
};
assert_eq!(read().total_memory_mb, bytes / 1024 / 1024);
}
#[test]
#[cfg(target_os = "linux")]
fn total_memory_matches_proc_meminfo() {
let Ok(meminfo) = std::fs::read_to_string("/proc/meminfo") else {
return;
};
assert_eq!(
read().total_memory_mb,
meminfo_kb(&meminfo, "MemTotal") / 1024
);
}
}
pub fn cpu_name() -> String {
#[cfg(target_os = "macos")]
{
if let Ok(out) = std::process::Command::new("sysctl")
.args(["-n", "machdep.cpu.brand_string"])
.output()
{
let name = String::from_utf8_lossy(&out.stdout).trim().to_string();
if !name.is_empty() {
return name;
}
}
}
#[cfg(target_os = "linux")]
{
if let Ok(info) = std::fs::read_to_string("/proc/cpuinfo")
&& let Some(name) = info
.lines()
.find(|l| l.starts_with("model name"))
.and_then(|l| l.split_once(':'))
.map(|(_, v)| v.trim().to_string())
{
return name;
}
}
#[cfg(windows)]
{
if let Ok(name) = std::env::var("PROCESSOR_IDENTIFIER") {
return name;
}
}
format!("{} CPU", std::env::consts::ARCH)
}