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

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")
}