Files
searchi/crates/searchi-extract-worker/src/resources.rs
T

65 lines
2.4 KiB
Rust

use anyhow::{ensure, Result};
#[cfg(not(target_os = "macos"))]
const MEMORY_BYTES: usize = 512 * 1024 * 1024;
#[cfg(windows)]
pub fn apply() -> Result<()> {
use windows_sys::Win32::System::JobObjects::{
JobObjectExtendedLimitInformation, QueryInformationJobObject,
JOBOBJECT_EXTENDED_LIMIT_INFORMATION, JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE,
JOB_OBJECT_LIMIT_PROCESS_MEMORY,
};
let mut information: JOBOBJECT_EXTENDED_LIMIT_INFORMATION = unsafe { std::mem::zeroed() };
// NULL addresses the job associated with the current process. The parent
// must assign this process before writing its first input bytes.
let success = unsafe {
QueryInformationJobObject(
std::ptr::null_mut(),
JobObjectExtendedLimitInformation,
&mut information as *mut _ as *mut _,
std::mem::size_of_val(&information) as u32,
std::ptr::null_mut(),
)
};
ensure!(
success != 0,
"Extraction worker requires parent resource job"
);
let required = JOB_OBJECT_LIMIT_PROCESS_MEMORY | JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
ensure!(
information.BasicLimitInformation.LimitFlags & required == required
&& information.ProcessMemoryLimit > 0
&& information.ProcessMemoryLimit <= MEMORY_BYTES,
"Extraction worker job resource policy is missing"
);
Ok(())
}
#[cfg(unix)]
pub fn apply() -> Result<()> {
// Darwin rejects finite AS/DATA limits. The parent enforces a sampled RSS
// budget there; retain the CPU limit and bounded path-free transport here.
let set_limit = |resource, maximum| -> Result<()> {
let value = libc::rlimit {
rlim_cur: maximum,
rlim_max: maximum,
};
// Applied before allocating input or invoking any parser. Platform
// enforcement semantics, especially macOS, require native evidence.
ensure!(
unsafe { libc::setrlimit(resource, &value) } == 0,
"Cannot install extraction resource limit"
);
Ok(())
};
#[cfg(not(target_os = "macos"))]
set_limit(libc::RLIMIT_AS, MEMORY_BYTES as libc::rlim_t)?;
set_limit(libc::RLIMIT_CPU, 30 as libc::rlim_t)
}
#[cfg(not(any(windows, unix)))]
pub fn apply() -> Result<()> {
anyhow::bail!("Extraction resource policy unsupported on this platform")
}