Skip to main content

lean_ctx/core/
memory_guard.rs

1//! Process-level RAM guardian with adaptive eviction and hard OOM protection.
2//!
3//! Monitors RSS via platform-specific APIs and triggers tiered cache eviction
4//! when memory usage exceeds configurable thresholds (default: 5% of system RAM,
5//! lowered to `mcp_max_rss_mb` — 512 MB — inside a stdio MCP server).
6//! At critical levels, performs aggressive eviction and signals background tasks
7//! to abort. It never exits the process — recovery is always via eviction.
8
9use std::sync::Arc;
10use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU64, Ordering};
11
12static PEAK_RSS: AtomicU64 = AtomicU64::new(0);
13static GUARD_RUNNING: AtomicBool = AtomicBool::new(false);
14static ABORT_REQUESTED: AtomicBool = AtomicBool::new(false);
15static CURRENT_PRESSURE: AtomicU8 = AtomicU8::new(0);
16/// Absolute per-process RSS target in bytes (0 = none). The stdio MCP server
17/// sets it at startup: hosts such as the Codex app-server keep one server per
18/// loaded thread alive for days, and on a large machine the percentage target
19/// alone let each of them settle near 1 GB before any eviction ran.
20static PROCESS_RSS_CAP: AtomicU64 = AtomicU64::new(0);
21
22/// Install an absolute RSS target for this process; the effective limit is the
23/// lower of it and the `max_ram_percent` target. `0` removes it.
24pub fn set_process_rss_cap(bytes: u64) {
25    PROCESS_RSS_CAP.store(bytes, Ordering::Relaxed);
26}
27
28/// Current process memory in bytes, or `None` if unavailable: resident set
29/// size on Linux and Windows, physical footprint on macOS (see
30/// `macos_footprint` — resident size there hides compressed pages).
31pub fn get_rss_bytes() -> Option<u64> {
32    #[cfg(target_os = "linux")]
33    {
34        linux_rss()
35    }
36    #[cfg(target_os = "macos")]
37    {
38        macos_rss()
39    }
40    #[cfg(windows)]
41    {
42        windows_rss()
43    }
44    #[cfg(not(any(target_os = "linux", target_os = "macos", windows)))]
45    {
46        None
47    }
48}
49
50/// RSS of an arbitrary process by PID, or `None` if unavailable/dead.
51pub fn get_rss_bytes_for_pid(pid: u32) -> Option<u64> {
52    #[cfg(target_os = "linux")]
53    {
54        linux_rss_for_pid(pid)
55    }
56    #[cfg(target_os = "macos")]
57    {
58        macos_rss_for_pid(pid)
59    }
60    #[cfg(windows)]
61    {
62        windows_rss_for_pid(pid)
63    }
64    #[cfg(not(any(target_os = "linux", target_os = "macos", windows)))]
65    {
66        let _ = pid;
67        None
68    }
69}
70
71/// Total physical RAM in bytes, or `None` if unavailable.
72pub fn get_system_ram_bytes() -> Option<u64> {
73    #[cfg(target_os = "linux")]
74    {
75        linux_memtotal()
76    }
77    #[cfg(target_os = "macos")]
78    {
79        macos_memsize()
80    }
81    #[cfg(windows)]
82    {
83        windows_system_ram()
84    }
85    #[cfg(not(any(target_os = "linux", target_os = "macos", windows)))]
86    {
87        None
88    }
89}
90
91#[cfg(windows)]
92fn windows_rss() -> Option<u64> {
93    use windows_sys::Win32::System::Threading::GetCurrentProcess;
94
95    // SAFETY: GetCurrentProcess returns a pseudo-handle valid for the current
96    // process and it must not be closed.
97    let process = unsafe { GetCurrentProcess() };
98    windows_rss_for_handle(process)
99}
100
101#[cfg(windows)]
102fn windows_rss_for_pid(pid: u32) -> Option<u64> {
103    use windows_sys::Win32::Foundation::CloseHandle;
104    use windows_sys::Win32::System::Threading::{OpenProcess, PROCESS_QUERY_LIMITED_INFORMATION};
105
106    // SAFETY: OpenProcess is called with query-only access. A successful real
107    // handle is closed before returning.
108    let process = unsafe { OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, 0, pid) };
109
110    if process.is_null() {
111        return None;
112    }
113
114    let rss = windows_rss_for_handle(process);
115
116    // SAFETY: `process` is a real handle returned successfully by OpenProcess.
117    unsafe {
118        let _ = CloseHandle(process);
119    }
120
121    rss
122}
123
124#[cfg(windows)]
125fn windows_rss_for_handle(process: windows_sys::Win32::Foundation::HANDLE) -> Option<u64> {
126    use windows_sys::Win32::System::ProcessStatus::{
127        K32GetProcessMemoryInfo, PROCESS_MEMORY_COUNTERS,
128    };
129
130    let mut counters = PROCESS_MEMORY_COUNTERS {
131        cb: std::mem::size_of::<PROCESS_MEMORY_COUNTERS>() as u32,
132        ..PROCESS_MEMORY_COUNTERS::default()
133    };
134
135    // SAFETY: callers pass the current-process pseudo-handle or a live handle
136    // returned by OpenProcess; `counters` is correctly sized and writable.
137    let ok = unsafe { K32GetProcessMemoryInfo(process, &raw mut counters, counters.cb) };
138
139    if ok == 0 {
140        None
141    } else {
142        Some(counters.WorkingSetSize as u64)
143    }
144}
145
146#[cfg(windows)]
147fn windows_system_ram() -> Option<u64> {
148    use windows_sys::Win32::System::SystemInformation::{GlobalMemoryStatusEx, MEMORYSTATUSEX};
149
150    let mut status = MEMORYSTATUSEX {
151        dwLength: std::mem::size_of::<MEMORYSTATUSEX>() as u32,
152        ..MEMORYSTATUSEX::default()
153    };
154
155    // SAFETY: `status` has the required `dwLength` and is writable for the call.
156    let ok = unsafe { GlobalMemoryStatusEx(&raw mut status) };
157
158    if ok == 0 {
159        None
160    } else {
161        Some(status.ullTotalPhys)
162    }
163}
164
165/// Returns the RSS limit in bytes: the `max_ram_percent` target, lowered to
166/// the process cap when one is installed ([`set_process_rss_cap`]).
167pub fn rss_limit_bytes() -> Option<u64> {
168    let sys_ram = get_system_ram_bytes()?;
169    let cfg = super::config::Config::load();
170    let pct = super::config::MemoryGuardConfig::effective(&cfg).max_ram_percent;
171    Some(effective_limit(
172        sys_ram / 100 * u64::from(pct),
173        PROCESS_RSS_CAP.load(Ordering::Relaxed),
174    ))
175}
176
177fn effective_limit(percent_limit: u64, cap: u64) -> u64 {
178    if cap == 0 {
179        percent_limit
180    } else {
181        percent_limit.min(cap)
182    }
183}
184
185/// Pressure tier for `rss` against the base `limit` (#790 multipliers:
186/// Soft above 1×, Medium above 1.2×, Hard above 1.5×, Critical above 2×).
187fn pressure_for(rss: u64, limit: u64) -> PressureLevel {
188    if limit == 0 {
189        return PressureLevel::Normal;
190    }
191    let ratio = rss as f64 / limit as f64;
192    if ratio > 2.0 {
193        PressureLevel::Critical
194    } else if ratio > 1.5 {
195        PressureLevel::Hard
196    } else if ratio > 1.2 {
197        PressureLevel::Medium
198    } else if ratio > 1.0 {
199        PressureLevel::Soft
200    } else {
201        PressureLevel::Normal
202    }
203}
204
205/// Computes a memory-bounded work batch from current RSS and the guardian's
206/// hard threshold (1.5× the configured base limit). Falls back to `max_files`
207/// when RSS is unavailable. The lower bound keeps progress deterministic while
208/// allowing callers to reduce in-flight work substantially under low headroom.
209pub fn adaptive_batch_size(
210    min_files: usize,
211    max_files: usize,
212    estimated_bytes_per_file: u64,
213) -> usize {
214    let headroom = match (rss_limit_bytes(), get_rss_bytes()) {
215        (Some(limit), Some(rss)) => hard_headroom_bytes(limit, rss),
216        _ => return max_files.max(1),
217    };
218    batch_size_for_headroom(headroom, min_files, max_files, estimated_bytes_per_file)
219}
220
221fn hard_headroom_bytes(base_limit: u64, rss: u64) -> u64 {
222    base_limit
223        .saturating_mul(3)
224        .saturating_div(2)
225        .saturating_sub(rss)
226}
227
228fn batch_size_for_headroom(
229    headroom_bytes: u64,
230    min_files: usize,
231    max_files: usize,
232    estimated_bytes_per_file: u64,
233) -> usize {
234    let min_files = min_files.max(1);
235    let max_files = max_files.max(min_files);
236    let estimate = estimated_bytes_per_file.max(1);
237    let by_headroom = (headroom_bytes / estimate).min(max_files as u64) as usize;
238    by_headroom.clamp(min_files, max_files)
239}
240
241#[cfg(test)]
242pub mod adaptive_batch_tests {
243    use super::{batch_size_for_headroom, hard_headroom_bytes};
244
245    #[test]
246    fn hard_headroom_uses_guardian_hard_threshold() {
247        assert_eq!(hard_headroom_bytes(1_000, 1_100), 400);
248        assert_eq!(hard_headroom_bytes(1_000, 1_500), 0);
249        assert_eq!(hard_headroom_bytes(1_000, 2_000), 0);
250    }
251
252    #[test]
253    fn batch_size_tracks_headroom_and_bounds() {
254        assert_eq!(batch_size_for_headroom(1_000, 1, 500, 100), 10);
255        assert_eq!(batch_size_for_headroom(0, 1, 500, 100), 1);
256        assert_eq!(batch_size_for_headroom(100_000, 1, 500, 100), 500);
257    }
258
259    #[test]
260    fn batch_size_sanitizes_zero_bounds_and_estimate() {
261        assert_eq!(batch_size_for_headroom(0, 0, 0, 0), 1);
262    }
263}
264
265/// Recorded peak RSS since process start.
266pub fn peak_rss_bytes() -> u64 {
267    PEAK_RSS.load(Ordering::Relaxed)
268}
269
270/// Snapshot of current memory state for diagnostics.
271#[derive(Debug, Clone, serde::Serialize)]
272pub struct MemorySnapshot {
273    pub rss_bytes: u64,
274    pub peak_rss_bytes: u64,
275    pub system_ram_bytes: u64,
276    pub rss_limit_bytes: u64,
277    pub rss_percent: f64,
278    pub pressure_level: PressureLevel,
279}
280
281#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, serde::Serialize)]
282#[serde(rename_all = "lowercase")]
283#[repr(u8)]
284pub enum PressureLevel {
285    Normal = 0,
286    Soft = 1,
287    Medium = 2,
288    Hard = 3,
289    Critical = 4,
290}
291
292impl PressureLevel {
293    fn from_u8(v: u8) -> Self {
294        match v {
295            1 => Self::Soft,
296            2 => Self::Medium,
297            3 => Self::Hard,
298            4 => Self::Critical,
299            _ => Self::Normal,
300        }
301    }
302}
303
304impl MemorySnapshot {
305    /// Capture memory snapshot of the **current** process.
306    pub fn capture() -> Option<Self> {
307        Self::capture_impl(get_rss_bytes()?)
308    }
309
310    /// Capture memory snapshot for the **daemon** process (by PID).
311    /// Falls back to the current process if the PID is dead or unreadable.
312    pub fn capture_for_pid(pid: u32) -> Option<Self> {
313        let rss = get_rss_bytes_for_pid(pid).or_else(get_rss_bytes)?;
314        Self::capture_impl(rss)
315    }
316
317    fn capture_impl(rss: u64) -> Option<Self> {
318        let sys = get_system_ram_bytes()?;
319        let limit = rss_limit_bytes()?;
320        let pct = if sys > 0 {
321            (rss as f64 / sys as f64) * 100.0
322        } else {
323            0.0
324        };
325
326        PEAK_RSS.fetch_max(rss, Ordering::Relaxed);
327
328        // #790: tightened multipliers so ABORT fires earlier — e.g. 10% config
329        // on 64 GB → Hard at 9.6 GB, Critical at 12.8 GB. Users expect the
330        // limit to be a meaningful cap, not a 3× suggestion. Measured against
331        // bytes so the per-process cap and the percentage share one ladder.
332        let level = pressure_for(rss, limit);
333
334        Some(Self {
335            rss_bytes: rss,
336            peak_rss_bytes: PEAK_RSS.load(Ordering::Relaxed),
337            system_ram_bytes: sys,
338            rss_limit_bytes: limit,
339            rss_percent: pct,
340            pressure_level: level,
341        })
342    }
343}
344
345/// Live heap bytes and allocator-resident bytes (`stats.allocated`,
346/// `stats.resident`), when jemalloc is the allocator. The gap between process
347/// RSS and `allocated` is what eviction cannot reach: allocator retention,
348/// code and thread stacks.
349pub fn allocator_stats() -> Option<(u64, u64)> {
350    #[cfg(all(feature = "jemalloc", not(windows), not(target_env = "musl")))]
351    {
352        tikv_jemalloc_ctl::epoch::advance().ok()?;
353        let allocated = tikv_jemalloc_ctl::stats::allocated::read().ok()?;
354        let resident = tikv_jemalloc_ctl::stats::resident::read().ok()?;
355        Some((allocated as u64, resident as u64))
356    }
357    #[cfg(not(all(feature = "jemalloc", not(windows), not(target_env = "musl"))))]
358    {
359        None
360    }
361}
362
363/// Force-purge all jemalloc arenas to return memory to the OS.
364/// Uses `MALLCTL_ARENAS_ALL` (value 4096) which is the jemalloc sentinel
365/// for "all arenas". Logs errors instead of silently swallowing them.
366pub fn jemalloc_purge() {
367    #[cfg(all(feature = "jemalloc", not(windows), not(target_env = "musl")))]
368    if let Err(rc) = purge_all_arenas() {
369        tracing::warn!("[memory_guard] jemalloc purge failed: mallctl rc={rc}");
370    }
371}
372
373/// `arena.<i>.purge` is a void ctl: it must be called with no old and no new
374/// value. Writing one (as this did before #2075) fails with EPERM, which left
375/// every purge a silent no-op: measured on macOS, a dropped 900 MB parse
376/// working set stayed resident instead of falling back to ~50 MB.
377#[cfg(all(feature = "jemalloc", not(windows), not(target_env = "musl")))]
378fn purge_all_arenas() -> Result<(), i32> {
379    // SAFETY: static NUL-terminated ctl name; null old/new pointers with zero
380    // length are exactly the calling convention of a void mallctl.
381    let rc = unsafe {
382        rjem::mallctl(
383            c"arena.4096.purge".as_ptr(),
384            std::ptr::null_mut(),
385            std::ptr::null_mut(),
386            std::ptr::null_mut(),
387            0,
388        )
389    };
390    if rc == 0 { Ok(()) } else { Err(rc) }
391}
392
393/// jemalloc's C entry points. tikv-jemalloc-sys builds jemalloc with the
394/// `_rjem_` symbol prefix (see `_rjem_malloc_conf` in `main.rs`) and
395/// tikv-jemallocator already links it, so declaring the few symbols needed
396/// here avoids a direct dependency on the -sys crate.
397#[cfg(all(feature = "jemalloc", not(windows), not(target_env = "musl")))]
398mod rjem {
399    use std::ffi::{c_char, c_int, c_void};
400
401    unsafe extern "C" {
402        #[cfg(feature = "tree-sitter")]
403        #[link_name = "_rjem_malloc"]
404        pub(super) fn malloc(size: usize) -> *mut c_void;
405        #[cfg(feature = "tree-sitter")]
406        #[link_name = "_rjem_calloc"]
407        pub(super) fn calloc(count: usize, size: usize) -> *mut c_void;
408        #[cfg(feature = "tree-sitter")]
409        #[link_name = "_rjem_realloc"]
410        pub(super) fn realloc(ptr: *mut c_void, size: usize) -> *mut c_void;
411        #[cfg(feature = "tree-sitter")]
412        #[link_name = "_rjem_free"]
413        pub(super) fn free(ptr: *mut c_void);
414        #[link_name = "_rjem_mallctl"]
415        pub(super) fn mallctl(
416            name: *const c_char,
417            oldp: *mut c_void,
418            oldlenp: *mut usize,
419            newp: *mut c_void,
420            newlen: usize,
421        ) -> c_int;
422    }
423}
424
425/// Returns `true` if the guardian has requested background tasks to abort.
426pub fn abort_requested() -> bool {
427    ABORT_REQUESTED.load(Ordering::Relaxed)
428}
429
430/// Quick, non-allocating memory pressure check for hot loops (scanners, indexers).
431/// Reads the cached atomic flag set by the guardian thread — O(1), no syscalls.
432pub fn is_under_pressure() -> bool {
433    current_pressure() >= PressureLevel::Soft
434}
435
436/// Returns the current pressure level as last observed by the guardian thread.
437pub fn current_pressure() -> PressureLevel {
438    PressureLevel::from_u8(CURRENT_PRESSURE.load(Ordering::Relaxed))
439}
440
441#[inline]
442const fn pressure_requests_abort(level: PressureLevel) -> bool {
443    level as u8 >= PressureLevel::Hard as u8
444}
445
446#[inline]
447fn publish_pressure(level: PressureLevel) {
448    CURRENT_PRESSURE.store(level as u8, Ordering::Relaxed);
449    ABORT_REQUESTED.store(pressure_requests_abort(level), Ordering::SeqCst);
450}
451
452#[derive(Debug, Clone, Copy, PartialEq, Eq)]
453enum EvictionSchedule {
454    RetryAfter(u64),
455    PauseFor(u64),
456}
457
458impl EvictionSchedule {
459    fn poll_secs(self) -> u64 {
460        match self {
461            Self::RetryAfter(secs) | Self::PauseFor(secs) => secs,
462        }
463    }
464}
465
466/// Backoff state for eviction rounds (any pressure level) that reclaimed no memory.
467#[derive(Default)]
468struct EvictionBackoff {
469    consecutive_zero_progress: u32,
470}
471
472impl EvictionBackoff {
473    const BASE_POLL_SECS: u64 = 1;
474    const MAX_POLL_SECS: u64 = 60;
475    const PAUSE_AFTER_ZERO_PROGRESS_ROUNDS: u32 = 10;
476    const PAUSE_SECS: u64 = 5 * 60;
477
478    fn record(&mut self, made_progress: bool) -> EvictionSchedule {
479        if made_progress {
480            self.consecutive_zero_progress = 0;
481            return EvictionSchedule::RetryAfter(Self::BASE_POLL_SECS);
482        }
483
484        self.consecutive_zero_progress = self.consecutive_zero_progress.saturating_add(1);
485        if self.consecutive_zero_progress >= Self::PAUSE_AFTER_ZERO_PROGRESS_ROUNDS {
486            return EvictionSchedule::PauseFor(Self::PAUSE_SECS);
487        }
488
489        let exponent = self.consecutive_zero_progress.saturating_sub(2).min(6);
490        let poll_secs = Self::BASE_POLL_SECS
491            .saturating_mul(1u64 << exponent)
492            .min(Self::MAX_POLL_SECS);
493        EvictionSchedule::RetryAfter(poll_secs)
494    }
495
496    fn reset(&mut self) {
497        self.consecutive_zero_progress = 0;
498    }
499}
500
501/// When the next eviction round may run. Sampling continues every second
502/// under pressure; the gate only spaces out rounds that reclaimed nothing,
503/// and lets a rise to a higher pressure level through at once.
504#[derive(Default)]
505struct EvictionGate {
506    backoff: EvictionBackoff,
507    last: Option<(PressureLevel, std::time::Instant)>,
508}
509
510impl EvictionGate {
511    fn due(&self, level: PressureLevel, now: std::time::Instant) -> bool {
512        self.last
513            .is_none_or(|(last_level, next_at)| level > last_level || now >= next_at)
514    }
515
516    fn record(
517        &mut self,
518        level: PressureLevel,
519        made_progress: bool,
520        now: std::time::Instant,
521    ) -> EvictionSchedule {
522        if self.last.is_some_and(|(last_level, _)| level > last_level) {
523            self.backoff.reset();
524        }
525        let schedule = self.backoff.record(made_progress);
526        self.last = Some((
527            level,
528            now + std::time::Duration::from_secs(schedule.poll_secs()),
529        ));
530        schedule
531    }
532}
533
534/// Start the background memory guardian task (idempotent).
535/// Polls every 3s (normal), 1s (under pressure), or up to 15s once RSS has been
536/// stably calm (idle backoff). Eviction rounds that reclaim nothing back off to
537/// 60s, then pause for five minutes, at every pressure level; pressure logs
538/// repeat at most once a minute per level. The callback returns whether it reclaimed
539/// memory so the guard can distinguish a real retry from an empty-cache loop.
540pub fn start_guard(eviction_callback: Arc<dyn Fn(PressureLevel) -> bool + Send + Sync>) {
541    // The guardian is a long-lived background monitor for the running
542    // server/daemon. Under `cargo test` a single OS process executes the entire
543    // suite, so its RSS routinely exceeds the per-operation pressure threshold
544    // (default 5% of system RAM). A test that constructs a server (e.g. the
545    // `http_server` tests via `new_shared_with_context`) would start this thread,
546    // which then flips the process-global `CURRENT_PRESSURE` / `ABORT_REQUESTED`
547    // flags. Unrelated later tests in the same binary read those flags and skip
548    // work — notably `graph_index::build_edges_with_cache` aborts edge-building
549    // under pressure, leaving indexed files with no edges. That manifested as an
550    // intermittent, macOS-only flake ("No files depend on Base.gd"). The guardian
551    // has no purpose inside the test harness, so never start it there. Production
552    // and the daemon compile without `cfg!(test)` and are unaffected.
553    if cfg!(test) {
554        return;
555    }
556    if GUARD_RUNNING.swap(true, Ordering::SeqCst) {
557        return;
558    }
559    std::thread::Builder::new()
560        .name("memory-guard".into())
561        .spawn(move || {
562            // Idle backoff: once RSS has stayed below the Soft threshold for
563            // CALM_TICKS_BEFORE_BACKOFF consecutive samples, stretch the poll
564            // interval to IDLE_POLL_SECS. An idle server allocates nothing, so 3s
565            // RSS sampling is just wasted wakeups; any pressure resets the cadence
566            // instantly (below), leaving OOM reaction time during real work
567            // unchanged (#453 idle hygiene).
568            const CALM_TICKS_BEFORE_BACKOFF: u64 = 5;
569            const IDLE_POLL_SECS: u64 = 15;
570            let mut poll_secs = 3u64;
571            let mut calm_ticks = 0u64;
572            let mut gate = EvictionGate::default();
573
574            // Pressure logging: on every level change, then at most once per
575            // LOG_EVERY while the level holds — a process whose baseline sits
576            // above a small cap must not write a line per second for hours.
577            const LOG_EVERY: std::time::Duration = std::time::Duration::from_mins(1);
578            let mut logged: Option<(PressureLevel, std::time::Instant)> = None;
579
580            // #790: the first sample runs immediately — close the 3s blind
581            // window so builders that start right after start_guard() see
582            // real pressure.
583            let mut first = true;
584            loop {
585                if !first {
586                    std::thread::sleep(std::time::Duration::from_secs(poll_secs));
587                }
588                first = false;
589                let Some(snap) = MemorySnapshot::capture() else {
590                    continue;
591                };
592
593                publish_pressure(snap.pressure_level);
594
595                if snap.pressure_level < PressureLevel::Soft {
596                    gate = EvictionGate::default();
597                    logged = None;
598                    calm_ticks = calm_ticks.saturating_add(1);
599                    poll_secs = if calm_ticks >= CALM_TICKS_BEFORE_BACKOFF {
600                        IDLE_POLL_SECS
601                    } else {
602                        3
603                    };
604                    continue;
605                }
606                calm_ticks = 0;
607                // Under pressure the guard samples every second, so
608                // `publish_pressure` (and with it build aborts) never lags;
609                // only eviction attempts follow the gate's backoff.
610                poll_secs = 1;
611
612                if logged.is_none_or(|(level, at)| {
613                    level != snap.pressure_level || at.elapsed() >= LOG_EVERY
614                }) {
615                    logged = Some((snap.pressure_level, std::time::Instant::now()));
616                    let heap = allocator_stats()
617                        .map(|(allocated, _)| format!(" heap={}MB", allocated >> 20))
618                        .unwrap_or_default();
619                    let line = format!(
620                        "[memory_guard] pressure={:?} RSS={:.0}MB{heap} limit={:.0}MB ({:.1}% of {:.0}GB)",
621                        snap.pressure_level,
622                        snap.rss_bytes as f64 / 1_048_576.0,
623                        snap.rss_limit_bytes as f64 / 1_048_576.0,
624                        snap.rss_percent,
625                        snap.system_ram_bytes as f64 / 1_073_741_824.0,
626                    );
627                    if snap.pressure_level == PressureLevel::Critical {
628                        tracing::error!("{line} — aggressive eviction to prevent OS OOM kill");
629                    } else {
630                        tracing::warn!("{line}");
631                    }
632                }
633
634                let now = std::time::Instant::now();
635                if !gate.due(snap.pressure_level, now) {
636                    continue;
637                }
638                let made_progress = (eviction_callback)(snap.pressure_level);
639                if snap.pressure_level >= PressureLevel::Hard {
640                    jemalloc_purge();
641                    release_system_heap();
642                }
643
644                // Rounds that reclaim nothing back off (and finally pause) at
645                // every level: what is left above the limit is not evictable.
646                // An escalation to a higher level is due at once regardless.
647                if let EvictionSchedule::PauseFor(secs) =
648                    gate.record(snap.pressure_level, made_progress, now)
649                {
650                    tracing::warn!(
651                        "[memory_guard] eviction made no progress for {} rounds at {:?}; \
652                         pausing eviction for {secs}s unless pressure rises",
653                        gate.backoff.consecutive_zero_progress,
654                        snap.pressure_level,
655                    );
656                }
657            }
658        })
659        .ok();
660}
661
662/// Force immediate purge of all caches and jemalloc arenas.
663pub fn force_purge() {
664    jemalloc_purge();
665    release_system_heap();
666    tracing::info!("[memory_guard] force_purge completed");
667}
668
669/// Return freed pages to the OS now, then measure and publish the pressure
670/// again. For work that would otherwise refuse under pressure the guardian
671/// published earlier: memory a finished request dropped is not pressure.
672pub fn reclaim_now() -> PressureLevel {
673    jemalloc_purge();
674    release_system_heap();
675    match MemorySnapshot::capture() {
676        Some(snap) => {
677            publish_pressure(snap.pressure_level);
678            snap.pressure_level
679        }
680        None => current_pressure(),
681    }
682}
683
684/// Route tree-sitter's C allocations (parse trees, parser stacks, query
685/// cursors) through jemalloc, whose arenas [`jemalloc_purge`] returns to the
686/// OS. The platform malloc keeps them instead: measured on macOS, freed
687/// parser-sized blocks stayed in the footprint and `malloc_zone_pressure_relief`
688/// released nothing, so a daemon that parses every file of every root grew
689/// with each new root (#2075).
690///
691/// Must run once at process start, before the first tree-sitter object
692/// exists: a block allocated by one allocator must never be freed by the other.
693pub fn route_tree_sitter_to_jemalloc() {
694    #[cfg(all(
695        feature = "jemalloc",
696        feature = "tree-sitter",
697        not(windows),
698        not(target_env = "musl")
699    ))]
700    // SAFETY: called first thing in `main`, before any thread or tree-sitter
701    // object exists, so no allocation can cross allocators. The jemalloc entry
702    // points have exactly the C `malloc` family's signatures and semantics.
703    unsafe {
704        tree_sitter::set_allocator(
705            Some(rjem::malloc),
706            Some(rjem::calloc),
707            Some(rjem::realloc),
708            Some(rjem::free),
709        );
710    }
711}
712
713/// Return glibc's free pages to the OS. C dependencies other than tree-sitter
714/// (SQLite, regex engines) still allocate through the platform malloc, whose
715/// per-thread arenas otherwise keep every page an index build touched.
716/// macOS has no equivalent: `malloc_zone_pressure_relief` released nothing
717/// in measurement, which is why tree-sitter is routed to jemalloc instead.
718pub fn release_system_heap() {
719    #[cfg(all(target_os = "linux", target_env = "gnu"))]
720    // SAFETY: `malloc_trim` only walks glibc's own arenas and releases pages
721    // that hold no live allocation; it is thread-safe.
722    unsafe {
723        libc::malloc_trim(0);
724    }
725}
726
727// --- Platform-specific implementations ---
728
729#[cfg(target_os = "linux")]
730fn linux_rss() -> Option<u64> {
731    linux_rss_for_pid(std::process::id())
732}
733
734#[cfg(target_os = "linux")]
735fn linux_rss_for_pid(pid: u32) -> Option<u64> {
736    let path = format!("/proc/{pid}/status");
737    let status = std::fs::read_to_string(path).ok()?;
738    for line in status.lines() {
739        if let Some(val) = line.strip_prefix("VmRSS:") {
740            let kb: u64 = val.trim().trim_end_matches(" kB").trim().parse().ok()?;
741            return Some(kb * 1024);
742        }
743    }
744    None
745}
746
747#[cfg(target_os = "linux")]
748fn linux_memtotal() -> Option<u64> {
749    let info = std::fs::read_to_string("/proc/meminfo").ok()?;
750    for line in info.lines() {
751        if let Some(val) = line.strip_prefix("MemTotal:") {
752            let kb: u64 = val.trim().trim_end_matches(" kB").trim().parse().ok()?;
753            return Some(kb * 1024);
754        }
755    }
756    None
757}
758
759#[cfg(target_os = "macos")]
760fn macos_rss() -> Option<u64> {
761    macos_footprint(std::process::id()).or_else(macos_resident_size)
762}
763
764/// Physical footprint (`ri_phys_footprint`): the figure Activity Monitor and
765/// the kernel's memory-pressure (jetsam) accounting use. Unlike
766/// `resident_size` it keeps counting pages the memory compressor squeezed out
767/// of RAM, so a long-idle process cannot hide a large heap from the guardian —
768/// measured: an idle MCP server at 40 MB resident still had a ~180 MB heap.
769/// `proc_pid_rusage` is public libproc API and works for any same-user pid.
770#[cfg(target_os = "macos")]
771fn macos_footprint(pid: u32) -> Option<u64> {
772    // SAFETY: `rusage_info_v2` is a plain C struct for which all-zero bytes
773    // are a valid value.
774    let mut info: libc::rusage_info_v2 = unsafe { std::mem::zeroed() };
775    let pid = libc::c_int::try_from(pid).ok()?;
776    // SAFETY: `info` is a live, correctly sized `rusage_info_v2` for the
777    // requested `RUSAGE_INFO_V2` flavour; libproc writes at most that struct.
778    let rc = unsafe {
779        libc::proc_pid_rusage(
780            pid,
781            libc::RUSAGE_INFO_V2,
782            std::ptr::from_mut(&mut info).cast::<libc::rusage_info_t>(),
783        )
784    };
785    (rc == 0).then_some(info.ri_phys_footprint)
786}
787
788#[cfg(target_os = "macos")]
789#[allow(deprecated, clippy::borrow_as_ptr, clippy::ptr_as_ptr)]
790fn macos_resident_size() -> Option<u64> {
791    use std::mem;
792    // SAFETY: `mach_task_basic_info_data_t` is a plain C struct for which an
793    // all-zero bit pattern is a valid initial value.
794    let mut info: libc::mach_task_basic_info_data_t = unsafe { mem::zeroed() };
795    let mut count = (mem::size_of::<libc::mach_task_basic_info_data_t>()
796        / mem::size_of::<libc::natural_t>()) as libc::mach_msg_type_number_t;
797    // SAFETY: `mach_task_self()` returns the current task port; `info` and
798    // `count` are live stack locals passed as out-pointers, sized to match the
799    // requested `MACH_TASK_BASIC_INFO` flavour.
800    let kr = unsafe {
801        libc::task_info(
802            libc::mach_task_self(),
803            libc::MACH_TASK_BASIC_INFO,
804            std::ptr::from_mut(&mut info).cast::<i32>(),
805            std::ptr::from_mut(&mut count),
806        )
807    };
808    if kr == libc::KERN_SUCCESS {
809        Some(info.resident_size)
810    } else {
811        None
812    }
813}
814
815#[cfg(target_os = "macos")]
816fn macos_rss_for_pid(pid: u32) -> Option<u64> {
817    if let Some(footprint) = macos_footprint(pid) {
818        return Some(footprint);
819    }
820    // Fallback: `ps -o rss= -p <pid>` (`task_for_pid` needs entitlements).
821    let output = std::process::Command::new("ps")
822        .args(["-o", "rss=", "-p", &pid.to_string()])
823        .output()
824        .ok()?;
825    if !output.status.success() {
826        return None;
827    }
828    let text = String::from_utf8_lossy(&output.stdout);
829    let kb: u64 = text.trim().parse().ok()?;
830    Some(kb * 1024)
831}
832
833#[cfg(target_os = "macos")]
834#[allow(clippy::borrow_as_ptr, clippy::ptr_as_ptr)]
835fn macos_memsize() -> Option<u64> {
836    use std::mem;
837    let mut memsize: u64 = 0;
838    let mut len = mem::size_of::<u64>();
839    let name = b"hw.memsize\0";
840    // SAFETY: `name` is a static, NUL-terminated sysctl name; `memsize` and
841    // `len` are live stack out-pointers whose sizes match the queried value.
842    let ret = unsafe {
843        libc::sysctlbyname(
844            name.as_ptr().cast(),
845            std::ptr::from_mut(&mut memsize).cast::<libc::c_void>(),
846            std::ptr::from_mut(&mut len),
847            std::ptr::null_mut(),
848            0,
849        )
850    };
851    if ret == 0 { Some(memsize) } else { None }
852}
853
854#[cfg(test)]
855pub mod tests {
856    use super::*;
857
858    /// #1899: jemalloc reads the conf as a C string and rejects
859    /// `background_thread` wherever it lacks pthread support.
860    #[cfg(all(feature = "jemalloc", not(windows), not(target_env = "musl")))]
861    #[test]
862    fn jemalloc_conf_is_c_string_with_platform_safe_options() {
863        let conf = crate::JEMALLOC_CONF;
864        assert_eq!(conf.last(), Some(&0), "must be NUL-terminated");
865        let body = std::str::from_utf8(&conf[..conf.len() - 1]).unwrap();
866        assert!(!body.contains('\0'), "no interior NUL");
867        assert!(body.contains("dirty_decay_ms:1000"));
868        assert_eq!(
869            body.contains("background_thread"),
870            cfg!(target_os = "linux"),
871            "background_thread is only supported on Linux: {body}"
872        );
873    }
874
875    /// #2075: the purge ctl is void; passing a value made it fail with EPERM,
876    /// so no purge ever returned a page to the OS.
877    #[cfg(all(feature = "jemalloc", not(windows), not(target_env = "musl")))]
878    #[test]
879    fn jemalloc_purge_is_accepted_by_the_allocator() {
880        assert_eq!(purge_all_arenas(), Ok(()));
881    }
882
883    #[test]
884    fn rss_returns_some_on_supported_os() {
885        if cfg!(any(target_os = "linux", target_os = "macos", windows)) {
886            let rss = get_rss_bytes();
887            assert!(rss.is_some(), "RSS should be readable");
888            assert!(rss.unwrap() > 0, "RSS should be > 0");
889        }
890    }
891
892    #[test]
893    fn system_ram_returns_some_on_supported_os() {
894        if cfg!(any(target_os = "linux", target_os = "macos", windows)) {
895            let ram = get_system_ram_bytes();
896            assert!(ram.is_some(), "System RAM should be readable");
897            assert!(ram.unwrap() > 1_000_000, "System RAM should be > 1MB");
898        }
899    }
900
901    #[test]
902    fn snapshot_captures_correctly() {
903        if cfg!(any(target_os = "linux", target_os = "macos", windows)) {
904            let snap = MemorySnapshot::capture();
905            assert!(snap.is_some());
906            let s = snap.unwrap();
907            assert!(s.rss_bytes > 0);
908            assert!(s.system_ram_bytes > s.rss_bytes);
909            assert!(s.rss_percent > 0.0 && s.rss_percent < 100.0);
910        }
911    }
912
913    #[test]
914    fn peak_rss_tracks_maximum() {
915        PEAK_RSS.store(0, Ordering::Relaxed);
916        PEAK_RSS.fetch_max(100, Ordering::Relaxed);
917        PEAK_RSS.fetch_max(50, Ordering::Relaxed);
918        assert_eq!(PEAK_RSS.load(Ordering::Relaxed), 100);
919    }
920
921    #[test]
922    fn pressure_level_roundtrip() {
923        for level in [
924            PressureLevel::Normal,
925            PressureLevel::Soft,
926            PressureLevel::Medium,
927            PressureLevel::Hard,
928            PressureLevel::Critical,
929        ] {
930            assert_eq!(PressureLevel::from_u8(level as u8), level);
931        }
932    }
933
934    #[test]
935    fn process_cap_lowers_but_never_raises_the_percent_limit() {
936        const MB: u64 = 1024 * 1024;
937        assert_eq!(effective_limit(3200 * MB, 0), 3200 * MB);
938        assert_eq!(effective_limit(3200 * MB, 512 * MB), 512 * MB);
939        assert_eq!(effective_limit(256 * MB, 512 * MB), 256 * MB);
940    }
941
942    #[test]
943    fn pressure_ladder_against_a_512mb_cap() {
944        const MB: u64 = 1024 * 1024;
945        let limit = 512 * MB;
946        assert_eq!(pressure_for(500 * MB, limit), PressureLevel::Normal);
947        assert_eq!(pressure_for(520 * MB, limit), PressureLevel::Soft);
948        assert_eq!(pressure_for(650 * MB, limit), PressureLevel::Medium);
949        assert_eq!(pressure_for(800 * MB, limit), PressureLevel::Hard);
950        // The ~1 GB Codex instances land in Critical → emergency drop.
951        assert_eq!(pressure_for(1100 * MB, limit), PressureLevel::Critical);
952        assert_eq!(pressure_for(u64::MAX, 0), PressureLevel::Normal);
953    }
954
955    #[test]
956    fn hard_pressure_requests_abort_immediately() {
957        assert!(!pressure_requests_abort(PressureLevel::Normal));
958        assert!(!pressure_requests_abort(PressureLevel::Soft));
959        assert!(!pressure_requests_abort(PressureLevel::Medium));
960        assert!(pressure_requests_abort(PressureLevel::Hard));
961        assert!(pressure_requests_abort(PressureLevel::Critical));
962    }
963
964    #[test]
965    fn eviction_gate_spaces_fruitless_rounds_but_lets_escalation_through() {
966        let start = std::time::Instant::now();
967        let at = |secs: u64| start + std::time::Duration::from_secs(secs);
968        let mut gate = EvictionGate::default();
969        assert!(gate.due(PressureLevel::Soft, start));
970
971        // Ten fruitless Soft rounds: 1,1,2,4,…,60 s apart, then a 5 min pause.
972        let mut now = start;
973        let mut paused_from = start;
974        let mut schedule = EvictionSchedule::RetryAfter(0);
975        for _ in 0..10 {
976            assert!(gate.due(PressureLevel::Soft, now));
977            paused_from = now;
978            schedule = gate.record(PressureLevel::Soft, false, now);
979            now += std::time::Duration::from_secs(schedule.poll_secs());
980        }
981        assert_eq!(schedule, EvictionSchedule::PauseFor(5 * 60));
982        assert!(!gate.due(
983            PressureLevel::Soft,
984            paused_from + std::time::Duration::from_secs(1)
985        ));
986        assert!(gate.due(
987            PressureLevel::Medium,
988            paused_from + std::time::Duration::from_secs(1)
989        ));
990
991        // A rise to Critical during the pause is evicted at once, with a fresh backoff.
992        assert!(gate.due(
993            PressureLevel::Critical,
994            paused_from + std::time::Duration::from_secs(1)
995        ));
996        assert_eq!(
997            gate.record(PressureLevel::Critical, true, at(1000)),
998            EvictionSchedule::RetryAfter(1)
999        );
1000        assert!(!gate.due(PressureLevel::Critical, at(1000)));
1001        assert!(gate.due(PressureLevel::Critical, at(1001)));
1002    }
1003
1004    #[test]
1005    fn critical_zero_progress_backoff_doubles_then_pauses_and_resets() {
1006        let mut backoff = EvictionBackoff::default();
1007
1008        assert_eq!(backoff.record(false), EvictionSchedule::RetryAfter(1));
1009        assert_eq!(backoff.record(false), EvictionSchedule::RetryAfter(1));
1010        assert_eq!(backoff.record(false), EvictionSchedule::RetryAfter(2));
1011        assert_eq!(backoff.record(false), EvictionSchedule::RetryAfter(4));
1012
1013        for _ in 0..4 {
1014            backoff.record(false);
1015        }
1016        assert_eq!(backoff.record(false), EvictionSchedule::RetryAfter(60));
1017        assert_eq!(backoff.record(false), EvictionSchedule::PauseFor(5 * 60));
1018
1019        assert_eq!(backoff.record(true), EvictionSchedule::RetryAfter(1));
1020        assert_eq!(backoff.consecutive_zero_progress, 0);
1021    }
1022
1023    #[test]
1024    fn atomic_pressure_defaults_to_normal() {
1025        assert_eq!(current_pressure(), PressureLevel::Normal);
1026    }
1027
1028    #[test]
1029    fn start_guard_is_noop_under_test() {
1030        // Regression guard: the background guardian must never run inside the
1031        // test harness. If it did, its 3s poll would observe the suite's large
1032        // RSS, flip the global pressure/abort flags, and silently make unrelated
1033        // tests (e.g. graph edge-building) skip work — an order/timing-dependent
1034        // flake. `start_guard` must be a no-op under `cfg!(test)`.
1035        let fired = Arc::new(AtomicBool::new(false));
1036        let fired_cb = fired.clone();
1037        start_guard(Arc::new(move |_| {
1038            fired_cb.store(true, Ordering::SeqCst);
1039            false
1040        }));
1041
1042        assert!(
1043            !GUARD_RUNNING.load(Ordering::Relaxed),
1044            "guardian thread must not start under cfg!(test)"
1045        );
1046        assert_eq!(current_pressure(), PressureLevel::Normal);
1047        assert!(!abort_requested());
1048        assert!(
1049            !fired.load(Ordering::Relaxed),
1050            "eviction callback must never fire in tests"
1051        );
1052    }
1053
1054    #[test]
1055    fn rss_for_own_pid_matches_self() {
1056        if cfg!(any(target_os = "linux", target_os = "macos", windows)) {
1057            let self_rss = get_rss_bytes().unwrap();
1058            let pid_rss = get_rss_bytes_for_pid(std::process::id()).unwrap();
1059            let ratio = self_rss as f64 / pid_rss as f64;
1060            assert!(
1061                (0.5..2.0).contains(&ratio),
1062                "self RSS ({self_rss}) and pid-based RSS ({pid_rss}) should be within 2x"
1063            );
1064        }
1065    }
1066
1067    #[test]
1068    fn rss_for_dead_pid_returns_none() {
1069        let dead_pid = 999_999_999u32;
1070        assert!(get_rss_bytes_for_pid(dead_pid).is_none());
1071    }
1072
1073    #[test]
1074    fn capture_for_pid_falls_back_on_dead_pid() {
1075        if cfg!(any(target_os = "linux", target_os = "macos", windows)) {
1076            let snap = MemorySnapshot::capture_for_pid(999_999_999);
1077            assert!(snap.is_some(), "should fall back to self RSS");
1078        }
1079    }
1080}