From 36c25962e6201f0d35fd0d77193fc45c05907eaf Mon Sep 17 00:00:00 2001 From: Shengyu Fu Date: Tue, 30 Jun 2026 08:40:07 -0700 Subject: [PATCH 01/12] Add memory cap to indexer to prevent OOM on large monorepos The background indexer (tgrep serve) now checks the process RSS after each batch and stops indexing early when it exceeds a configurable memory budget. This prevents the host from being OOM-killed on very large monorepos that produce an unbounded in-memory trigram overlay. Changes: - Add --max-memory flag to 'tgrep serve' (defaults to 50% of physical RAM, clamped between 512 MB and 16 GB) - Add cross-platform mem.rs module (Windows/Linux/macOS) for querying process RSS and total physical memory - When the cap is hit, indexing stops and the on-disk index is marked as partial (complete=false) so subsequent starts will continue indexing the remaining files incrementally Fixes: github/copilot-cli#3976 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- Cargo.lock | 2 + tgrep-cli/Cargo.toml | 6 ++ tgrep-cli/src/main.rs | 16 +++++- tgrep-cli/src/mem.rs | 127 +++++++++++++++++++++++++++++++++++++++++ tgrep-cli/src/serve.rs | 41 ++++++++++++- 5 files changed, 188 insertions(+), 4 deletions(-) create mode 100644 tgrep-cli/src/mem.rs diff --git a/Cargo.lock b/Cargo.lock index 98cf1fa..070e7a0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1045,6 +1045,7 @@ dependencies = [ "fancy-regex", "fs2", "globset", + "libc", "lru", "notify", "predicates", @@ -1054,6 +1055,7 @@ dependencies = [ "serde_json", "tempfile", "tgrep-core", + "windows-sys 0.61.2", "winresource", ] diff --git a/tgrep-cli/Cargo.toml b/tgrep-cli/Cargo.toml index bd987df..e8271ec 100644 --- a/tgrep-cli/Cargo.toml +++ b/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" diff --git a/tgrep-cli/src/main.rs b/tgrep-cli/src/main.rs index 935f881..4c1daaf 100644 --- a/tgrep-cli/src/main.rs +++ b/tgrep-cli/src/main.rs @@ -8,6 +8,7 @@ mod glob_filter; mod index; mod matching; +mod mem; mod output; mod search; mod serve; @@ -225,6 +226,13 @@ enum Command { #[arg(long)] no_watch: bool, + /// Maximum memory budget in megabytes for the indexing process. + /// When the process RSS exceeds this limit, indexing stops early + /// and produces a partial index. Defaults to 50% of physical RAM + /// (clamped between 512 MB and 16 GB). + #[arg(long = "max-memory", value_name = "MB")] + max_memory_mb: Option, + /// Exclude directories from indexing (can be specified multiple times). #[arg(long = "exclude", action = clap::ArgAction::Append)] exclude: Vec, @@ -321,8 +329,14 @@ fn main() { Some(Command::Serve { path, no_watch, + max_memory_mb, exclude, - }) => serve::run(&path, cli.index_path.as_deref(), no_watch, &exclude), + }) => { + let memory_cap = max_memory_mb + .map(|mb| mb * 1024 * 1024) + .unwrap_or_else(mem::default_memory_cap_bytes); + serve::run(&path, cli.index_path.as_deref(), no_watch, &exclude, memory_cap) + } Some(Command::Search { ref pattern, ref paths, diff --git a/tgrep-cli/src/mem.rs b/tgrep-cli/src/mem.rs new file mode 100644 index 0000000..a739495 --- /dev/null +++ b/tgrep-cli/src/mem.rs @@ -0,0 +1,127 @@ +//! 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 { + 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::::zeroed(); + let size = std::mem::size_of::() as u32; + 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 { + use std::mem::MaybeUninit; + use windows_sys::Win32::System::SystemInformation::{ + GlobalMemoryStatusEx, MEMORYSTATUSEX, + }; + + unsafe { + let mut status = MaybeUninit::::zeroed(); + (*status.as_mut_ptr()).dwLength = std::mem::size_of::() 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 { + 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 { + 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 { + use std::mem::MaybeUninit; + unsafe { + let mut rusage = MaybeUninit::::zeroed(); + if libc::getrusage(libc::RUSAGE_SELF, rusage.as_mut_ptr()) == 0 { + // On macOS, ru_maxrss is in bytes + Some(rusage.assume_init().ru_maxrss as u64) + } else { + None + } + } +} + +#[cfg(target_os = "macos")] +pub fn total_physical_memory_bytes() -> Option { + use std::mem::MaybeUninit; + unsafe { + let mut size: u64 = 0; + let mut len = std::mem::size_of::(); + let mib = [libc::CTL_HW, libc::HW_MEMSIZE]; + let ret = libc::sysctl( + mib.as_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 { + None +} + +#[cfg(not(any(target_os = "windows", target_os = "linux", target_os = "macos")))] +pub fn total_physical_memory_bytes() -> Option { + 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) +} diff --git a/tgrep-cli/src/serve.rs b/tgrep-cli/src/serve.rs index 0c0499e..962c336 100644 --- a/tgrep-cli/src/serve.rs +++ b/tgrep-cli/src/serve.rs @@ -155,6 +155,9 @@ struct ServerState { /// loading or if no matcher could be built; during that window the watcher /// falls back to hidden / exclude filtering. gitignore: RwLock>, + /// Maximum RSS budget (bytes). When the process exceeds this during + /// indexing, the build stops early and produces a partial index. + memory_cap_bytes: u64, } struct SearchOpts { @@ -171,6 +174,7 @@ pub fn run( index_path: Option<&Path>, no_watch: bool, exclude_dirs: &[String], + memory_cap_bytes: u64, ) -> Result<()> { let serve_start = Instant::now(); let root = std::fs::canonicalize(root)?; @@ -241,6 +245,7 @@ pub fn run( file_stamps: RwLock::new(tgrep_core::meta::read_filestamps(&index_dir).unwrap_or_default()), snapshot_gate: RwLock::new(()), gitignore: RwLock::new(None), + memory_cap_bytes, }); // Bind TCP listener on a random port @@ -255,10 +260,11 @@ pub fn run( info.save(&index_dir)?; eprintln!( - "[trace] serve ready in {:.1}ms. TCP on port {}. Cache: max {} entries.", + "[trace] serve ready in {:.1}ms. TCP on port {}. Cache: max {} entries. Memory cap: {} MB.", serve_start.elapsed().as_secs_f64() * 1000.0, port, - CACHE_CAPACITY + CACHE_CAPACITY, + memory_cap_bytes / (1024 * 1024), ); // If no pre-existing index, build into the LiveIndex in background @@ -1438,7 +1444,23 @@ fn background_index_build(state: &Arc, root: &Path, index_dir: &Pat ); // Phase 2: Process new files in parallel batches + let mut memory_capped = false; for (batch_idx, batch) in new_files.chunks(BATCH_SIZE).enumerate() { + // Check memory budget before processing the next batch. + if let Some(rss) = crate::mem::process_rss_bytes() + && rss > state.memory_cap_bytes + { + eprintln!( + "[trace] memory cap reached ({} MB RSS > {} MB cap) after {} batches — \ + stopping indexing early", + rss / (1024 * 1024), + state.memory_cap_bytes / (1024 * 1024), + batch_idx, + ); + memory_capped = true; + break; + } + // Parallel: read files + extract trigrams (no locks held) let batch_results: Vec<(String, Vec)> = batch .par_iter() @@ -1485,7 +1507,8 @@ fn background_index_build(state: &Arc, root: &Path, index_dir: &Pat } eprintln!( - "[trace] background indexing complete: {} total files ({} new, {} seeded) in {:.1}s", + "[trace] background indexing {}: {} total files ({} new, {} seeded) in {:.1}s", + if memory_capped { "stopped (memory cap)" } else { "complete" }, total, new_count, seeded_count, @@ -1527,6 +1550,18 @@ fn background_index_build(state: &Arc, root: &Path, index_dir: &Pat // Final (and only) flush to disk for the bulk build. eprintln!("[trace] persisting final index to disk..."); let pruned = flush_index_to_disk(state, root, index_dir, Some(&stamps)); + + // When the build was cut short by the memory cap, mark the on-disk + // index as partial so the next server start will attempt to continue + // indexing the remaining files. + if memory_capped + && pruned + && let Ok(mut meta) = tgrep_core::meta::IndexMeta::load(index_dir) + { + meta.complete = false; + let _ = meta.save(index_dir); + } + state .flushing .store(false, std::sync::atomic::Ordering::Relaxed); From 5b947e7e133daf96b802cf6e57034846112408bb Mon Sep 17 00:00:00 2001 From: Shengyu Fu Date: Tue, 30 Jun 2026 11:07:22 -0700 Subject: [PATCH 02/12] Bound peak indexing memory via flush-and-continue (complete index) Replace the 'stop at the memory cap (partial index)' behavior with a memory-bounded flush-and-continue strategy so the bulk indexer keeps peak memory under the budget AND still produces a COMPLETE index. The naive approach -- reuse the existing flush_index_to_disk mid-build -- does not bound peak memory, because HybridIndex::full_snapshot() rebuilds the entire merged inverted index in heap (O(total index size)). Key insight: during the initial bulk build the live overlay is strictly append-only (the file watcher and auto-save are both suppressed while indexing==true), so a flush can copy the existing on-disk postings verbatim from the reader's mmap and only the bounded overlay needs to live in heap. Changes: - tgrep-core: add builder::append_overlay_to_index -- a streaming 2-way append-merge of the reader's sorted lookup table (read from mmap, posting bytes copied verbatim) with the live overlay's sorted postings. Peak heap stays bounded to the overlay snapshot, independent of total index size. - tgrep-core: add IndexReader::nth_trigram_raw (zero-copy access to a trigram's raw posting bytes) and HybridIndex::reader_arc. - serve.rs: in background_index_build, when RSS exceeds the cap, perform an incremental_flush (publish complete=false) and clear the overlay, then continue. The final flush publishes complete=true. Extract the shared move+reopen+swap+prune logic into publish_staged_index, reused by flush_index_to_disk and incremental_flush. - main.rs: update --max-memory help to reflect flush-and-continue. Tests: - core unit test for append-merge correctness (file IDs, sorted postings, base-mask preservation, overlay sentinel masks). - CLI integration test: serve with --max-memory 1 over 1200 files forces multiple incremental flushes; asserts meta complete=true, all files present, and the last-indexed file is searchable. Fixes: github/copilot-cli#3976 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- tgrep-cli/src/main.rs | 9 +- tgrep-cli/src/serve.rs | 175 ++++++++++++++++++------ tgrep-cli/tests/memory_cap.rs | 195 ++++++++++++++++++++++++++ tgrep-core/src/builder.rs | 251 ++++++++++++++++++++++++++++++++++ tgrep-core/src/hybrid.rs | 9 ++ tgrep-core/src/reader.rs | 24 ++++ 6 files changed, 618 insertions(+), 45 deletions(-) create mode 100644 tgrep-cli/tests/memory_cap.rs diff --git a/tgrep-cli/src/main.rs b/tgrep-cli/src/main.rs index 4c1daaf..f7e8ab6 100644 --- a/tgrep-cli/src/main.rs +++ b/tgrep-cli/src/main.rs @@ -226,10 +226,11 @@ enum Command { #[arg(long)] no_watch: bool, - /// Maximum memory budget in megabytes for the indexing process. - /// When the process RSS exceeds this limit, indexing stops early - /// and produces a partial index. Defaults to 50% of physical RAM - /// (clamped between 512 MB and 16 GB). + /// 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")] max_memory_mb: Option, diff --git a/tgrep-cli/src/serve.rs b/tgrep-cli/src/serve.rs index 962c336..aea788d 100644 --- a/tgrep-cli/src/serve.rs +++ b/tgrep-cli/src/serve.rs @@ -1444,23 +1444,8 @@ fn background_index_build(state: &Arc, root: &Path, index_dir: &Pat ); // Phase 2: Process new files in parallel batches - let mut memory_capped = false; + let mut incremental_flushes = 0u32; for (batch_idx, batch) in new_files.chunks(BATCH_SIZE).enumerate() { - // Check memory budget before processing the next batch. - if let Some(rss) = crate::mem::process_rss_bytes() - && rss > state.memory_cap_bytes - { - eprintln!( - "[trace] memory cap reached ({} MB RSS > {} MB cap) after {} batches — \ - stopping indexing early", - rss / (1024 * 1024), - state.memory_cap_bytes / (1024 * 1024), - batch_idx, - ); - memory_capped = true; - break; - } - // Parallel: read files + extract trigrams (no locks held) let batch_results: Vec<(String, Vec)> = batch .par_iter() @@ -1504,14 +1489,42 @@ fn background_index_build(state: &Arc, root: &Path, index_dir: &Pat start.elapsed().as_secs_f64() ); } + + // Memory-bounded build: if the in-heap overlay has pushed RSS past the + // budget, persist what we've indexed so far to disk and reclaim the + // heap before continuing. This keeps peak memory bounded (the flush + // copies existing on-disk postings verbatim from mmap rather than into + // heap) while still converging to a *complete* index — unlike simply + // stopping, which would leave a partial index. + if let Some(rss) = crate::mem::process_rss_bytes() + && rss > state.memory_cap_bytes + { + eprintln!( + "[trace] memory cap reached ({} MB RSS > {} MB cap) — flushing \ + overlay to disk to reclaim memory and continuing", + rss / (1024 * 1024), + state.memory_cap_bytes / (1024 * 1024), + ); + if incremental_flush(state, index_dir) { + incremental_flushes += 1; + let mut index = state.index.write().unwrap(); + index.live.shrink_to_fit(); + } else { + eprintln!( + "[trace] warning: incremental flush did not reclaim memory; \ + continuing (build may still exceed the budget)" + ); + } + } } eprintln!( - "[trace] background indexing {}: {} total files ({} new, {} seeded) in {:.1}s", - if memory_capped { "stopped (memory cap)" } else { "complete" }, + "[trace] background indexing complete: {} total files ({} new, {} seeded, \ + {} incremental flushes) in {:.1}s", total, new_count, seeded_count, + incremental_flushes, start.elapsed().as_secs_f64() ); @@ -1547,21 +1560,13 @@ fn background_index_build(state: &Arc, root: &Path, index_dir: &Pat .indexing .store(false, std::sync::atomic::Ordering::Relaxed); - // Final (and only) flush to disk for the bulk build. + // Final flush to disk for the bulk build. This always publishes with + // `complete = true`; any intermediate incremental flushes published + // `complete = false` so a mid-build kill would resume rather than be + // treated as finished. eprintln!("[trace] persisting final index to disk..."); let pruned = flush_index_to_disk(state, root, index_dir, Some(&stamps)); - // When the build was cut short by the memory cap, mark the on-disk - // index as partial so the next server start will attempt to continue - // indexing the remaining files. - if memory_capped - && pruned - && let Ok(mut meta) = tgrep_core::meta::IndexMeta::load(index_dir) - { - meta.complete = false; - let _ = meta.save(index_dir); - } - state .flushing .store(false, std::sync::atomic::Ordering::Relaxed); @@ -1656,16 +1661,109 @@ fn flush_index_to_disk( eprintln!("[trace] warning: failed to write staging filestamps: {e}"); } - // Lock-free publish: rename staging files into place, build new reader, - // then swap. Search queries continue to be served throughout. - // + // Lock-free publish: rename staging files into place, build the new + // reader, swap it in, and prune the now-persisted overlay. Search queries + // continue to be served by the previous reader throughout. + let pruned = publish_staged_index(state, index_dir, &staging_dir, num_files); + + eprintln!( + "[trace] index flushed: {num_files} files on disk in {:.1}s", + flush_start.elapsed().as_secs_f64() + ); + pruned +} + +/// Memory-bounded incremental flush used during the initial bulk build. +/// +/// Unlike [`flush_index_to_disk`] (which builds a full reader+overlay snapshot +/// in heap via `full_snapshot`, costing O(total index size) memory), this +/// streams the live overlay onto the existing on-disk index via +/// [`builder::append_overlay_to_index`]: the existing postings are copied +/// verbatim from the reader's mmap and never enter the heap. Peak heap stays +/// bounded to the overlay snapshot, so repeated calls keep the whole build +/// under the memory budget while still converging to a complete index. +/// +/// Relies on the bulk-build invariant that the overlay is **append-only** +/// (watcher + auto-save suppressed while `indexing == true`), so every overlay +/// file is new and the merge is a pure append. +/// +/// The on-disk index is published with `complete = false`: a kill mid-build +/// must leave the index marked partial so the next start resumes indexing the +/// remaining files. The final end-of-build flush re-publishes with +/// `complete = true`. +/// +/// Returns `true` if the new reader was published and the overlay pruned. +fn incremental_flush(state: &ServerState, index_dir: &Path) -> bool { + let flush_start = Instant::now(); + + // Hold the snapshot gate for the whole snapshot → publish → prune cycle, + // mirroring flush_index_to_disk. (During the bulk build the watcher is + // already suppressed, but auto-save coordination and future-proofing make + // the gate the right call.) + let _gate = state.snapshot_gate.write().unwrap(); + + // Snapshot the overlay (bounded heap) and the current reader (cheap Arc). + let (overlay_paths, overlay_inverted, reader) = { + let index = state.index.read().unwrap(); + let (paths, inverted) = index.live.snapshot_for_disk(); + (paths, inverted, index.reader_arc()) + }; + if overlay_paths.is_empty() { + return false; + } + let num_files = reader.num_files() + overlay_paths.len(); + + let staging_dir = index_dir.with_file_name(".tgrep_flush_staging"); + let _ = std::fs::remove_dir_all(&staging_dir); + + // Stream-merge overlay onto the existing on-disk index. `complete = false` + // so an interrupted build resumes rather than being treated as finished. + if let Err(e) = builder::append_overlay_to_index( + &state.root, + &staging_dir, + &reader, + &overlay_paths, + &overlay_inverted, + false, + ) { + eprintln!("[trace] warning: incremental flush write failed: {e}"); + let _ = std::fs::remove_dir_all(&staging_dir); + return false; + } + + let pruned = publish_staged_index(state, index_dir, &staging_dir, num_files); + eprintln!( + "[trace] incremental flush: {num_files} files on disk in {:.1}s", + flush_start.elapsed().as_secs_f64() + ); + pruned +} + +/// Publish a staged index directory: move the staged files into `index_dir`, +/// reopen the on-disk reader (with Windows stale-NTFS-metadata retries), +/// validate + warm it, swap it in without blocking searches, and prune the +/// now-persisted overlay entries. +/// +/// Shared by [`flush_index_to_disk`] and [`incremental_flush`]. The +/// `publish_lock` is held across move + open + swap so concurrent publishers +/// cannot interleave renames or swap readers out of order. `num_files` is the +/// expected on-disk file count used to reject a partially-published reader. +/// +/// Returns `true` when the swap + prune succeeded, `false` on any failure (the +/// previous reader and the live overlay are retained as the fallback). +fn publish_staged_index( + state: &ServerState, + index_dir: &Path, + staging_dir: &Path, + num_files: usize, +) -> bool { // Held across move + open + swap so concurrent publishers (auto-save / // background-build / watcher reindex flush) cannot interleave renames // or swap readers out of order. Searches do not take this lock. let _publish = state.publish_lock.lock().unwrap(); - if let Err(e) = move_staged_files(&staging_dir, index_dir) { + if let Err(e) = move_staged_files(staging_dir, index_dir) { eprintln!("[trace] warning: flush move failed: {e}"); - let _ = std::fs::remove_dir_all(&staging_dir); + let _ = std::fs::remove_dir_all(staging_dir); return false; } @@ -1772,12 +1870,7 @@ fn flush_index_to_disk( } false }; - let _ = std::fs::remove_dir_all(&staging_dir); - - eprintln!( - "[trace] index flushed: {num_files} files on disk in {:.1}s", - flush_start.elapsed().as_secs_f64() - ); + let _ = std::fs::remove_dir_all(staging_dir); pruned } diff --git a/tgrep-cli/tests/memory_cap.rs b/tgrep-cli/tests/memory_cap.rs new file mode 100644 index 0000000..d5a9a9c --- /dev/null +++ b/tgrep-cli/tests/memory_cap.rs @@ -0,0 +1,195 @@ +//! Integration test: the memory-bounded indexer must still produce a +//! **complete** index. +//! +//! Starts `tgrep serve` with a tiny `--max-memory` budget so the bulk indexer +//! is forced to flush its in-memory overlay to disk repeatedly mid-build, then +//! verifies that (a) the published index ends up marked complete and (b) every +//! file — including ones indexed in the very last batch, after several +//! incremental flushes — is searchable. + +use std::fs; +use std::io::{BufRead, BufReader, Write}; +use std::net::TcpStream; +use std::path::{Path, PathBuf}; +use std::process::{Child, Command}; +use std::time::{Duration, Instant}; + +use tempfile::TempDir; + +/// Number of files to index. Comfortably exceeds the indexer's internal batch +/// size (500) so that, with a 1 MB cap, several incremental flushes fire. +const NUM_FILES: usize = 1200; + +fn tgrep_bin() -> PathBuf { + assert_cmd::cargo::cargo_bin("tgrep") +} + +/// Create a fixture where every file contains a unique, greppable token +/// (`UNIQUETOKEN`) so we can assert exact per-file coverage after the build. +fn setup_fixture() -> TempDir { + let dir = TempDir::new().unwrap(); + let src = dir.path().join("src"); + fs::create_dir_all(&src).unwrap(); + for i in 0..NUM_FILES { + let content = format!( + "fn handler_{i}() {{\n // marker UNIQUETOKEN{i}\n let value_{i} = {i};\n}}\n" + ); + fs::write(src.join(format!("mod_{i:05}.rs")), content).unwrap(); + } + dir +} + +struct ServerGuard { + child: Child, + port: u16, + index_dir: PathBuf, +} + +impl Drop for ServerGuard { + fn drop(&mut self) { + let _ = self.child.kill(); + let _ = self.child.wait(); + } +} + +fn send_request(port: u16, request: &str) -> std::io::Result { + let mut stream = TcpStream::connect(format!("127.0.0.1:{port}"))?; + stream.set_read_timeout(Some(Duration::from_secs(30)))?; + writeln!(stream, "{request}")?; + stream.flush()?; + let mut reader = BufReader::new(stream); + let mut response = String::new(); + reader.read_line(&mut response)?; + Ok(response) +} + +/// Start `tgrep serve` with a 1 MB memory cap so the indexer flushes mid-build. +fn start_capped_server(root: &Path) -> ServerGuard { + let index_dir = root.join(".tgrep_memcap_index"); + fs::create_dir_all(&index_dir).unwrap(); + + let child = Command::new(tgrep_bin()) + .args([ + "serve", + "--no-watch", + "--max-memory", + "1", // 1 MB — process baseline already exceeds this, forcing flushes + "--index-path", + index_dir.to_str().unwrap(), + root.to_str().unwrap(), + ]) + .stderr(std::process::Stdio::null()) + .stdout(std::process::Stdio::null()) + .spawn() + .expect("failed to start tgrep serve"); + + let serve_json = index_dir.join("serve.json"); + let start = Instant::now(); + let port = loop { + assert!( + start.elapsed() <= Duration::from_secs(30), + "tgrep serve did not start within 30s" + ); + if let Ok(data) = fs::read_to_string(&serve_json) + && let Ok(info) = serde_json::from_str::(&data) + && let Some(p) = info.get("port").and_then(|v| v.as_u64()) + && TcpStream::connect(format!("127.0.0.1:{p}")).is_ok() + { + break p as u16; + } + std::thread::sleep(Duration::from_millis(100)); + }; + + ServerGuard { + child, + port, + index_dir, + } +} + +/// Poll status until the background build reports it has finished indexing. +fn wait_for_indexing_done(port: u16) { + let start = Instant::now(); + loop { + assert!( + start.elapsed() <= Duration::from_secs(120), + "indexing did not finish within 120s" + ); + if let Ok(resp) = send_request(port, r#"{"jsonrpc":"2.0","method":"status","id":0}"#) + && let Ok(v) = serde_json::from_str::(&resp) + && v.pointer("/result/indexing").and_then(|v| v.as_bool()) == Some(false) + { + return; + } + std::thread::sleep(Duration::from_millis(150)); + } +} + +/// Poll the on-disk meta.json until the final flush marks the index complete. +fn wait_for_complete_meta(index_dir: &Path) -> serde_json::Value { + let meta_path = index_dir.join("meta.json"); + let start = Instant::now(); + loop { + assert!( + start.elapsed() <= Duration::from_secs(120), + "index never reached complete=true" + ); + if let Ok(data) = fs::read_to_string(&meta_path) + && let Ok(v) = serde_json::from_str::(&data) + && v.get("complete").and_then(|c| c.as_bool()) == Some(true) + { + return v; + } + std::thread::sleep(Duration::from_millis(150)); + } +} + +fn search_count(port: u16, pattern: &str) -> u64 { + let req = serde_json::json!({ + "jsonrpc": "2.0", + "method": "search", + "id": 1, + "params": { "pattern": pattern } + }) + .to_string(); + let resp = send_request(port, &req).expect("search request failed"); + let v: serde_json::Value = serde_json::from_str(&resp).expect("invalid JSON"); + assert!(v.get("error").is_none(), "search error: {v}"); + v.pointer("/result/num_matches") + .and_then(|n| n.as_u64()) + .expect("missing num_matches") +} + +#[test] +fn memory_capped_build_still_produces_complete_searchable_index() { + let dir = setup_fixture(); + let root = dir.path().join("src"); + let server = start_capped_server(&root); + let port = server.port; + + wait_for_indexing_done(port); + let meta = wait_for_complete_meta(&server.index_dir); + + // Every file made it into the on-disk index. + assert_eq!( + meta.get("num_files").and_then(|n| n.as_u64()), + Some(NUM_FILES as u64), + "complete index must contain all files; meta = {meta}" + ); + + // A token shared by every file resolves to every file — full coverage. + assert_eq!( + search_count(port, "UNIQUETOKEN"), + NUM_FILES as u64, + "all files should be searchable after a memory-capped build" + ); + + // Tokens unique to the very first and very last files (the last one indexed + // only after several incremental flushes) are both findable. + assert_eq!(search_count(port, "UNIQUETOKEN0\\b"), 1); + assert_eq!( + search_count(port, &format!("UNIQUETOKEN{}\\b", NUM_FILES - 1)), + 1, + "file indexed in the final batch (post-flush) must be searchable" + ); +} diff --git a/tgrep-core/src/builder.rs b/tgrep-core/src/builder.rs index d73e46c..f137863 100644 --- a/tgrep-core/src/builder.rs +++ b/tgrep-core/src/builder.rs @@ -7,6 +7,7 @@ use std::path::Path; use crate::Result; use crate::meta::{self, IndexMeta}; use crate::ondisk::{self, LookupEntry, PostingEntry}; +use crate::reader::IndexReader; use crate::trigram::{self, TrigramMasks}; use crate::walker; @@ -347,6 +348,166 @@ pub fn write_index_from_snapshot( ) } +/// Append a live overlay of **brand-new** files onto an existing on-disk index, +/// writing a fresh index into `out_dir` without ever materializing the existing +/// postings on the heap. +/// +/// This is the memory-bounded flush used by the bulk indexer: rather than +/// merging reader + overlay into a single in-heap `HashMap` (which costs +/// O(total index size) memory and would defeat a memory cap), it streams a +/// 2-way merge of the reader's sorted lookup table (read straight from its mmap) +/// with the overlay's sorted trigram postings. The reader's posting bytes are +/// copied **verbatim** — they have the identical on-disk layout — so peak heap +/// stays bounded to the size of the overlay snapshot plus small write buffers, +/// independent of how large the existing index already is. +/// +/// ## Append-only precondition +/// Every overlay file must be **new** (not already present in `reader`, and not +/// a deletion/supersession of a reader file). The bulk indexer guarantees this: +/// the file watcher and auto-save are both suppressed while the initial build is +/// in progress, so the overlay only ever accumulates fresh files. Under this +/// precondition the merge is a pure append: +/// - existing files keep their IDs `[0, base)`, +/// - overlay files take IDs `[base, base + overlay_paths.len())` in the order +/// given by `overlay_paths`, +/// - for any trigram, `reader_postings (ids < base) ++ overlay_postings +/// (ids >= base)` is already globally sorted by `file_id`. +/// +/// `overlay_inverted` maps each trigram to the overlay's sorted, **zero-based** +/// file indices (as produced by [`crate::live::LiveIndex::snapshot_for_disk`]); +/// each index `k` refers to `overlay_paths[k]` and is written with disk ID +/// `base + k`. Overlay entries carry the no-filter sentinel masks +/// `(u8::MAX, u8::MAX)`, matching the bulk indexer's mask-free fast path. +pub fn append_overlay_to_index( + root: &Path, + out_dir: &Path, + reader: &IndexReader, + overlay_paths: &[String], + overlay_inverted: &HashMap>, + complete: bool, +) -> Result<()> { + std::fs::create_dir_all(out_dir)?; + + let base = reader.num_files() as u32; + + // Overlay trigrams in ascending order for the 2-way merge. + let mut overlay_trigrams: Vec = overlay_inverted.keys().copied().collect(); + overlay_trigrams.sort_unstable(); + + let mut postings_file = + std::io::BufWriter::new(std::fs::File::create(out_dir.join("index.bin"))?); + let mut lookup_file = + std::io::BufWriter::new(std::fs::File::create(out_dir.join("lookup.bin"))?); + let mut lookup_scratch = + Vec::with_capacity(LOOKUP_WRITE_CHUNK_ENTRIES * ondisk::LOOKUP_ENTRY_SIZE); + let mut posting_scratch = + Vec::with_capacity(POSTING_WRITE_CHUNK_ENTRIES * ondisk::POSTING_ENTRY_SIZE); + + let reader_trigram_count = reader.num_trigrams(); + let mut ri = 0usize; + let mut oi = 0usize; + let mut offset: u64 = 0; + let mut trigram_count = 0usize; + + // Standard 2-way merge over two ascending trigram streams. + while ri < reader_trigram_count || oi < overlay_trigrams.len() { + let reader_next = (ri < reader_trigram_count) + .then(|| reader.nth_trigram_raw(ri)) + .flatten(); + let overlay_next = overlay_trigrams.get(oi).copied(); + + let (trigram, reader_bytes, overlay_seq) = match (reader_next, overlay_next) { + (Some((rt, rbytes)), Some(ot)) => match rt.cmp(&ot) { + std::cmp::Ordering::Less => { + ri += 1; + (rt, Some(rbytes), None) + } + std::cmp::Ordering::Greater => { + oi += 1; + (ot, None, overlay_inverted.get(&ot)) + } + std::cmp::Ordering::Equal => { + ri += 1; + oi += 1; + (rt, Some(rbytes), overlay_inverted.get(&rt)) + } + }, + (Some((rt, rbytes)), None) => { + ri += 1; + (rt, Some(rbytes), None) + } + (None, Some(ot)) => { + oi += 1; + (ot, None, overlay_inverted.get(&ot)) + } + // A reader entry whose raw bytes could not be read (truncated mmap). + // Skip it rather than emit a corrupt posting range. + (None, None) => break, + }; + + let reader_len = reader_bytes.map_or(0, |b| b.len() / ondisk::POSTING_ENTRY_SIZE); + let overlay_len = overlay_seq.map_or(0, |v| v.len()); + let length = (reader_len + overlay_len) as u32; + if length == 0 { + continue; + } + + write_lookup_entry( + &mut lookup_file, + LookupEntry { + trigram, + offset, + length, + }, + &mut lookup_scratch, + )?; + + // Reader postings: copy the on-disk bytes verbatim (zero decode). + if let Some(rbytes) = reader_bytes { + postings_file.write_all(rbytes)?; + } + // Overlay postings: encode with sentinel masks, IDs offset by `base`. + if let Some(seq) = overlay_seq { + for chunk in seq.chunks(POSTING_WRITE_CHUNK_ENTRIES) { + posting_scratch.clear(); + for &k in chunk { + let entry = PostingEntry { + file_id: base + k, + loc_mask: u8::MAX, + next_mask: u8::MAX, + }; + posting_scratch.extend_from_slice(&entry.encode()); + } + postings_file.write_all(&posting_scratch)?; + } + } + + offset += length as u64 * ondisk::POSTING_ENTRY_SIZE as u64; + trigram_count += 1; + } + + flush_lookup_entries(&mut lookup_file, &mut lookup_scratch)?; + postings_file.flush()?; + lookup_file.flush()?; + + // files.bin + meta.json: existing files keep IDs [0, base), overlay files + // follow at [base, base + N). Reader paths stream from its already-loaded + // file table; overlay paths from the snapshot. + let paths = reader + .all_paths() + .iter() + .map(String::as_str) + .chain(overlay_paths.iter().map(String::as_str)); + write_files_and_meta( + out_dir, + root, + base as usize + overlay_paths.len(), + paths, + trigram_count, + Some(complete), + ) +} + fn write_index_v2_from_postings( index_dir: &Path, root: &Path, @@ -423,4 +584,94 @@ mod tests { needle_paths.sort_unstable(); assert_eq!(needle_paths, vec!["src/a.txt", "src/b.txt"]); } + + #[test] + fn append_overlay_merges_new_files_into_complete_index() { + use crate::live::LiveIndex; + + // Base index with two files. + let repo = tempfile::tempdir().unwrap(); + let src = repo.path().join("src"); + std::fs::create_dir_all(&src).unwrap(); + std::fs::write(src.join("a.txt"), "hello world\nneedle one\n").unwrap(); + std::fs::write(src.join("b.txt"), "needle two\nother content\n").unwrap(); + + let base_dir = tempfile::tempdir().unwrap(); + build_index(repo.path(), Some(base_dir.path()), false, &[]).unwrap(); + let base_reader = IndexReader::open(base_dir.path()).unwrap(); + assert_eq!(base_reader.num_files(), 2); + + // Build a live overlay of two brand-new files (append-only invariant). + let mut live = LiveIndex::new(); + live.upsert_file_with_trigrams( + "src/c.txt", + crate::trigram::extract(b"needle three\n"), + ); + live.upsert_file_with_trigrams("src/d.txt", crate::trigram::extract(b"zzz unique\n")); + let (overlay_paths, overlay_inverted) = live.snapshot_for_disk(); + + // Stream-merge overlay onto the base index into a fresh dir. + let merged_dir = tempfile::tempdir().unwrap(); + append_overlay_to_index( + repo.path(), + merged_dir.path(), + &base_reader, + &overlay_paths, + &overlay_inverted, + true, + ) + .unwrap(); + + let merged = IndexReader::open(merged_dir.path()).unwrap(); + merged.validate_lookup().unwrap(); + assert_eq!(merged.num_files(), 4, "all base + overlay files present"); + + // Base file IDs are preserved (copied verbatim at the front). + assert_eq!(merged.file_path(0), base_reader.file_path(0)); + assert_eq!(merged.file_path(1), base_reader.file_path(1)); + // Overlay files follow in insertion order. + assert_eq!(merged.file_path(2), Some("src/c.txt")); + assert_eq!(merged.file_path(3), Some("src/d.txt")); + + // A trigram shared by base + overlay returns all three files, sorted. + let needle = crate::trigram::hash(b'n', b'e', b'e'); + let mut needle_paths: Vec<&str> = merged + .lookup_trigram(needle) + .iter() + .filter_map(|&id| merged.file_path(id)) + .collect(); + needle_paths.sort_unstable(); + assert_eq!(needle_paths, vec!["src/a.txt", "src/b.txt", "src/c.txt"]); + + // An overlay-only trigram resolves to just the overlay file. + let uni = crate::trigram::hash(b'u', b'n', b'i'); + let uni_paths: Vec<&str> = merged + .lookup_trigram(uni) + .iter() + .filter_map(|&id| merged.file_path(id)) + .collect(); + assert_eq!(uni_paths, vec!["src/d.txt"]); + + // Posting lists stay globally sorted by file_id after the merge. + let needle_ids = merged.lookup_trigram(needle); + let mut sorted = needle_ids.clone(); + sorted.sort_unstable(); + assert_eq!(needle_ids, sorted, "merged posting list must be sorted"); + + // Masks: base entries keep their real masks; overlay entries carry the + // no-filter sentinel (the bulk path stores no masks). + let c_id = (0..merged.num_files() as u32) + .find(|&id| merged.file_path(id) == Some("src/c.txt")) + .unwrap(); + let entries = merged.lookup_trigram_with_masks(needle); + let c_entry = entries.iter().find(|e| e.file_id == c_id).unwrap(); + assert_eq!(c_entry.loc_mask, u8::MAX); + assert_eq!(c_entry.next_mask, u8::MAX); + let base_entry = entries.iter().find(|e| e.file_id < 2).unwrap(); + assert_ne!( + base_entry.loc_mask, + u8::MAX, + "base file's real loc_mask must be preserved verbatim" + ); + } } diff --git a/tgrep-core/src/hybrid.rs b/tgrep-core/src/hybrid.rs index e63de39..e81198d 100644 --- a/tgrep-core/src/hybrid.rs +++ b/tgrep-core/src/hybrid.rs @@ -56,6 +56,15 @@ impl HybridIndex { Arc::clone(&self.reader.read().unwrap()) } + /// Public snapshot of the current on-disk reader (cheap `Arc` clone). + /// + /// Exposed so callers performing a memory-bounded incremental flush can + /// stream-merge the live overlay onto the existing on-disk index without + /// materializing the reader's postings on the heap. + pub fn reader_arc(&self) -> Arc { + self.reader() + } + /// Atomically replace the on-disk reader with `new_reader`. /// /// Takes `&self` (not `&mut self`) so that callers can perform the swap diff --git a/tgrep-core/src/reader.rs b/tgrep-core/src/reader.rs index 4bd2473..56f2d45 100644 --- a/tgrep-core/src/reader.rs +++ b/tgrep-core/src/reader.rs @@ -311,6 +311,30 @@ impl IndexReader { None } + /// Return the `i`-th trigram (in ascending sorted order) together with the + /// raw, on-disk posting bytes for that trigram. Zero-copy: the returned + /// slice points directly into the mmap, so callers can copy the bytes + /// verbatim into a new index without decoding them into heap. + /// + /// Used by the streaming append-merge in the builder to keep the existing + /// on-disk postings out of heap while merging a live overlay into a new + /// index. Returns `None` if `i` is out of range or the postings mmap is + /// absent/truncated for the entry. + pub fn nth_trigram_raw(&self, i: usize) -> Option<(u32, &[u8])> { + if i >= self.num_entries { + return None; + } + let entry = self.read_lookup_entry(i); + let postings = self.postings.as_ref()?; + let start = entry.offset as usize; + let byte_len = entry.length as usize * POSTING_ENTRY_SIZE; + let end = start.checked_add(byte_len)?; + if end > postings.len() { + return None; + } + Some((entry.trigram, &postings[start..end])) + } + fn read_lookup_entry(&self, index: usize) -> LookupEntry { let lookup = self.lookup.as_ref().unwrap(); let start = index * LOOKUP_ENTRY_SIZE; From 88ffcd33c4c237c25f794804047c82b7bcc6e30c Mon Sep 17 00:00:00 2001 From: Shengyu Fu Date: Tue, 30 Jun 2026 12:54:45 -0700 Subject: [PATCH 03/12] Cap CPU usage during the initial index build MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The initial bulk build's CPU-heavy work (file reads + trigram extraction) ran via rayon across every logical core, saturating the host. Confine it to a bounded worker pool so tgrep stays a good neighbor — especially when embedded in another tool. - Add --max-cpu to 'tgrep serve' (percentage of logical cores, default 50%). - New cpu.rs: index_thread_count(percent) maps the budget to a worker count clamped to 1..=cores. - background_index_build runs the parallel extraction par_iter inside a dedicated rayon ThreadPool sized to the budget (pool.install), falling back to the global pool if one can't be built. The single-threaded merge/flush paths need no pooling. Tests: cpu:: unit tests for clamping/bounds; the memory_cap integration test exercises the bounded pool end-to-end at the default 50%. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- tgrep-cli/src/cpu.rs | 46 +++++++++++++++++++++++ tgrep-cli/src/main.rs | 19 +++++++++- tgrep-cli/src/serve.rs | 84 +++++++++++++++++++++++++++++------------- 3 files changed, 122 insertions(+), 27 deletions(-) create mode 100644 tgrep-cli/src/cpu.rs diff --git a/tgrep-cli/src/cpu.rs b/tgrep-cli/src/cpu.rs new file mode 100644 index 0000000..d7e83d2 --- /dev/null +++ b/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); + } +} diff --git a/tgrep-cli/src/main.rs b/tgrep-cli/src/main.rs index f7e8ab6..58fccf7 100644 --- a/tgrep-cli/src/main.rs +++ b/tgrep-cli/src/main.rs @@ -5,6 +5,7 @@ /// tgrep serve [path] Start the search server /// tgrep [path] Search (auto-delegates to server) /// tgrep status [path] Show index/server status +mod cpu; mod glob_filter; mod index; mod matching; @@ -234,6 +235,13 @@ enum Command { #[arg(long = "max-memory", value_name = "MB")] max_memory_mb: Option, + /// 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, + /// Exclude directories from indexing (can be specified multiple times). #[arg(long = "exclude", action = clap::ArgAction::Append)] exclude: Vec, @@ -331,12 +339,21 @@ fn main() { path, no_watch, max_memory_mb, + max_cpu_percent, exclude, }) => { let memory_cap = max_memory_mb .map(|mb| mb * 1024 * 1024) .unwrap_or_else(mem::default_memory_cap_bytes); - serve::run(&path, cli.index_path.as_deref(), no_watch, &exclude, memory_cap) + 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, diff --git a/tgrep-cli/src/serve.rs b/tgrep-cli/src/serve.rs index aea788d..cc21570 100644 --- a/tgrep-cli/src/serve.rs +++ b/tgrep-cli/src/serve.rs @@ -155,9 +155,13 @@ struct ServerState { /// loading or if no matcher could be built; during that window the watcher /// falls back to hidden / exclude filtering. gitignore: RwLock>, - /// Maximum RSS budget (bytes). When the process exceeds this during - /// indexing, the build stops early and produces a partial index. + /// Maximum RSS budget (bytes). When the process exceeds this during the + /// initial build, the indexer flushes the overlay to disk and continues so + /// peak memory stays bounded while still producing a complete index. memory_cap_bytes: u64, + /// Number of worker threads used for the parallel file-reading/trigram + /// extraction during the initial build. Caps indexing CPU usage. + index_threads: usize, } struct SearchOpts { @@ -175,6 +179,7 @@ pub fn run( no_watch: bool, exclude_dirs: &[String], memory_cap_bytes: u64, + index_threads: usize, ) -> Result<()> { let serve_start = Instant::now(); let root = std::fs::canonicalize(root)?; @@ -246,6 +251,7 @@ pub fn run( snapshot_gate: RwLock::new(()), gitignore: RwLock::new(None), memory_cap_bytes, + index_threads, }); // Bind TCP listener on a random port @@ -260,11 +266,13 @@ pub fn run( info.save(&index_dir)?; eprintln!( - "[trace] serve ready in {:.1}ms. TCP on port {}. Cache: max {} entries. Memory cap: {} MB.", + "[trace] serve ready in {:.1}ms. TCP on port {}. Cache: max {} entries. \ + Memory cap: {} MB. Index threads: {}.", serve_start.elapsed().as_secs_f64() * 1000.0, port, CACHE_CAPACITY, memory_cap_bytes / (1024 * 1024), + index_threads, ); // If no pre-existing index, build into the LiveIndex in background @@ -1443,31 +1451,55 @@ fn background_index_build(state: &Arc, root: &Path, index_dir: &Pat t_walk.elapsed().as_secs_f64() * 1000.0 ); - // Phase 2: Process new files in parallel batches + // Phase 2: Process new files in parallel batches. + // + // Confine the CPU-heavy file-read + trigram-extraction work to a bounded + // worker pool (sized from the `--max-cpu` budget) so the initial build + // doesn't saturate every core and starve the host. Falls back to the + // global rayon pool if a dedicated pool can't be built. + let index_pool = rayon::ThreadPoolBuilder::new() + .num_threads(state.index_threads) + .thread_name(|i| format!("tgrep-index-{i}")) + .build() + .ok(); + if index_pool.is_some() { + eprintln!( + "[trace] indexing with {} worker thread(s)", + state.index_threads + ); + } + let mut incremental_flushes = 0u32; for (batch_idx, batch) in new_files.chunks(BATCH_SIZE).enumerate() { - // Parallel: read files + extract trigrams (no locks held) - let batch_results: Vec<(String, Vec)> = batch - .par_iter() - .filter_map(|path| { - let data = std::fs::read(path).ok()?; - if tgrep_core::trigram::is_binary(&data) { - return None; - } - let rel_path = path - .strip_prefix(root) - .unwrap_or(path) - .to_string_lossy() - .replace('\\', "/"); - - let mut trigrams = tgrep_core::trigram::extract(&data); - let lower = data.to_ascii_lowercase(); - if lower != data { - trigrams.extend(tgrep_core::trigram::extract(&lower)); - } - Some((rel_path, trigrams)) - }) - .collect(); + // Parallel: read files + extract trigrams (no locks held). Run inside + // the bounded pool when available so indexing CPU stays capped. + let extract = || { + batch + .par_iter() + .filter_map(|path| { + let data = std::fs::read(path).ok()?; + if tgrep_core::trigram::is_binary(&data) { + return None; + } + let rel_path = path + .strip_prefix(root) + .unwrap_or(path) + .to_string_lossy() + .replace('\\', "/"); + + let mut trigrams = tgrep_core::trigram::extract(&data); + let lower = data.to_ascii_lowercase(); + if lower != data { + trigrams.extend(tgrep_core::trigram::extract(&lower)); + } + Some((rel_path, trigrams)) + }) + .collect::)>>() + }; + let batch_results: Vec<(String, Vec)> = match &index_pool { + Some(pool) => pool.install(extract), + None => extract(), + }; // Sequential: insert into LiveIndex (brief write lock per batch) { From b4899c65289bb90d41c15bcf3fc42a1e3eaf9e15 Mon Sep 17 00:00:00 2001 From: Shengyu Fu Date: Tue, 30 Jun 2026 13:06:04 -0700 Subject: [PATCH 04/12] Fix memory-bounded bulk index flush coordination Review found two correctness issues in the resource-cap path: - The indexing-to-flushing handoff relied on relaxed atomics even though the auto-save loop must not observe both flags as false. Use sequentially consistent accesses for the indexing/flushing coordination flags. - The final bulk-build publish still used the full_snapshot flush path, which materializes the entire reader+overlay inverted index in heap and defeats the memory cap. Reuse the append-only streaming flush for the final complete publish as well, including staged filestamps. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- tgrep-cli/src/serve.rs | 94 +++++++++++++++++++++++------------------- 1 file changed, 51 insertions(+), 43 deletions(-) diff --git a/tgrep-cli/src/serve.rs b/tgrep-cli/src/serve.rs index cc21570..b19c54e 100644 --- a/tgrep-cli/src/serve.rs +++ b/tgrep-cli/src/serve.rs @@ -7,6 +7,7 @@ use std::io::{BufRead, BufReader, Write}; use std::net::{TcpListener, TcpStream}; use std::num::NonZeroUsize; use std::path::{Path, PathBuf}; +use std::sync::atomic::Ordering; use std::sync::{Arc, Mutex, RwLock}; use std::thread; use std::time::{Duration, Instant}; @@ -90,7 +91,9 @@ fn try_acquire_server_lock(index_dir: &Path) -> Result { /// 5. `cache` — guards the file content LRU cache /// 6. `file_stamps` — guards per-file mtime/size stamps /// -/// `flushing` is an AtomicBool, not a lock — no ordering constraint. +/// `indexing` and `flushing` coordinate the handoff between bulk indexing and +/// final flush with the auto-save loop; use sequentially consistent accesses +/// for those flags so auto-save never observes both as false during handoff. /// Searches only acquire `index` (read) and `cache` (read then write); /// they never take `snapshot_gate` or `publish_lock`. struct ServerState { @@ -777,7 +780,7 @@ fn search_file_matches( fn handle_status(id: Option, state: &ServerState) -> String { let index = state.index.read().unwrap(); let cache = state.cache.read().unwrap(); - let indexing = state.indexing.load(std::sync::atomic::Ordering::Relaxed); + let indexing = state.indexing.load(Ordering::SeqCst); let result = serde_json::json!({ "num_files": index.num_files(), @@ -920,7 +923,7 @@ fn handle_fs_event(state: &ServerState, root: &Path, event: &Event) { // Skip file events while the initial background index build is in progress — // the indexer will pick up all files itself, and the watcher would just // cause duplicate reindex work. - if state.indexing.load(std::sync::atomic::Ordering::Relaxed) { + if state.indexing.load(Ordering::SeqCst) { return; } @@ -1062,9 +1065,7 @@ fn auto_save_loop(state: Arc, index_dir: &Path) { // active — those paths handle their own publication and an // auto-save fired in parallel would just snapshot the same // overlay redundantly. - if state.indexing.load(std::sync::atomic::Ordering::Relaxed) - || state.flushing.load(std::sync::atomic::Ordering::Relaxed) - { + if state.indexing.load(Ordering::SeqCst) || state.flushing.load(Ordering::SeqCst) { continue; } @@ -1345,13 +1346,9 @@ fn background_refresh_stale(state: &Arc, root: &Path, index_dir: &P ) }) .collect(); - state - .flushing - .store(true, std::sync::atomic::Ordering::Relaxed); + state.flushing.store(true, Ordering::SeqCst); flush_index_to_disk(state, root, index_dir, Some(&new_stamps)); - state - .flushing - .store(false, std::sync::atomic::Ordering::Relaxed); + state.flushing.store(false, Ordering::SeqCst); // Refresh in-memory stamps so the watcher can dedupe spurious notify // events for files that already match what we just published. @@ -1537,7 +1534,7 @@ fn background_index_build(state: &Arc, root: &Path, index_dir: &Pat rss / (1024 * 1024), state.memory_cap_bytes / (1024 * 1024), ); - if incremental_flush(state, index_dir) { + if flush_append_only_overlay(state, index_dir, false, None) { incremental_flushes += 1; let mut index = state.index.write().unwrap(); index.live.shrink_to_fit(); @@ -1585,23 +1582,19 @@ fn background_index_build(state: &Arc, root: &Path, index_dir: &Pat // repos. Set `flushing` *before* clearing `indexing` so the auto-save // loop never observes both flags as false during the handoff and // races us into a redundant parallel snapshot of the bulk overlay. - state - .flushing - .store(true, std::sync::atomic::Ordering::Relaxed); - state - .indexing - .store(false, std::sync::atomic::Ordering::Relaxed); - - // Final flush to disk for the bulk build. This always publishes with - // `complete = true`; any intermediate incremental flushes published - // `complete = false` so a mid-build kill would resume rather than be - // treated as finished. + state.flushing.store(true, Ordering::SeqCst); + state.indexing.store(false, Ordering::SeqCst); + + // Final flush to disk for the bulk build. Use the same streaming + // append-only path as incremental flushes so the final complete publish + // does not materialize the whole reader+overlay in heap and violate the + // memory cap. This always publishes with `complete = true`; any + // intermediate incremental flushes published `complete = false` so a + // mid-build kill would resume rather than be treated as finished. eprintln!("[trace] persisting final index to disk..."); - let pruned = flush_index_to_disk(state, root, index_dir, Some(&stamps)); + let pruned = flush_append_only_overlay(state, index_dir, true, Some(&stamps)); - state - .flushing - .store(false, std::sync::atomic::Ordering::Relaxed); + state.flushing.store(false, Ordering::SeqCst); // Refresh the in-memory file_stamps so the file watcher can recognize // unchanged files and skip spurious notify events (e.g. atime/attribute @@ -1705,27 +1698,32 @@ fn flush_index_to_disk( pruned } -/// Memory-bounded incremental flush used during the initial bulk build. +/// Memory-bounded append-only flush used during the initial bulk build. /// /// Unlike [`flush_index_to_disk`] (which builds a full reader+overlay snapshot /// in heap via `full_snapshot`, costing O(total index size) memory), this /// streams the live overlay onto the existing on-disk index via /// [`builder::append_overlay_to_index`]: the existing postings are copied /// verbatim from the reader's mmap and never enter the heap. Peak heap stays -/// bounded to the overlay snapshot, so repeated calls keep the whole build -/// under the memory budget while still converging to a complete index. +/// bounded to the overlay snapshot, so repeated checkpoint flushes and the +/// final complete publish keep the whole build under the memory budget. /// /// Relies on the bulk-build invariant that the overlay is **append-only** /// (watcher + auto-save suppressed while `indexing == true`), so every overlay /// file is new and the merge is a pure append. /// -/// The on-disk index is published with `complete = false`: a kill mid-build -/// must leave the index marked partial so the next start resumes indexing the -/// remaining files. The final end-of-build flush re-publishes with -/// `complete = true`. +/// Checkpoint flushes pass `complete = false`: a kill mid-build must leave the +/// index marked partial so the next start resumes indexing the remaining files. +/// The final end-of-build flush passes `complete = true` and may pass file +/// stamps to publish alongside the index. /// /// Returns `true` if the new reader was published and the overlay pruned. -fn incremental_flush(state: &ServerState, index_dir: &Path) -> bool { +fn flush_append_only_overlay( + state: &ServerState, + index_dir: &Path, + complete: bool, + stamps: Option<&std::collections::HashMap>, +) -> bool { let flush_start = Instant::now(); // Hold the snapshot gate for the whole snapshot → publish → prune cycle, @@ -1740,7 +1738,7 @@ fn incremental_flush(state: &ServerState, index_dir: &Path) -> bool { let (paths, inverted) = index.live.snapshot_for_disk(); (paths, inverted, index.reader_arc()) }; - if overlay_paths.is_empty() { + if overlay_paths.is_empty() && !complete && stamps.is_none() { return false; } let num_files = reader.num_files() + overlay_paths.len(); @@ -1748,24 +1746,34 @@ fn incremental_flush(state: &ServerState, index_dir: &Path) -> bool { let staging_dir = index_dir.with_file_name(".tgrep_flush_staging"); let _ = std::fs::remove_dir_all(&staging_dir); - // Stream-merge overlay onto the existing on-disk index. `complete = false` - // so an interrupted build resumes rather than being treated as finished. + // Stream-merge overlay onto the existing on-disk index. Incremental + // checkpoint flushes publish `complete = false`; the final bulk-build flush + // republishes the same stream with `complete = true` and stamps. if let Err(e) = builder::append_overlay_to_index( &state.root, &staging_dir, &reader, &overlay_paths, &overlay_inverted, - false, + complete, ) { - eprintln!("[trace] warning: incremental flush write failed: {e}"); + eprintln!("[trace] warning: append-only flush write failed: {e}"); let _ = std::fs::remove_dir_all(&staging_dir); return false; } + // Stage filestamps alongside the final complete index. If this fails we + // still publish the index: losing incremental stale-check state on next + // start is preferable to dropping the completed build. + if let Some(stamps) = stamps + && let Err(e) = tgrep_core::meta::write_filestamps(stamps, &staging_dir) + { + eprintln!("[trace] warning: failed to write staging filestamps: {e}"); + } + let pruned = publish_staged_index(state, index_dir, &staging_dir, num_files); eprintln!( - "[trace] incremental flush: {num_files} files on disk in {:.1}s", + "[trace] append-only flush: {num_files} files on disk (complete={complete}) in {:.1}s", flush_start.elapsed().as_secs_f64() ); pruned @@ -1776,7 +1784,7 @@ fn incremental_flush(state: &ServerState, index_dir: &Path) -> bool { /// validate + warm it, swap it in without blocking searches, and prune the /// now-persisted overlay entries. /// -/// Shared by [`flush_index_to_disk`] and [`incremental_flush`]. The +/// Shared by [`flush_index_to_disk`] and [`flush_append_only_overlay`]. The /// `publish_lock` is held across move + open + swap so concurrent publishers /// cannot interleave renames or swap readers out of order. `num_files` is the /// expected on-disk file count used to reject a partially-published reader. From 876c0413c19f5273e0029eb9d5ca97f11f8099fc Mon Sep 17 00:00:00 2001 From: Shengyu Fu Date: Tue, 30 Jun 2026 16:36:58 -0700 Subject: [PATCH 05/12] Address review: skip unreadable reader entries, guard overflow, use current RSS on macOS - builder: advance past truncated/unreadable reader trigrams during the append-overlay merge so remaining reader entries are never dropped - reader: use checked_mul for posting byte length to avoid 32-bit/corrupt overflow - mem: report current resident size on macOS via proc_pidinfo instead of peak ru_maxrss Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- tgrep-cli/src/mem.rs | 262 ++++++++++++++++++++------------------ tgrep-core/src/builder.rs | 19 ++- tgrep-core/src/reader.rs | 2 +- 3 files changed, 149 insertions(+), 134 deletions(-) diff --git a/tgrep-cli/src/mem.rs b/tgrep-cli/src/mem.rs index a739495..f092270 100644 --- a/tgrep-cli/src/mem.rs +++ b/tgrep-cli/src/mem.rs @@ -1,127 +1,135 @@ -//! 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 { - 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::::zeroed(); - let size = std::mem::size_of::() as u32; - 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 { - use std::mem::MaybeUninit; - use windows_sys::Win32::System::SystemInformation::{ - GlobalMemoryStatusEx, MEMORYSTATUSEX, - }; - - unsafe { - let mut status = MaybeUninit::::zeroed(); - (*status.as_mut_ptr()).dwLength = std::mem::size_of::() 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 { - 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 { - 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 { - use std::mem::MaybeUninit; - unsafe { - let mut rusage = MaybeUninit::::zeroed(); - if libc::getrusage(libc::RUSAGE_SELF, rusage.as_mut_ptr()) == 0 { - // On macOS, ru_maxrss is in bytes - Some(rusage.assume_init().ru_maxrss as u64) - } else { - None - } - } -} - -#[cfg(target_os = "macos")] -pub fn total_physical_memory_bytes() -> Option { - use std::mem::MaybeUninit; - unsafe { - let mut size: u64 = 0; - let mut len = std::mem::size_of::(); - let mib = [libc::CTL_HW, libc::HW_MEMSIZE]; - let ret = libc::sysctl( - mib.as_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 { - None -} - -#[cfg(not(any(target_os = "windows", target_os = "linux", target_os = "macos")))] -pub fn total_physical_memory_bytes() -> Option { - 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) -} +//! 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 { + 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::::zeroed(); + let size = std::mem::size_of::() as u32; + 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 { + use std::mem::MaybeUninit; + use windows_sys::Win32::System::SystemInformation::{GlobalMemoryStatusEx, MEMORYSTATUSEX}; + + unsafe { + let mut status = MaybeUninit::::zeroed(); + (*status.as_mut_ptr()).dwLength = std::mem::size_of::() 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 { + 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 { + 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 { + 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::::zeroed(); + let size = std::mem::size_of::() 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 { + use std::mem::MaybeUninit; + unsafe { + let mut size: u64 = 0; + let mut len = std::mem::size_of::(); + let mib = [libc::CTL_HW, libc::HW_MEMSIZE]; + let ret = libc::sysctl( + mib.as_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 { + None +} + +#[cfg(not(any(target_os = "windows", target_os = "linux", target_os = "macos")))] +pub fn total_physical_memory_bytes() -> Option { + 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) +} diff --git a/tgrep-core/src/builder.rs b/tgrep-core/src/builder.rs index f137863..b8301f2 100644 --- a/tgrep-core/src/builder.rs +++ b/tgrep-core/src/builder.rs @@ -414,6 +414,17 @@ pub fn append_overlay_to_index( let reader_next = (ri < reader_trigram_count) .then(|| reader.nth_trigram_raw(ri)) .flatten(); + + // A reader entry that exists in the lookup table but whose raw posting + // bytes can't be read (truncated/corrupt mmap) yields `None` here while + // `ri` is still in range. Advance past it so we never stall consuming + // overlay keys against a dead reader slot, and never break the merge + // early while reader trigrams remain. + if ri < reader_trigram_count && reader_next.is_none() { + ri += 1; + continue; + } + let overlay_next = overlay_trigrams.get(oi).copied(); let (trigram, reader_bytes, overlay_seq) = match (reader_next, overlay_next) { @@ -440,8 +451,7 @@ pub fn append_overlay_to_index( oi += 1; (ot, None, overlay_inverted.get(&ot)) } - // A reader entry whose raw bytes could not be read (truncated mmap). - // Skip it rather than emit a corrupt posting range. + // Both streams exhausted (reader entries already skipped above). (None, None) => break, }; @@ -603,10 +613,7 @@ mod tests { // Build a live overlay of two brand-new files (append-only invariant). let mut live = LiveIndex::new(); - live.upsert_file_with_trigrams( - "src/c.txt", - crate::trigram::extract(b"needle three\n"), - ); + live.upsert_file_with_trigrams("src/c.txt", crate::trigram::extract(b"needle three\n")); live.upsert_file_with_trigrams("src/d.txt", crate::trigram::extract(b"zzz unique\n")); let (overlay_paths, overlay_inverted) = live.snapshot_for_disk(); diff --git a/tgrep-core/src/reader.rs b/tgrep-core/src/reader.rs index 56f2d45..4abddec 100644 --- a/tgrep-core/src/reader.rs +++ b/tgrep-core/src/reader.rs @@ -327,7 +327,7 @@ impl IndexReader { let entry = self.read_lookup_entry(i); let postings = self.postings.as_ref()?; let start = entry.offset as usize; - let byte_len = entry.length as usize * POSTING_ENTRY_SIZE; + let byte_len = (entry.length as usize).checked_mul(POSTING_ENTRY_SIZE)?; let end = start.checked_add(byte_len)?; if end > postings.len() { return None; From db8d36f57e962a3de331c7dd3f6242983fe7b7c5 Mon Sep 17 00:00:00 2001 From: Shengyu Fu Date: Tue, 30 Jun 2026 16:56:56 -0700 Subject: [PATCH 06/12] Fix macOS sysctl mib pointer mutability and drop unused import sysctl expects a mutable name pointer; use a mut mib array with as_mut_ptr. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- tgrep-cli/src/mem.rs | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tgrep-cli/src/mem.rs b/tgrep-cli/src/mem.rs index f092270..00db062 100644 --- a/tgrep-cli/src/mem.rs +++ b/tgrep-cli/src/mem.rs @@ -93,13 +93,12 @@ pub fn process_rss_bytes() -> Option { #[cfg(target_os = "macos")] pub fn total_physical_memory_bytes() -> Option { - use std::mem::MaybeUninit; unsafe { let mut size: u64 = 0; let mut len = std::mem::size_of::(); - let mib = [libc::CTL_HW, libc::HW_MEMSIZE]; + let mut mib = [libc::CTL_HW, libc::HW_MEMSIZE]; let ret = libc::sysctl( - mib.as_ptr(), + mib.as_mut_ptr(), 2, &mut size as *mut u64 as *mut libc::c_void, &mut len, From 0214ff94757efbad3ab43871a8d14e35c0b3b85c Mon Sep 17 00:00:00 2001 From: Shengyu Fu Date: Tue, 30 Jun 2026 17:27:10 -0700 Subject: [PATCH 07/12] Harden append-overlay merge and CLI memory flag against overflow/corruption - builder: fail the flush (IndexCorrupted) on an unreadable in-range reader entry instead of silently dropping trigrams; serve keeps previous reader + live overlay as fallback - builder: reject >u32::MAX reader files, and use checked add + try_from for posting-list length and overlay file-id math - main: reject --max-memory 0 via clap range and use saturating_mul when converting MB to bytes Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- tgrep-cli/src/main.rs | 4 ++-- tgrep-core/src/builder.rs | 41 +++++++++++++++++++++++++++++---------- 2 files changed, 33 insertions(+), 12 deletions(-) diff --git a/tgrep-cli/src/main.rs b/tgrep-cli/src/main.rs index 58fccf7..5fc37ff 100644 --- a/tgrep-cli/src/main.rs +++ b/tgrep-cli/src/main.rs @@ -232,7 +232,7 @@ enum Command { /// 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")] + #[arg(long = "max-memory", value_name = "MB", value_parser = clap::value_parser!(u64).range(1..))] max_memory_mb: Option, /// Maximum CPU budget for the initial index build, as a percentage of @@ -343,7 +343,7 @@ fn main() { exclude, }) => { let memory_cap = max_memory_mb - .map(|mb| mb * 1024 * 1024) + .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( diff --git a/tgrep-core/src/builder.rs b/tgrep-core/src/builder.rs index b8301f2..f70586f 100644 --- a/tgrep-core/src/builder.rs +++ b/tgrep-core/src/builder.rs @@ -4,12 +4,12 @@ use std::collections::HashMap; use std::io::Write; use std::path::Path; -use crate::Result; use crate::meta::{self, IndexMeta}; use crate::ondisk::{self, LookupEntry, PostingEntry}; use crate::reader::IndexReader; use crate::trigram::{self, TrigramMasks}; use crate::walker; +use crate::{Error, Result}; const INDEX_DIR_NAME: &str = ".tgrep"; const INDEX_BUILD_BATCH_SIZE: usize = 1024; @@ -388,7 +388,13 @@ pub fn append_overlay_to_index( ) -> Result<()> { std::fs::create_dir_all(out_dir)?; - let base = reader.num_files() as u32; + // File IDs are `u32` on disk. Fail loudly rather than truncate. + let base = u32::try_from(reader.num_files()).map_err(|_| { + Error::IndexCorrupted(format!( + "reader has {} files, exceeding the u32 file-id limit", + reader.num_files() + )) + })?; // Overlay trigrams in ascending order for the 2-way merge. let mut overlay_trigrams: Vec = overlay_inverted.keys().copied().collect(); @@ -417,12 +423,15 @@ pub fn append_overlay_to_index( // A reader entry that exists in the lookup table but whose raw posting // bytes can't be read (truncated/corrupt mmap) yields `None` here while - // `ri` is still in range. Advance past it so we never stall consuming - // overlay keys against a dead reader slot, and never break the merge - // early while reader trigrams remain. + // `ri` is still in range. Silently skipping would drop that trigram yet + // still publish an index, turning reader corruption into silent data + // loss. Fail the flush instead so the caller keeps the previous reader + // plus the live overlay as a safe fallback. if ri < reader_trigram_count && reader_next.is_none() { - ri += 1; - continue; + return Err(Error::IndexCorrupted(format!( + "reader trigram entry {ri} of {reader_trigram_count} has unreadable \ + postings; refusing to publish an incomplete merged index" + ))); } let overlay_next = overlay_trigrams.get(oi).copied(); @@ -451,13 +460,22 @@ pub fn append_overlay_to_index( oi += 1; (ot, None, overlay_inverted.get(&ot)) } - // Both streams exhausted (reader entries already skipped above). + // Reader exhausted (in-range unreadable entries already errored above). (None, None) => break, }; let reader_len = reader_bytes.map_or(0, |b| b.len() / ondisk::POSTING_ENTRY_SIZE); let overlay_len = overlay_seq.map_or(0, |v| v.len()); - let length = (reader_len + overlay_len) as u32; + let length = u32::try_from( + reader_len + .checked_add(overlay_len) + .ok_or_else(|| Error::IndexCorrupted("posting list length overflow".into()))?, + ) + .map_err(|_| { + Error::IndexCorrupted(format!( + "posting list for trigram {trigram} exceeds the u32 length limit" + )) + })?; if length == 0 { continue; } @@ -481,8 +499,11 @@ pub fn append_overlay_to_index( for chunk in seq.chunks(POSTING_WRITE_CHUNK_ENTRIES) { posting_scratch.clear(); for &k in chunk { + let file_id = base.checked_add(k).ok_or_else(|| { + Error::IndexCorrupted("overlay file id overflow beyond u32".into()) + })?; let entry = PostingEntry { - file_id: base + k, + file_id, loc_mask: u8::MAX, next_mask: u8::MAX, }; From 61a8768b928d7f549a68962287c40f04e618e475 Mon Sep 17 00:00:00 2001 From: Shengyu Fu Date: Tue, 30 Jun 2026 17:43:22 -0700 Subject: [PATCH 08/12] Bump version to 0.1.21 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- Cargo.lock | 4 ++-- Cargo.toml | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 070e7a0..85368c0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1037,7 +1037,7 @@ checksum = "8f50febec83f5ee1df3015341d8bd429f2d1cc62bcba7ea2076759d315084683" [[package]] name = "tgrep-cli" -version = "0.1.20" +version = "0.1.21" dependencies = [ "anyhow", "assert_cmd", @@ -1061,7 +1061,7 @@ dependencies = [ [[package]] name = "tgrep-core" -version = "0.1.20" +version = "0.1.21" dependencies = [ "anyhow", "criterion", diff --git a/Cargo.toml b/Cargo.toml index 548ec32..3b8a97e 100644 --- a/Cargo.toml +++ b/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" From 1dc9e52ec9644da3b623b705b27d035531012265 Mon Sep 17 00:00:00 2001 From: Shengyu Fu Date: Tue, 30 Jun 2026 18:13:52 -0700 Subject: [PATCH 09/12] Close watcher race during final append-only publish handoff The final bulk-build handoff set flushing=true/indexing=false and only then let flush_append_only_overlay acquire snapshot_gate. handle_fs_event only skips while indexing is true, so a filesystem event could slip into that gap, take snapshot_gate.read() first, and mutate the LiveIndex before the final snapshot. If it touched a path already present in the on-disk reader, append_overlay_to_index would append a duplicate/stale path, breaking its brand-new-paths precondition and corrupting the persisted index. Acquire snapshot_gate.write() *before* clearing indexing and hold it across the whole final flush by splitting the flush into flush_append_only_overlay (acquires the gate) and flush_append_only_overlay_locked (assumes it is held). Any watcher event that observes indexing=false now blocks on the gate until the publish completes and then applies to the newly published reader, so no event is lost. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- tgrep-cli/src/serve.rs | 32 +++++++++++++++++++++++++++++--- 1 file changed, 29 insertions(+), 3 deletions(-) diff --git a/tgrep-cli/src/serve.rs b/tgrep-cli/src/serve.rs index b19c54e..dd44046 100644 --- a/tgrep-cli/src/serve.rs +++ b/tgrep-cli/src/serve.rs @@ -1582,6 +1582,17 @@ fn background_index_build(state: &Arc, root: &Path, index_dir: &Pat // repos. Set `flushing` *before* clearing `indexing` so the auto-save // loop never observes both flags as false during the handoff and // races us into a redundant parallel snapshot of the bulk overlay. + // + // Acquire the publish gate *before* clearing `indexing`. `handle_fs_event` + // only skips while `indexing` is true; once it's false a watcher event can + // run, and if it grabbed `snapshot_gate.read()` before our flush grabbed + // the write lock it could mutate the overlay in the gap between the flag + // flip and the final snapshot — updating/deleting a path already on disk in + // the reader and violating `append_overlay_to_index`'s brand-new-paths + // precondition. Holding the gate across the flip makes any such event block + // (not skip) until the flush publishes, after which it applies safely to + // the newly published reader; no event is lost. + let gate = state.snapshot_gate.write().unwrap(); state.flushing.store(true, Ordering::SeqCst); state.indexing.store(false, Ordering::SeqCst); @@ -1592,7 +1603,8 @@ fn background_index_build(state: &Arc, root: &Path, index_dir: &Pat // intermediate incremental flushes published `complete = false` so a // mid-build kill would resume rather than be treated as finished. eprintln!("[trace] persisting final index to disk..."); - let pruned = flush_append_only_overlay(state, index_dir, true, Some(&stamps)); + let pruned = flush_append_only_overlay_locked(state, index_dir, true, Some(&stamps)); + drop(gate); state.flushing.store(false, Ordering::SeqCst); @@ -1724,13 +1736,27 @@ fn flush_append_only_overlay( complete: bool, stamps: Option<&std::collections::HashMap>, ) -> bool { - let flush_start = Instant::now(); - // Hold the snapshot gate for the whole snapshot → publish → prune cycle, // mirroring flush_index_to_disk. (During the bulk build the watcher is // already suppressed, but auto-save coordination and future-proofing make // the gate the right call.) let _gate = state.snapshot_gate.write().unwrap(); + flush_append_only_overlay_locked(state, index_dir, complete, stamps) +} + +/// Body of [`flush_append_only_overlay`] that assumes `snapshot_gate` is +/// **already held for write** by the caller. Split out so the final bulk-build +/// handoff can acquire the gate *before* clearing the `indexing` flag, closing +/// the window where a filesystem event could observe `indexing == false`, take +/// the gate first, and mutate the overlay between the flag flip and the final +/// snapshot (which would break the append-only precondition). +fn flush_append_only_overlay_locked( + state: &ServerState, + index_dir: &Path, + complete: bool, + stamps: Option<&std::collections::HashMap>, +) -> bool { + let flush_start = Instant::now(); // Snapshot the overlay (bounded heap) and the current reader (cheap Arc). let (overlay_paths, overlay_inverted, reader) = { From 6ae59204cc2f4274d86f0ef84e19fd7a0f8f9ec4 Mon Sep 17 00:00:00 2001 From: Shengyu Fu Date: Tue, 30 Jun 2026 18:22:19 -0700 Subject: [PATCH 10/12] Validate posting offset in u64 before narrowing; shrink memory_cap fixture - reader: compute nth_trigram_raw range in u64 and bounds-check against the mmap length before converting to usize, so a large/corrupt offset can't truncate on 32-bit targets and slice the wrong region - tests: reduce memory_cap fixture from 1200 to 700 files (still exceeds the 500 batch size and forces multiple flushes) to keep CI fast Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- tgrep-cli/tests/memory_cap.rs | 7 ++++--- tgrep-core/src/reader.rs | 11 ++++++++--- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/tgrep-cli/tests/memory_cap.rs b/tgrep-cli/tests/memory_cap.rs index d5a9a9c..3f0d67d 100644 --- a/tgrep-cli/tests/memory_cap.rs +++ b/tgrep-cli/tests/memory_cap.rs @@ -16,9 +16,10 @@ use std::time::{Duration, Instant}; use tempfile::TempDir; -/// Number of files to index. Comfortably exceeds the indexer's internal batch -/// size (500) so that, with a 1 MB cap, several incremental flushes fire. -const NUM_FILES: usize = 1200; +/// Number of files to index. Exceeds the indexer's internal batch size (500) +/// so that, with a 1 MB cap, multiple incremental flushes fire, while staying +/// small enough to keep the test fast and CI-friendly. +const NUM_FILES: usize = 700; fn tgrep_bin() -> PathBuf { assert_cmd::cargo::cargo_bin("tgrep") diff --git a/tgrep-core/src/reader.rs b/tgrep-core/src/reader.rs index 4abddec..9629f83 100644 --- a/tgrep-core/src/reader.rs +++ b/tgrep-core/src/reader.rs @@ -326,12 +326,17 @@ impl IndexReader { } let entry = self.read_lookup_entry(i); let postings = self.postings.as_ref()?; - let start = entry.offset as usize; - let byte_len = (entry.length as usize).checked_mul(POSTING_ENTRY_SIZE)?; + // Do the range math in u64 and validate against the mmap length before + // narrowing to usize, so a large/corrupt offset can't truncate on + // 32-bit targets and yield an in-bounds slice from the wrong region. + let start = entry.offset; + let byte_len = (entry.length as u64).checked_mul(POSTING_ENTRY_SIZE as u64)?; let end = start.checked_add(byte_len)?; - if end > postings.len() { + if end > postings.len() as u64 { return None; } + let start = usize::try_from(start).ok()?; + let end = usize::try_from(end).ok()?; Some((entry.trigram, &postings[start..end])) } From 409c4be83e24fca881d63a511376e019764298df Mon Sep 17 00:00:00 2001 From: Shengyu Fu Date: Tue, 30 Jun 2026 18:27:39 -0700 Subject: [PATCH 11/12] Guard posting start offset against mmap length in nth_trigram_raw Explicitly reject start > len (not just end > len) so a corrupt zero-length entry with an out-of-range offset can never slice past the mmap. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- tgrep-core/src/reader.rs | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tgrep-core/src/reader.rs b/tgrep-core/src/reader.rs index 9629f83..cf9149f 100644 --- a/tgrep-core/src/reader.rs +++ b/tgrep-core/src/reader.rs @@ -332,7 +332,12 @@ impl IndexReader { let start = entry.offset; let byte_len = (entry.length as u64).checked_mul(POSTING_ENTRY_SIZE as u64)?; let end = start.checked_add(byte_len)?; - if end > postings.len() as u64 { + // Guard both ends against the mmap length. `end >= start` holds because + // `byte_len` is unsigned and `checked_add` rejects overflow, but check + // `start` explicitly so a zero-length entry with an out-of-range offset + // can never reach the slice below. + let len = postings.len() as u64; + if start > len || end > len { return None; } let start = usize::try_from(start).ok()?; From 38e23bf08abadd45487473cbd5c7c6d9fdd0dd55 Mon Sep 17 00:00:00 2001 From: Shengyu Fu Date: Tue, 30 Jun 2026 18:33:31 -0700 Subject: [PATCH 12/12] Initialize PROCESS_MEMORY_COUNTERS.cb before GetProcessMemoryInfo The Windows RSS query zeroed the struct but left cb at 0. GetProcessMemoryInfo requires cb to hold the struct size on input; without it the call can fail, silently disabling the memory cap on Windows. Set cb before the call and add mem-query regression tests asserting non-zero RSS/physical memory. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- tgrep-cli/src/mem.rs | 36 ++++++++++++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/tgrep-cli/src/mem.rs b/tgrep-cli/src/mem.rs index 00db062..0411c17 100644 --- a/tgrep-cli/src/mem.rs +++ b/tgrep-cli/src/mem.rs @@ -17,6 +17,8 @@ pub fn process_rss_bytes() -> Option { let handle = GetCurrentProcess(); let mut counters = MaybeUninit::::zeroed(); let size = std::mem::size_of::() 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(); @@ -132,3 +134,37 @@ pub fn default_memory_cap_bytes() -> u64 { 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); + } +}