Add memory cap to indexer to prevent OOM on large monorepos

Cargo.toml

@@ -3,7 +3,7 @@ members = ["tgrep-core", "tgrep-cli"]
 resolver = "2"
 
 [workspace.package]
-version = "0.1.20"
+version = "0.1.21"
 edition = "2024"
 license = "MIT"
 repository = "https://github.com/microsoft/tgrep"

tgrep-cli/Cargo.toml

@@ -28,6 +28,12 @@ lru = "0.16.3"
 fs2 = "0.4.3"
 globset = "0.4.18"
 
+[target.'cfg(windows)'.dependencies]
+windows-sys = { version = "0.61", features = ["Win32_System_ProcessStatus", "Win32_System_SystemInformation", "Win32_System_Threading"] }
+
+[target.'cfg(unix)'.dependencies]
+libc = "0.2"
+
 [dev-dependencies]
 assert_cmd = "2"
 tempfile = "3"

tgrep-cli/src/cpu.rs

@@ -0,0 +1,46 @@
+//! CPU budget for the background indexer.
+//!
+//! Indexing's CPU-heavy work (reading files + extracting trigrams) runs in
+//! parallel via rayon, which by default fans out across every logical core and
+//! can saturate the host. To stay a good neighbor — especially when tgrep is
+//! embedded in another tool — the indexer confines that work to a bounded
+//! worker pool sized from a CPU budget expressed as a percentage of cores.
+
+use std::thread::available_parallelism;
+
+/// Number of worker threads to use for the parallel indexing work, given a CPU
+/// budget expressed as a percentage of logical cores (1–100).
+///
+/// Always at least 1 and never more than the available core count, so a 50%
+/// budget on an 8-core host yields 4 threads, and any budget on a single-core
+/// host yields 1.
+pub fn index_thread_count(max_cpu_percent: u8) -> usize {
+    let cores = available_parallelism().map(|n| n.get()).unwrap_or(1);
+    let pct = max_cpu_percent.clamp(1, 100) as usize;
+    (cores * pct).div_ceil(100).clamp(1, cores)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn clamps_to_at_least_one_thread() {
+        // Even a 1% budget must keep indexing alive with a single worker.
+        assert!(index_thread_count(1) >= 1);
+    }
+
+    #[test]
+    fn never_exceeds_core_count() {
+        let cores = available_parallelism().map(|n| n.get()).unwrap_or(1);
+        assert!(index_thread_count(100) <= cores);
+        // Out-of-range percentages are clamped to the 1..=100 budget.
+        assert!(index_thread_count(u8::MAX) <= cores);
+    }
+
+    #[test]
+    fn full_budget_uses_all_cores() {
+        let cores = available_parallelism().map(|n| n.get()).unwrap_or(1);
+        assert_eq!(index_thread_count(100), cores);
+    }
+}

tgrep-cli/src/main.rs

@@ -5,9 +5,11 @@
 ///   tgrep serve [path]           Start the search server
 ///   tgrep <pattern> [path]       Search (auto-delegates to server)
 ///   tgrep status [path]          Show index/server status
+mod cpu;
 mod glob_filter;
 mod index;
 mod matching;
+mod mem;
 mod output;
 mod search;
 mod serve;
