use std::path::{Path, PathBuf};
pub const PARALLEL_BREAK_EVEN_ENTRIES: usize = 16;
const CPU_MARGIN_DIVISOR: usize = 8;
const MEMORY_SHARE_DIVISOR: u64 = 2;
const MEMORY_INFORMATION: &str = "/proc/meminfo";
const CONTROL_GROUP_NAME: &str = "/proc/self/cgroup";
const CONTROL_GROUP_ROOT: &str = "/sys/fs/cgroup";
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct HostResources {
pub parallelism: usize,
pub memory_available: Option<u64>,
}
#[derive(Clone, Copy, Debug)]
pub struct BudgetRequest<'host> {
pub items: usize,
pub per_item_memory: u64,
pub ceiling: Option<usize>,
pub cap: Option<usize>,
pub host: &'host HostResources,
}
impl HostResources {
pub fn read() -> HostResources {
let parallelism: usize = std::thread::available_parallelism().map_or(1, std::num::NonZeroUsize::get);
let system: Option<u64> = std::fs::read_to_string(MEMORY_INFORMATION).ok().and_then(|text: String| { return memory_available_from(text.as_str()); });
let group: Option<u64> = std::fs::read_to_string(CONTROL_GROUP_NAME)
.ok()
.and_then(|text: String| { return control_group_path(text.as_str()); })
.and_then(|path: String| { return group_headroom(Path::new(CONTROL_GROUP_ROOT), path.as_str()); });
let memory_available: Option<u64> = match (system, group) {
(Some(left), Some(right)) => Some(left.min(right)),
(left, right) => left.or(right),
};
return HostResources { parallelism, memory_available };
}
}
pub fn memory_available_from(text: &str) -> Option<u64> {
let line: &str = text.lines().find(|candidate: &&str| { return candidate.starts_with("MemAvailable:"); })?;
let kibibytes: u64 = line.split_whitespace().nth(1)?.parse().ok()?;
return kibibytes.checked_mul(1024);
}
pub fn control_group_path(text: &str) -> Option<String> {
return text.lines().find_map(|line: &str| { return line.strip_prefix("0::").map(String::from); });
}
pub fn headroom_from(max: &str, current: &str) -> Option<u64> {
let limit: u64 = max.trim().parse().ok()?;
let used: u64 = current.trim().parse().ok()?;
return Some(limit.saturating_sub(used));
}
fn group_headroom(root: &Path, path: &str) -> Option<u64> {
let mut least: Option<u64> = None;
let mut directory: PathBuf = root.join(path.trim_start_matches('/'));
loop {
let limit: Option<String> = std::fs::read_to_string(directory.join("memory.max")).ok();
let current: Option<String> = std::fs::read_to_string(directory.join("memory.current")).ok();
if let (Some(max_text), Some(current_text)) = (limit, current)
&& let Some(room) = headroom_from(max_text.as_str(), current_text.as_str())
{
least = Some(least.map_or(room, |known: u64| { return known.min(room); }));
}
if directory == root || !directory.pop() || !directory.starts_with(root) {
break;
}
}
return least;
}
pub fn cpu_margin(parallelism: usize) -> usize {
if parallelism <= 2 {
return 0;
}
return (parallelism / CPU_MARGIN_DIVISOR).max(1);
}
pub fn worker_budget(request: &BudgetRequest<'_>) -> usize {
let by_cpu: usize = request.host.parallelism.saturating_sub(cpu_margin(request.host.parallelism)).max(1);
let by_memory: usize = match request.host.memory_available {
Some(bytes) => usize::try_from((bytes / MEMORY_SHARE_DIVISOR) / request.per_item_memory.max(1)).unwrap_or(usize::MAX),
None => usize::MAX,
};
return request
.items
.min(by_cpu)
.min(by_memory)
.min(request.ceiling.unwrap_or(usize::MAX))
.min(request.cap.unwrap_or(usize::MAX))
.max(1);
}
pub const SLICE_MEMORY: u64 = 16 << 20;
pub fn sliced_workers(items: usize, ceiling: usize, host: &HostResources) -> usize {
return worker_budget(&BudgetRequest { items, per_item_memory: SLICE_MEMORY, ceiling: Some(ceiling), cap: None, host });
}
#[cfg(test)]
#[path = "parallel_budget_tests.rs"]
mod tests;