Skip to main content

aft/
memory.rs

1use std::collections::BTreeMap;
2use std::path::Path;
3
4use serde::Serialize;
5use serde_json::Value;
6
7/// A cold-path estimate of memory AFT can attribute without allocator hooks.
8///
9/// `estimated_bytes` is `None` when a subsystem is busy or its resident bytes
10/// are not cheaply observable. Counts remain available in those cases so the
11/// status response never substitutes a fabricated byte estimate.
12#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
13pub struct MemoryEstimate {
14    pub status: &'static str,
15    pub bytes_status: &'static str,
16    pub estimated_bytes: Option<u64>,
17    #[serde(skip_serializing_if = "Vec::is_empty")]
18    pub not_estimated: Vec<String>,
19    #[serde(flatten)]
20    pub counts: BTreeMap<String, u64>,
21}
22
23impl MemoryEstimate {
24    pub fn estimated(bytes: u64) -> Self {
25        Self {
26            status: "ready",
27            bytes_status: "estimated",
28            estimated_bytes: Some(bytes),
29            not_estimated: Vec::new(),
30            counts: BTreeMap::new(),
31        }
32    }
33
34    pub fn partial(bytes: u64) -> Self {
35        Self {
36            status: "ready",
37            bytes_status: "partial",
38            estimated_bytes: Some(bytes),
39            not_estimated: Vec::new(),
40            counts: BTreeMap::new(),
41        }
42    }
43
44    pub fn not_estimated() -> Self {
45        Self {
46            status: "ready",
47            bytes_status: "not_estimated",
48            estimated_bytes: None,
49            not_estimated: Vec::new(),
50            counts: BTreeMap::new(),
51        }
52    }
53
54    pub fn busy() -> Self {
55        Self {
56            status: "busy",
57            bytes_status: "not_estimated",
58            estimated_bytes: None,
59            not_estimated: Vec::new(),
60            counts: BTreeMap::new(),
61        }
62    }
63
64    pub fn count(mut self, name: impl Into<String>, value: usize) -> Self {
65        self.counts.insert(name.into(), usize_to_u64(value));
66        self
67    }
68
69    pub fn count_u64(mut self, name: impl Into<String>, value: u64) -> Self {
70        self.counts.insert(name.into(), value);
71        self
72    }
73
74    pub fn gap(mut self, name: impl Into<String>) -> Self {
75        self.not_estimated.push(name.into());
76        self
77    }
78}
79
80#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
81pub struct RootMemorySnapshot {
82    pub status: &'static str,
83    pub attributed_bytes: u64,
84    pub semantic: MemoryEstimate,
85    pub trigram: MemoryEstimate,
86    pub symbols: MemoryEstimate,
87    pub callgraph: MemoryEstimate,
88    pub callgraph_projection: MemoryEstimate,
89    pub inspect: MemoryEstimate,
90    pub bash: MemoryEstimate,
91    pub lsp: MemoryEstimate,
92    pub parser_pool: MemoryEstimate,
93}
94
95#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
96pub(crate) struct RootMemoryRollup {
97    pub(crate) status: &'static str,
98    pub(crate) attributed_bytes: u64,
99    pub(crate) busy_subsystems: usize,
100    pub(crate) not_estimated_subsystems: usize,
101}
102
103impl RootMemoryRollup {
104    pub(crate) fn from_estimates(estimates: &[&MemoryEstimate]) -> Self {
105        let attributed_bytes = estimates
106            .iter()
107            .filter_map(|estimate| estimate.estimated_bytes)
108            .fold(0u64, u64::saturating_add);
109        let busy_subsystems = estimates
110            .iter()
111            .filter(|estimate| estimate.status == "busy")
112            .count();
113        let not_estimated_subsystems = estimates
114            .iter()
115            .filter(|estimate| estimate.estimated_bytes.is_none())
116            .count();
117        Self {
118            status: if busy_subsystems > 0 { "busy" } else { "ready" },
119            attributed_bytes,
120            busy_subsystems,
121            not_estimated_subsystems,
122        }
123    }
124}
125
126impl RootMemorySnapshot {
127    pub fn new(
128        semantic: MemoryEstimate,
129        trigram: MemoryEstimate,
130        symbols: MemoryEstimate,
131        callgraph: MemoryEstimate,
132        callgraph_projection: MemoryEstimate,
133        inspect: MemoryEstimate,
134        bash: MemoryEstimate,
135        lsp: MemoryEstimate,
136        parser_pool: MemoryEstimate,
137    ) -> Self {
138        let rollup = RootMemoryRollup::from_estimates(&[
139            &semantic,
140            &trigram,
141            &symbols,
142            &callgraph,
143            &callgraph_projection,
144            &inspect,
145            &bash,
146            &lsp,
147            &parser_pool,
148        ]);
149        Self {
150            status: rollup.status,
151            attributed_bytes: rollup.attributed_bytes,
152            semantic,
153            trigram,
154            symbols,
155            callgraph,
156            callgraph_projection,
157            inspect,
158            bash,
159            lsp,
160            parser_pool,
161        }
162    }
163
164    fn rollup(&self) -> RootMemoryRollup {
165        RootMemoryRollup::from_estimates(&[
166            &self.semantic,
167            &self.trigram,
168            &self.symbols,
169            &self.callgraph,
170            &self.callgraph_projection,
171            &self.inspect,
172            &self.bash,
173            &self.lsp,
174            &self.parser_pool,
175        ])
176    }
177
178    pub fn busy_subsystem_count(&self) -> usize {
179        self.rollup().busy_subsystems
180    }
181
182    pub fn not_estimated_subsystem_count(&self) -> usize {
183        self.rollup().not_estimated_subsystems
184    }
185}
186
187#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
188pub struct SqliteMemorySnapshot {
189    pub status: &'static str,
190    pub memory_used_bytes: u64,
191    pub memory_highwater_bytes: u64,
192}
193
194impl SqliteMemorySnapshot {
195    fn measure() -> Self {
196        // SQLite's allocator counters are process-wide and internally synchronized.
197        // They intentionally replace per-connection guesses in root estimates.
198        let memory_used = unsafe { rusqlite::ffi::sqlite3_memory_used() };
199        let memory_highwater = unsafe { rusqlite::ffi::sqlite3_memory_highwater(0) };
200        Self {
201            status: "measured",
202            memory_used_bytes: nonnegative_i64_to_u64(memory_used),
203            memory_highwater_bytes: nonnegative_i64_to_u64(memory_highwater),
204        }
205    }
206}
207
208#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
209pub struct AllocatorMemorySnapshot {
210    pub status: &'static str,
211    pub bytes_in_use: Option<u64>,
212    pub size_allocated: Option<u64>,
213    pub retained_slack_bytes: Option<u64>,
214    #[serde(skip_serializing_if = "Option::is_none")]
215    pub not_estimated: Option<&'static str>,
216}
217
218impl AllocatorMemorySnapshot {
219    #[cfg(any(target_os = "macos", all(target_os = "linux", target_env = "gnu")))]
220    fn measured(bytes_in_use: u64, size_allocated: u64) -> Self {
221        Self {
222            status: "measured",
223            bytes_in_use: Some(bytes_in_use),
224            size_allocated: Some(size_allocated),
225            retained_slack_bytes: Some(size_allocated.saturating_sub(bytes_in_use)),
226            not_estimated: None,
227        }
228    }
229
230    // Not cfg-gated to the fallback platforms: linux-gnu also uses this at
231    // RUNTIME when the host glibc predates mallinfo2 (< 2.33), which only
232    // manifests on release binaries built against an old glibc floor.
233    #[cfg_attr(target_os = "macos", allow(dead_code))]
234    fn not_estimated(reason: &'static str) -> Self {
235        Self {
236            status: "not_estimated_on_this_platform",
237            bytes_in_use: None,
238            size_allocated: None,
239            retained_slack_bytes: None,
240            not_estimated: Some(reason),
241        }
242    }
243}
244
245#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
246pub struct ProcessMemorySnapshot {
247    pub rss_status: &'static str,
248    pub rss_bytes: Option<u64>,
249    /// Kernel physical footprint (macOS `phys_footprint`): dirty + compressed +
250    /// IOKit pages, excluding clean/reclaimable ones. This is the number
251    /// Activity Monitor's "Real Memory" and the OOM killer use. RSS counts
252    /// MADV_FREE pages the allocator has already surrendered (the kernel
253    /// reclaims them lazily), so RSS can read gigabytes above what the process
254    /// actually holds — observed 5.1 GB RSS over a 610 MB footprint. None on
255    /// non-macOS platforms (Linux RSS does not have this skew; MADV_FREE'd
256    /// pages leave Linux RSS on reclaim, not on advice).
257    #[serde(skip_serializing_if = "Option::is_none")]
258    pub phys_footprint_bytes: Option<u64>,
259    #[serde(skip_serializing_if = "Option::is_none")]
260    pub rss_not_estimated: Option<&'static str>,
261    pub sqlite: SqliteMemorySnapshot,
262    /// Allocator bytes overlap the attributed subsystem totals and are an
263    /// allocation envelope, not another amount to subtract from RSS.
264    pub allocator: AllocatorMemorySnapshot,
265    pub total_attributed_bytes: u64,
266    pub unattributed_bytes: Option<i64>,
267    pub root_count: usize,
268    pub busy_subsystems: usize,
269    pub not_estimated_subsystems: usize,
270}
271
272#[derive(Debug, Clone, PartialEq, Eq)]
273pub struct AllocatorPressureRelief {
274    pub bytes_released: u64,
275    pub rss_before_bytes: Option<u64>,
276    pub rss_after_bytes: Option<u64>,
277    pub allocator_before: AllocatorMemorySnapshot,
278    pub allocator_after: AllocatorMemorySnapshot,
279}
280
281impl ProcessMemorySnapshot {
282    pub fn from_roots(
283        roots: &BTreeMap<String, RootMemorySnapshot>,
284        shared_semantic_bases: &MemoryEstimate,
285    ) -> Self {
286        let rollups: Vec<_> = roots.values().map(RootMemorySnapshot::rollup).collect();
287        Self::from_root_rollups(rollups.iter(), roots.len(), shared_semantic_bases)
288    }
289
290    fn from_root_rollups<'a>(
291        roots: impl Iterator<Item = &'a RootMemoryRollup>,
292        root_count: usize,
293        shared_semantic_bases: &MemoryEstimate,
294    ) -> Self {
295        let mut root_attributed_bytes = 0u64;
296        let mut busy_subsystems = 0usize;
297        let mut not_estimated_subsystems = 0usize;
298        for root in roots {
299            root_attributed_bytes = root_attributed_bytes.saturating_add(root.attributed_bytes);
300            busy_subsystems = busy_subsystems.saturating_add(root.busy_subsystems);
301            not_estimated_subsystems =
302                not_estimated_subsystems.saturating_add(root.not_estimated_subsystems);
303        }
304        let sqlite = SqliteMemorySnapshot::measure();
305        let allocator = allocator_memory_snapshot();
306        let total_attributed_bytes = root_attributed_bytes
307            .saturating_add(shared_semantic_bases.estimated_bytes.unwrap_or(0))
308            .saturating_add(sqlite.memory_used_bytes);
309        let rss_bytes = process_rss_bytes();
310        let phys_footprint_bytes = process_phys_footprint_bytes();
311        // Attribute against the footprint when available: it excludes
312        // already-surrendered pages, so the residual actually means
313        // "held memory we cannot explain" instead of allocator noise.
314        let unattributed_basis = phys_footprint_bytes.or(rss_bytes);
315        let unattributed_bytes =
316            unattributed_basis.map(|held| signed_difference(held, total_attributed_bytes));
317        Self {
318            rss_status: if rss_bytes.is_some() {
319                "estimated"
320            } else {
321                "not_estimated_on_this_platform"
322            },
323            rss_bytes,
324            phys_footprint_bytes,
325            rss_not_estimated: rss_bytes
326                .is_none()
327                .then_some("platform_process_rss_unavailable"),
328            sqlite,
329            allocator,
330            total_attributed_bytes,
331            unattributed_bytes,
332            root_count,
333            busy_subsystems,
334            not_estimated_subsystems,
335        }
336    }
337}
338
339/// Cap on per-root detail entries in serialized snapshots. Process totals
340/// always cover every root; only the per-root breakdown is capped so a
341/// many-root daemon process cannot balloon the status payload past
342/// downstream consumers' size limits (the daemon metrics cache truncates
343/// around 27 KB, and JSON keys serialize alphabetically, so an oversized
344/// `memory.roots` map pushes later sections past the cut).
345pub const MEMORY_SNAPSHOT_ROOT_DETAIL_CAP: usize = 8;
346
347#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
348pub struct MemorySnapshot {
349    pub roots_status: &'static str,
350    /// Top roots by attributed bytes, capped at
351    /// [`MEMORY_SNAPSHOT_ROOT_DETAIL_CAP`]; the remainder is summarized by
352    /// `roots_omitted` / `roots_omitted_bytes`.
353    pub roots: BTreeMap<String, RootMemorySnapshot>,
354    /// Total roots attributed (including omitted ones).
355    pub roots_total: usize,
356    /// Roots summarized out of the detail map.
357    pub roots_omitted: usize,
358    /// Attributed bytes carried by the omitted roots (already included in
359    /// `process.total_attributed_bytes`).
360    pub roots_omitted_bytes: u64,
361    /// Immutable borrowed semantic snapshots, attributed once process-wide.
362    pub shared_semantic_bases: MemoryEstimate,
363    pub process: ProcessMemorySnapshot,
364}
365
366#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
367pub(crate) struct MemoryRollupSnapshot {
368    pub(crate) roots_status: &'static str,
369    pub(crate) roots: BTreeMap<String, RootMemoryRollup>,
370    pub(crate) roots_total: usize,
371    pub(crate) roots_omitted: usize,
372    pub(crate) roots_omitted_bytes: u64,
373    pub(crate) process: ProcessMemorySnapshot,
374}
375
376impl MemoryRollupSnapshot {
377    pub(crate) fn new(
378        roots_status: &'static str,
379        roots: BTreeMap<String, RootMemoryRollup>,
380    ) -> Self {
381        let shared_semantic_bases = crate::semantic_index::shared_semantic_bases_memory();
382        let process = ProcessMemorySnapshot::from_root_rollups(
383            roots.values(),
384            roots.len(),
385            &shared_semantic_bases,
386        );
387        let roots_total = roots.len();
388        let (roots, roots_omitted, roots_omitted_bytes) =
389            cap_root_rollups(roots, MEMORY_SNAPSHOT_ROOT_DETAIL_CAP);
390        Self {
391            roots_status,
392            roots,
393            roots_total,
394            roots_omitted,
395            roots_omitted_bytes,
396            process,
397        }
398    }
399}
400
401impl MemorySnapshot {
402    pub fn new(roots_status: &'static str, roots: BTreeMap<String, RootMemorySnapshot>) -> Self {
403        let shared_semantic_bases = crate::semantic_index::shared_semantic_bases_memory();
404        // Totals cover EVERY root before the detail map is capped.
405        let process = ProcessMemorySnapshot::from_roots(&roots, &shared_semantic_bases);
406        let roots_total = roots.len();
407        let (roots, roots_omitted, roots_omitted_bytes) =
408            cap_root_detail(roots, MEMORY_SNAPSHOT_ROOT_DETAIL_CAP);
409        Self {
410            roots_status,
411            roots,
412            roots_total,
413            roots_omitted,
414            roots_omitted_bytes,
415            shared_semantic_bases,
416            process,
417        }
418    }
419}
420
421fn cap_roots<T>(
422    roots: BTreeMap<String, T>,
423    cap: usize,
424    attributed_bytes: impl Fn(&T) -> u64,
425) -> (BTreeMap<String, T>, usize, u64) {
426    if roots.len() <= cap {
427        return (roots, 0, 0);
428    }
429    let mut entries: Vec<_> = roots.into_iter().collect();
430    entries.sort_by(|a, b| {
431        attributed_bytes(&b.1)
432            .cmp(&attributed_bytes(&a.1))
433            .then_with(|| a.0.cmp(&b.0))
434    });
435    let omitted = entries.split_off(cap);
436    let omitted_bytes = omitted
437        .iter()
438        .map(|(_, snapshot)| attributed_bytes(snapshot))
439        .fold(0u64, u64::saturating_add);
440    (entries.into_iter().collect(), omitted.len(), omitted_bytes)
441}
442
443fn cap_root_rollups(
444    roots: BTreeMap<String, RootMemoryRollup>,
445    cap: usize,
446) -> (BTreeMap<String, RootMemoryRollup>, usize, u64) {
447    cap_roots(roots, cap, |snapshot| snapshot.attributed_bytes)
448}
449
450/// Keep the `cap` roots with the highest attributed bytes; report the rest
451/// as an omitted-count + omitted-bytes rollup.
452fn cap_root_detail(
453    roots: BTreeMap<String, RootMemorySnapshot>,
454    cap: usize,
455) -> (BTreeMap<String, RootMemorySnapshot>, usize, u64) {
456    cap_roots(roots, cap, |snapshot| snapshot.attributed_bytes)
457}
458
459#[cfg(test)]
460mod snapshot_cap_tests {
461    use super::*;
462
463    fn root_with_bytes(bytes: u64) -> RootMemorySnapshot {
464        let estimate = MemoryEstimate::estimated;
465        RootMemorySnapshot::new(
466            estimate(bytes),
467            estimate(0),
468            estimate(0),
469            estimate(0),
470            estimate(0),
471            estimate(0),
472            estimate(0),
473            estimate(0),
474            estimate(0),
475        )
476    }
477
478    #[test]
479    fn detail_map_capped_but_totals_cover_all_roots() {
480        let mut roots = BTreeMap::new();
481        for i in 0..(MEMORY_SNAPSHOT_ROOT_DETAIL_CAP + 4) {
482            // Distinct sizes so the kept set is deterministic: later roots larger.
483            roots.insert(
484                format!("/root/{i:02}"),
485                root_with_bytes((i as u64 + 1) * 1000),
486            );
487        }
488        let snapshot = MemorySnapshot::new("ready", roots);
489
490        assert_eq!(snapshot.roots.len(), MEMORY_SNAPSHOT_ROOT_DETAIL_CAP);
491        assert_eq!(snapshot.roots_total, MEMORY_SNAPSHOT_ROOT_DETAIL_CAP + 4);
492        assert_eq!(snapshot.roots_omitted, 4);
493        // The four smallest (1000..=4000) are the omitted ones.
494        assert_eq!(snapshot.roots_omitted_bytes, 1000 + 2000 + 3000 + 4000);
495        // Largest roots are the ones kept.
496        assert!(snapshot
497            .roots
498            .values()
499            .all(|root| root.attributed_bytes > 4000));
500        // Process totals include omitted roots' bytes.
501        let expected_total: u64 = (1..=(MEMORY_SNAPSHOT_ROOT_DETAIL_CAP as u64 + 4))
502            .map(|i| i * 1000)
503            .sum();
504        assert!(snapshot.process.total_attributed_bytes >= expected_total);
505    }
506
507    #[test]
508    fn rollup_cap_keeps_only_top_roots_while_totals_cover_all() {
509        let mut roots = BTreeMap::new();
510        for i in 0..(MEMORY_SNAPSHOT_ROOT_DETAIL_CAP + 4) {
511            roots.insert(
512                format!("/rollup/{i:02}"),
513                root_with_bytes((i as u64 + 1) * 1000).rollup(),
514            );
515        }
516        let snapshot = MemoryRollupSnapshot::new("ready", roots);
517
518        assert_eq!(snapshot.roots.len(), MEMORY_SNAPSHOT_ROOT_DETAIL_CAP);
519        assert_eq!(snapshot.roots_total, MEMORY_SNAPSHOT_ROOT_DETAIL_CAP + 4);
520        assert_eq!(snapshot.roots_omitted, 4);
521        assert_eq!(snapshot.roots_omitted_bytes, 1000 + 2000 + 3000 + 4000);
522        let expected_total: u64 = (1..=(MEMORY_SNAPSHOT_ROOT_DETAIL_CAP as u64 + 4))
523            .map(|i| i * 1000)
524            .sum();
525        assert!(snapshot.process.total_attributed_bytes >= expected_total);
526    }
527
528    #[test]
529    fn under_cap_keeps_everything_with_zero_omitted() {
530        let mut roots = BTreeMap::new();
531        roots.insert("/a".to_string(), root_with_bytes(10));
532        roots.insert("/b".to_string(), root_with_bytes(20));
533        let snapshot = MemorySnapshot::new("ready", roots);
534        assert_eq!(snapshot.roots.len(), 2);
535        assert_eq!(snapshot.roots_total, 2);
536        assert_eq!(snapshot.roots_omitted, 0);
537        assert_eq!(snapshot.roots_omitted_bytes, 0);
538    }
539}
540
541pub fn path_bytes(path: &Path) -> u64 {
542    #[cfg(unix)]
543    {
544        use std::os::unix::ffi::OsStrExt;
545        usize_to_u64(path.as_os_str().as_bytes().len())
546    }
547    #[cfg(windows)]
548    {
549        use std::os::windows::ffi::OsStrExt;
550        usize_to_u64(path.as_os_str().encode_wide().count())
551            .saturating_mul(std::mem::size_of::<u16>() as u64)
552    }
553    #[cfg(not(any(unix, windows)))]
554    {
555        usize_to_u64(path.to_string_lossy().len())
556    }
557}
558
559pub fn usize_to_u64(value: usize) -> u64 {
560    u64::try_from(value).unwrap_or(u64::MAX)
561}
562
563pub fn estimated_json_bytes(value: &Value) -> u64 {
564    match value {
565        Value::Null => 0,
566        Value::Bool(_) => std::mem::size_of::<bool>() as u64,
567        Value::Number(_) => std::mem::size_of::<serde_json::Number>() as u64,
568        Value::String(value) => usize_to_u64(value.len()),
569        Value::Array(values) => values
570            .iter()
571            .map(estimated_json_bytes)
572            .fold(0u64, u64::saturating_add),
573        Value::Object(values) => values.iter().fold(0u64, |bytes, (key, value)| {
574            bytes
575                .saturating_add(usize_to_u64(key.len()))
576                .saturating_add(estimated_json_bytes(value))
577        }),
578    }
579}
580
581fn signed_difference(lhs: u64, rhs: u64) -> i64 {
582    let difference = i128::from(lhs) - i128::from(rhs);
583    difference.clamp(i128::from(i64::MIN), i128::from(i64::MAX)) as i64
584}
585
586fn nonnegative_i64_to_u64(value: i64) -> u64 {
587    u64::try_from(value).unwrap_or(0)
588}
589
590#[cfg(target_os = "macos")]
591fn allocator_memory_snapshot() -> AllocatorMemorySnapshot {
592    let mut statistics = std::mem::MaybeUninit::<libc::malloc_statistics_t>::zeroed();
593    unsafe {
594        libc::malloc_zone_statistics(libc::malloc_default_zone(), statistics.as_mut_ptr());
595    }
596    let statistics = unsafe { statistics.assume_init() };
597    AllocatorMemorySnapshot::measured(
598        usize_to_u64(statistics.size_in_use),
599        usize_to_u64(statistics.size_allocated),
600    )
601}
602
603#[cfg(all(target_os = "linux", target_env = "gnu"))]
604fn allocator_memory_snapshot() -> AllocatorMemorySnapshot {
605    // mallinfo2 exists only in glibc >= 2.33. Release Linux binaries link
606    // against an older glibc floor (cross gnu images, kept old so dlopen and
607    // wide distro compatibility hold), so a link-time reference to the symbol
608    // fails the release build even though native CI (glibc 2.35) links fine.
609    // Resolve it at runtime instead and report honestly when it is absent.
610    use std::sync::OnceLock;
611    type Mallinfo2Fn = unsafe extern "C" fn() -> libc::mallinfo2;
612    static MALLINFO2: OnceLock<Option<Mallinfo2Fn>> = OnceLock::new();
613    let resolved = MALLINFO2.get_or_init(|| {
614        let symbol = unsafe { libc::dlsym(libc::RTLD_DEFAULT, c"mallinfo2".as_ptr()) };
615        if symbol.is_null() {
616            None
617        } else {
618            // SAFETY: glibc declares mallinfo2 as `struct mallinfo2 (*)(void)`;
619            // the signature matches Mallinfo2Fn exactly.
620            Some(unsafe { std::mem::transmute::<*mut libc::c_void, Mallinfo2Fn>(symbol) })
621        }
622    });
623    let Some(mallinfo2) = resolved else {
624        return AllocatorMemorySnapshot::not_estimated("mallinfo2_requires_glibc_2_33");
625    };
626    let statistics = unsafe { mallinfo2() };
627    let mapped_bytes = statistics.hblkhd as u64;
628    let bytes_in_use = (statistics.uordblks as u64).saturating_add(mapped_bytes);
629    let size_allocated = (statistics.arena as u64).saturating_add(mapped_bytes);
630    AllocatorMemorySnapshot::measured(bytes_in_use, size_allocated)
631}
632
633#[cfg(all(target_os = "linux", target_env = "gnu"))]
634type MallocTrimFn = unsafe extern "C" fn(libc::size_t) -> libc::c_int;
635
636/// Resolve glibc's optional trimming primitive without creating a link-time
637/// dependency on a symbol that musl and alternate allocators do not provide.
638#[cfg(all(target_os = "linux", target_env = "gnu"))]
639fn resolved_malloc_trim() -> Option<MallocTrimFn> {
640    use std::sync::OnceLock;
641    static MALLOC_TRIM: OnceLock<Option<MallocTrimFn>> = OnceLock::new();
642    MALLOC_TRIM
643        .get_or_init(|| {
644            let symbol = unsafe { libc::dlsym(libc::RTLD_DEFAULT, c"malloc_trim".as_ptr()) };
645            if symbol.is_null() {
646                None
647            } else {
648                // SAFETY: glibc declares malloc_trim as `int (size_t)`;
649                // the signature matches MallocTrimFn exactly.
650                Some(unsafe { std::mem::transmute::<*mut libc::c_void, MallocTrimFn>(symbol) })
651            }
652        })
653        .as_ref()
654        .copied()
655}
656
657#[cfg(not(any(target_os = "macos", all(target_os = "linux", target_env = "gnu"))))]
658fn allocator_memory_snapshot() -> AllocatorMemorySnapshot {
659    AllocatorMemorySnapshot::not_estimated("platform_allocator_statistics_unavailable")
660}
661
662#[cfg(target_os = "macos")]
663unsafe extern "C" {
664    fn malloc_zone_pressure_relief(zone: *mut libc::malloc_zone_t, goal: usize) -> usize;
665}
666
667/// Allocator slack (mapped-but-unused arena bytes) above which opportunistic
668/// pressure relief is worth the zone-lock contention it briefly causes.
669pub const ALLOCATOR_SLACK_RELIEF_THRESHOLD_BYTES: u64 = 1024 * 1024 * 1024;
670
671/// Minimum spacing between opportunistic relief passes so a workload that
672/// legitimately cycles through large allocations does not thrash the allocator.
673pub const ALLOCATOR_SLACK_RELIEF_MIN_INTERVAL: std::time::Duration =
674    std::time::Duration::from_secs(300);
675
676/// Decide whether an opportunistic allocator relief pass is due.
677///
678/// Pure so the policy is unit-testable: fires only when the allocator reports
679/// at least the threshold of retained slack AND the previous pass is old
680/// enough. Callers own actually measuring the snapshot and running the pass.
681pub fn allocator_slack_relief_due(
682    retained_slack_bytes: Option<u64>,
683    last_relief: Option<std::time::Instant>,
684    now: std::time::Instant,
685) -> bool {
686    let Some(slack) = retained_slack_bytes else {
687        return false;
688    };
689    if slack < ALLOCATOR_SLACK_RELIEF_THRESHOLD_BYTES {
690        return false;
691    }
692    match last_relief {
693        None => true,
694        Some(at) => now.duration_since(at) >= ALLOCATOR_SLACK_RELIEF_MIN_INTERVAL,
695    }
696}
697
698/// Opportunistically return unused allocator pages when slack is large, even
699/// while sessions are active. The whole-process idle sweep only fires when
700/// every root has been quiet, so one long-lived chatty session used to block
701/// reclamation for the process lifetime (observed: 5.1 GB RSS over ~600 MB of
702/// live data). Runs the relief on a detached thread because allocator trimming
703/// walks allocator state under its lock and must not stall the dispatch loop or
704/// health probes.
705///
706/// Returns true when a pass was spawned (caller records the timestamp).
707#[cfg(any(target_os = "macos", target_os = "linux"))]
708pub fn spawn_allocator_slack_relief_if_due(
709    last_relief: Option<std::time::Instant>,
710    now: std::time::Instant,
711) -> bool {
712    let slack = allocator_memory_snapshot().retained_slack_bytes;
713    if !allocator_slack_relief_due(slack, last_relief, now) {
714        return false;
715    }
716    std::thread::Builder::new()
717        .name("aft-mem-relief".to_string())
718        .spawn(|| {
719            let relief = relieve_allocator_pressure();
720            log::info!(
721                "allocator slack relief: released={} allocator_slack_bytes_before={:?} allocator_slack_bytes_after={:?} rss_bytes_before={:?} rss_bytes_after={:?}",
722                relief.bytes_released,
723                relief.allocator_before.retained_slack_bytes,
724                relief.allocator_after.retained_slack_bytes,
725                relief.rss_before_bytes,
726                relief.rss_after_bytes,
727            );
728        })
729        .is_ok()
730}
731
732/// Ask the platform allocator to return unused pages after a process-wide idle
733/// gate. Callers own that gate because allocator pressure relief can add
734/// latency. Linux invokes glibc's optional `malloc_trim(0)` when the symbol is
735/// available; non-glibc allocators intentionally remain a no-op.
736#[cfg(target_os = "macos")]
737pub fn relieve_allocator_pressure() -> AllocatorPressureRelief {
738    let rss_before_bytes = process_rss_bytes();
739    let allocator_before = allocator_memory_snapshot();
740    let bytes_released = unsafe { malloc_zone_pressure_relief(std::ptr::null_mut(), 0) };
741    let allocator_after = allocator_memory_snapshot();
742    let rss_after_bytes = process_rss_bytes();
743    AllocatorPressureRelief {
744        bytes_released: usize_to_u64(bytes_released),
745        rss_before_bytes,
746        rss_after_bytes,
747        allocator_before,
748        allocator_after,
749    }
750}
751
752#[cfg(target_os = "linux")]
753pub fn relieve_allocator_pressure() -> AllocatorPressureRelief {
754    let rss_before_bytes = process_rss_bytes();
755    let allocator_before = allocator_memory_snapshot();
756    #[cfg(target_env = "gnu")]
757    if let Some(malloc_trim) = resolved_malloc_trim() {
758        // SAFETY: resolved_malloc_trim verifies the symbol and its C ABI
759        // signature before returning the function pointer.
760        unsafe { malloc_trim(0) };
761    }
762    let allocator_after = allocator_memory_snapshot();
763    let rss_after_bytes = process_rss_bytes();
764    let bytes_released = allocator_before
765        .size_allocated
766        .zip(allocator_after.size_allocated)
767        .map(|(before, after)| before.saturating_sub(after))
768        .unwrap_or(0);
769    AllocatorPressureRelief {
770        bytes_released,
771        rss_before_bytes,
772        rss_after_bytes,
773        allocator_before,
774        allocator_after,
775    }
776}
777
778#[cfg(target_os = "macos")]
779fn process_rss_bytes() -> Option<u64> {
780    let mut info = std::mem::MaybeUninit::<libc::proc_taskinfo>::zeroed();
781    let size = std::mem::size_of::<libc::proc_taskinfo>();
782    let written = unsafe {
783        libc::proc_pidinfo(
784            libc::getpid(),
785            libc::PROC_PIDTASKINFO,
786            0,
787            info.as_mut_ptr().cast(),
788            i32::try_from(size).ok()?,
789        )
790    };
791    if written != i32::try_from(size).ok()? {
792        return None;
793    }
794    Some(unsafe { info.assume_init() }.pti_resident_size)
795}
796
797/// Kernel physical footprint via `proc_pid_rusage` (`ri_phys_footprint`).
798/// See `phys_footprint_bytes` for why this, not RSS, is the headline number.
799#[cfg(target_os = "macos")]
800fn process_phys_footprint_bytes() -> Option<u64> {
801    let mut usage = std::mem::MaybeUninit::<libc::rusage_info_v4>::zeroed();
802    let rc = unsafe {
803        libc::proc_pid_rusage(
804            libc::getpid(),
805            libc::RUSAGE_INFO_V4,
806            usage.as_mut_ptr().cast(),
807        )
808    };
809    if rc != 0 {
810        return None;
811    }
812    Some(unsafe { usage.assume_init() }.ri_phys_footprint)
813}
814
815#[cfg(not(target_os = "macos"))]
816fn process_phys_footprint_bytes() -> Option<u64> {
817    None
818}
819
820#[cfg(target_os = "linux")]
821fn process_rss_bytes() -> Option<u64> {
822    let statm = std::fs::read_to_string("/proc/self/statm").ok()?;
823    let resident_pages = statm.split_whitespace().nth(1)?.parse::<u64>().ok()?;
824    let page_size = unsafe { libc::sysconf(libc::_SC_PAGESIZE) };
825    if page_size <= 0 {
826        return None;
827    }
828    resident_pages.checked_mul(page_size as u64)
829}
830
831#[cfg(not(any(target_os = "macos", target_os = "linux")))]
832fn process_rss_bytes() -> Option<u64> {
833    None
834}
835
836#[cfg(test)]
837mod tests {
838    use super::*;
839
840    #[test]
841    fn process_snapshot_preserves_negative_residuals() {
842        assert_eq!(signed_difference(5, 8), -3);
843    }
844
845    #[test]
846    fn slack_relief_fires_on_large_slack_and_respects_spacing() {
847        use std::time::{Duration, Instant};
848        let now = Instant::now();
849        let big = Some(ALLOCATOR_SLACK_RELIEF_THRESHOLD_BYTES);
850        // Unknown slack (allocator stats unavailable) never fires.
851        assert!(!allocator_slack_relief_due(None, None, now));
852        // Below threshold never fires.
853        assert!(!allocator_slack_relief_due(
854            Some(ALLOCATOR_SLACK_RELIEF_THRESHOLD_BYTES - 1),
855            None,
856            now
857        ));
858        // At threshold with no prior pass fires.
859        assert!(allocator_slack_relief_due(big, None, now));
860        // A recent pass suppresses the next one...
861        let recent = now - Duration::from_secs(10);
862        assert!(!allocator_slack_relief_due(big, Some(recent), now));
863        // ...until the minimum spacing has elapsed.
864        let stale = now - ALLOCATOR_SLACK_RELIEF_MIN_INTERVAL;
865        assert!(allocator_slack_relief_due(big, Some(stale), now));
866    }
867
868    #[test]
869    fn json_estimator_scales_with_payload_content() {
870        let empty = estimated_json_bytes(&serde_json::json!({}));
871        let populated = estimated_json_bytes(&serde_json::json!({"message": "hello"}));
872        assert_eq!(empty, 0);
873        assert!(populated >= 12);
874    }
875
876    #[test]
877    fn process_snapshot_exposes_sqlite_and_allocator_sections() {
878        let shared = MemoryEstimate::estimated(7);
879        let snapshot = ProcessMemorySnapshot::from_roots(&BTreeMap::new(), &shared);
880        assert_eq!(snapshot.sqlite.status, "measured");
881        assert!(snapshot.sqlite.memory_highwater_bytes >= snapshot.sqlite.memory_used_bytes);
882        assert_eq!(
883            snapshot.total_attributed_bytes,
884            snapshot.sqlite.memory_used_bytes.saturating_add(7)
885        );
886
887        let serialized = serde_json::to_value(&snapshot).expect("serialize process memory");
888        assert!(serialized["sqlite"]["memory_used_bytes"].is_u64());
889        assert!(serialized["allocator"].get("bytes_in_use").is_some());
890        assert!(serialized["allocator"].get("size_allocated").is_some());
891        assert!(serialized["allocator"]
892            .get("retained_slack_bytes")
893            .is_some());
894    }
895
896    #[cfg(any(target_os = "macos", all(target_os = "linux", target_env = "gnu")))]
897    #[test]
898    fn allocator_snapshot_reports_measured_slack() {
899        let allocator = allocator_memory_snapshot();
900        if allocator.status == "measured" {
901            let in_use = allocator.bytes_in_use.expect("allocator bytes in use");
902            let allocated = allocator.size_allocated.expect("allocator size allocated");
903            assert_eq!(
904                allocator.retained_slack_bytes,
905                Some(allocated.saturating_sub(in_use))
906            );
907        } else {
908            assert_eq!(allocator.status, "not_estimated_on_this_platform");
909            assert_eq!(allocator.bytes_in_use, None);
910            assert_eq!(allocator.size_allocated, None);
911            assert_eq!(allocator.retained_slack_bytes, None);
912            assert_eq!(
913                allocator.not_estimated,
914                Some("mallinfo2_requires_glibc_2_33")
915            );
916        }
917    }
918
919    #[cfg(target_os = "linux")]
920    #[test]
921    fn linux_allocator_pressure_relief_smoke() {
922        let mut allocation = vec![0u8; 32 * 1024 * 1024];
923        for byte in allocation.iter_mut().step_by(4096) {
924            *byte = 1;
925        }
926        std::hint::black_box(&allocation);
927        drop(allocation);
928
929        let relief = relieve_allocator_pressure();
930        std::hint::black_box(relief);
931
932        #[cfg(target_env = "gnu")]
933        assert!(
934            resolved_malloc_trim().is_some(),
935            "glibc malloc_trim must be available for the Linux relief path"
936        );
937    }
938
939    #[cfg(not(any(target_os = "macos", all(target_os = "linux", target_env = "gnu"))))]
940    #[test]
941    fn allocator_snapshot_is_honest_when_platform_counters_are_unavailable() {
942        let allocator = allocator_memory_snapshot();
943        assert_eq!(allocator.status, "not_estimated_on_this_platform");
944        assert_eq!(allocator.bytes_in_use, None);
945        assert_eq!(allocator.size_allocated, None);
946        assert_eq!(allocator.retained_slack_bytes, None);
947        assert_eq!(
948            allocator.not_estimated,
949            Some("platform_allocator_statistics_unavailable")
950        );
951    }
952
953    #[cfg(target_os = "macos")]
954    #[test]
955    #[ignore = "bounded live RSS experiment; run explicitly after allocator changes"]
956    fn allocator_pressure_relief_warm_then_idle_measurement() {
957        let warm_pages = (0..16 * 1024)
958            .map(|seed| {
959                let mut page = Box::new([0u8; 4096]);
960                page[0] = seed as u8;
961                page
962            })
963            .collect::<Vec<_>>();
964        std::hint::black_box(&warm_pages);
965        drop(warm_pages);
966
967        let relief = relieve_allocator_pressure();
968        let sqlite = SqliteMemorySnapshot::measure();
969        eprintln!(
970            "warm-then-idle pressure relief: rss_before={:?} rss_after={:?} allocator_in_use_before={:?} allocator_in_use_after={:?} allocator_allocated_before={:?} allocator_allocated_after={:?} allocator_slack_before={:?} allocator_slack_after={:?} allocator_reported_released={} sqlite_used={} sqlite_highwater={}",
971            relief.rss_before_bytes,
972            relief.rss_after_bytes,
973            relief.allocator_before.bytes_in_use,
974            relief.allocator_after.bytes_in_use,
975            relief.allocator_before.size_allocated,
976            relief.allocator_after.size_allocated,
977            relief.allocator_before.retained_slack_bytes,
978            relief.allocator_after.retained_slack_bytes,
979            relief.bytes_released,
980            sqlite.memory_used_bytes,
981            sqlite.memory_highwater_bytes,
982        );
983        assert_eq!(relief.allocator_before.status, "measured");
984        assert_eq!(relief.allocator_after.status, "measured");
985    }
986}