@@ -225,6 +227,21 @@ enum Command {
         #[arg(long)]
         no_watch: bool,
 
+        /// Maximum memory budget in megabytes for the in-memory index built
+        /// during the initial scan. When the indexer's working set exceeds
+        /// this, it flushes to disk and continues, keeping peak memory bounded
+        /// while still producing a complete index. Defaults to 50% of physical
+        /// RAM (clamped between 512 MB and 16 GB).
+        #[arg(long = "max-memory", value_name = "MB", value_parser = clap::value_parser!(u64).range(1..))]
+        max_memory_mb: Option<u64>,
+
+        /// Maximum CPU budget for the initial index build, as a percentage of
+        /// logical cores (1-100). The parallel file-reading/trigram-extraction
+        /// work is confined to this fraction of cores so the host stays
+        /// responsive. Defaults to 50%.
+        #[arg(long = "max-cpu", value_name = "PERCENT")]
+        max_cpu_percent: Option<u8>,
+
         /// Exclude directories from indexing (can be specified multiple times).
         #[arg(long = "exclude", action = clap::ArgAction::Append)]
         exclude: Vec<String>,
@@ -321,8 +338,23 @@ fn main() {
         Some(Command::Serve {
             path,
             no_watch,
+            max_memory_mb,
+            max_cpu_percent,
             exclude,
-        }) => serve::run(&path, cli.index_path.as_deref(), no_watch, &exclude),
+        }) => {
+            let memory_cap = max_memory_mb
+                .map(|mb| mb.saturating_mul(1024 * 1024))
+                .unwrap_or_else(mem::default_memory_cap_bytes);
+            let index_threads = cpu::index_thread_count(max_cpu_percent.unwrap_or(50));
+            serve::run(
+                &path,
+                cli.index_path.as_deref(),
+                no_watch,
+                &exclude,
+                memory_cap,
+                index_threads,
+            )
+        }
         Some(Command::Search {
             ref pattern,
             ref paths,

tgrep-cli/src/mem.rs

@@ -0,0 +1,170 @@
+//! Cross-platform process memory introspection.
+//!
+//! Provides functions to query the current process's resident set size (RSS)
+//! and the host's total physical memory — used to enforce an indexing memory
+//! budget so tgrep doesn't OOM-kill the host on large monorepos.
+
+/// Returns the current process's resident set size in bytes, or `None` if
+/// the platform query fails.
+#[cfg(target_os = "windows")]
+pub fn process_rss_bytes() -> Option<u64> {
+    use std::mem::MaybeUninit;
+    use windows_sys::Win32::System::ProcessStatus::GetProcessMemoryInfo;
+    use windows_sys::Win32::System::ProcessStatus::PROCESS_MEMORY_COUNTERS;
+    use windows_sys::Win32::System::Threading::GetCurrentProcess;
+
+    unsafe {
+        let handle = GetCurrentProcess();
+        let mut counters = MaybeUninit::<PROCESS_MEMORY_COUNTERS>::zeroed();
+        let size = std::mem::size_of::<PROCESS_MEMORY_COUNTERS>() as u32;
+        // GetProcessMemoryInfo requires `cb` to hold the struct size on input.
+        (*counters.as_mut_ptr()).cb = size;
+        let ok = GetProcessMemoryInfo(handle, counters.as_mut_ptr(), size);
+        if ok != 0 {
+            let counters = counters.assume_init();
+            Some(counters.WorkingSetSize as u64)
+        } else {
+            None
+        }
+    }
+}
+
+#[cfg(target_os = "windows")]
+pub fn total_physical_memory_bytes() -> Option<u64> {
+    use std::mem::MaybeUninit;
+    use windows_sys::Win32::System::SystemInformation::{GlobalMemoryStatusEx, MEMORYSTATUSEX};
+
+    unsafe {
+        let mut status = MaybeUninit::<MEMORYSTATUSEX>::zeroed();
+        (*status.as_mut_ptr()).dwLength = std::mem::size_of::<MEMORYSTATUSEX>() as u32;
+        let ok = GlobalMemoryStatusEx(status.as_mut_ptr());
+        if ok != 0 {
+            Some(status.assume_init().ullTotalPhys)
+        } else {
+            None
+        }
+    }
+}
+
+#[cfg(target_os = "linux")]
+pub fn process_rss_bytes() -> Option<u64> {
+    let statm = std::fs::read_to_string("/proc/self/statm").ok()?;
+    let rss_pages: u64 = statm.split_whitespace().nth(1)?.parse().ok()?;
+    let page_size = unsafe { libc::sysconf(libc::_SC_PAGESIZE) };
+    if page_size <= 0 {
+        return None;
+    }
+    Some(rss_pages * page_size as u64)
+}
+
+#[cfg(target_os = "linux")]
+pub fn total_physical_memory_bytes() -> Option<u64> {
+    let meminfo = std::fs::read_to_string("/proc/meminfo").ok()?;
+    for line in meminfo.lines() {
+        if let Some(rest) = line.strip_prefix("MemTotal:") {
+            let kb: u64 = rest.trim().strip_suffix("kB")?.trim().parse().ok()?;
+            return Some(kb * 1024);
+        }
+    }
+    None
+}
+
+#[cfg(target_os = "macos")]
+pub fn process_rss_bytes() -> Option<u64> {
+    use std::mem::MaybeUninit;
+    unsafe {
+        // `ru_maxrss` reports the *peak* RSS, so it never decreases after a
+        // flush reclaims memory and would keep the process looking over-budget
+        // forever. Query the *current* resident size via proc_pidinfo instead.
+        let mut info = MaybeUninit::<libc::proc_taskinfo>::zeroed();
+        let size = std::mem::size_of::<libc::proc_taskinfo>() as libc::c_int;
+        let ret = libc::proc_pidinfo(
+            libc::getpid(),
+            libc::PROC_PIDTASKINFO,
+            0,
+            info.as_mut_ptr() as *mut libc::c_void,
+            size,
+        );
+        if ret == size {
+            Some(info.assume_init().pti_resident_size)
+        } else {
+            None
+        }
+    }
+}
+
+#[cfg(target_os = "macos")]
+pub fn total_physical_memory_bytes() -> Option<u64> {
+    unsafe {
+        let mut size: u64 = 0;
+        let mut len = std::mem::size_of::<u64>();
+        let mut mib = [libc::CTL_HW, libc::HW_MEMSIZE];
+        let ret = libc::sysctl(
+            mib.as_mut_ptr(),
+            2,
+            &mut size as *mut u64 as *mut libc::c_void,
+            &mut len,
+            std::ptr::null_mut(),
+            0,
+        );
+        if ret == 0 { Some(size) } else { None }
+    }
+}
+
+// Fallback for unsupported platforms
+#[cfg(not(any(target_os = "windows", target_os = "linux", target_os = "macos")))]
+pub fn process_rss_bytes() -> Option<u64> {
+    None
+}
+
+#[cfg(not(any(target_os = "windows", target_os = "linux", target_os = "macos")))]
+pub fn total_physical_memory_bytes() -> Option<u64> {
+    None
+}
+
+/// Compute the default memory cap: 50% of physical RAM, with a floor of 512 MB
+/// and a ceiling of 16 GB. Returns bytes.
+pub fn default_memory_cap_bytes() -> u64 {
+    const FLOOR: u64 = 512 * 1024 * 1024; // 512 MB
+    const CEILING: u64 = 16 * 1024 * 1024 * 1024; // 16 GB
+
+    let half_ram = total_physical_memory_bytes()
+        .map(|total| total / 2)
+        .unwrap_or(4 * 1024 * 1024 * 1024); // fallback: 4 GB
+
+    half_ram.clamp(FLOOR, CEILING)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    // On supported platforms the queries must succeed and return non-zero
+    // values. This guards regressions like a missing `PROCESS_MEMORY_COUNTERS.cb`
+    // (or `MEMORYSTATUSEX.dwLength`) initialization, which would make the call
+    // fail and silently disable the memory cap.
+    #[cfg(any(target_os = "windows", target_os = "linux", target_os = "macos"))]
+    #[test]
+    fn process_rss_is_nonzero() {
+        let rss = process_rss_bytes().expect("process RSS query should succeed");
+        assert!(rss > 0, "process RSS should be non-zero, got {rss}");
+    }
+
+    #[cfg(any(target_os = "windows", target_os = "linux", target_os = "macos"))]
+    #[test]
+    fn total_physical_memory_is_nonzero() {
+        let total =
+            total_physical_memory_bytes().expect("total physical memory query should succeed");
+        assert!(
+            total > 0,
+            "total physical memory should be non-zero, got {total}"
+        );
+    }
+
+    #[test]
+    fn default_cap_is_within_bounds() {
+        let cap = default_memory_cap_bytes();
+        assert!(cap >= 512 * 1024 * 1024);
+        assert!(cap <= 16 * 1024 * 1024 * 1024);
+    }
+}