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 inspect: MemoryEstimate,
89    pub bash: MemoryEstimate,
90    pub lsp: MemoryEstimate,
91    pub parser_pool: MemoryEstimate,
92}
93
94impl RootMemorySnapshot {
95    pub fn new(
96        semantic: MemoryEstimate,
97        trigram: MemoryEstimate,
98        symbols: MemoryEstimate,
99        callgraph: MemoryEstimate,
100        inspect: MemoryEstimate,
101        bash: MemoryEstimate,
102        lsp: MemoryEstimate,
103        parser_pool: MemoryEstimate,
104    ) -> Self {
105        let estimates = [
106            &semantic,
107            &trigram,
108            &symbols,
109            &callgraph,
110            &inspect,
111            &bash,
112            &lsp,
113            &parser_pool,
114        ];
115        let attributed_bytes = estimates
116            .iter()
117            .filter_map(|estimate| estimate.estimated_bytes)
118            .fold(0u64, u64::saturating_add);
119        let status = if estimates.iter().any(|estimate| estimate.status == "busy") {
120            "busy"
121        } else {
122            "ready"
123        };
124        Self {
125            status,
126            attributed_bytes,
127            semantic,
128            trigram,
129            symbols,
130            callgraph,
131            inspect,
132            bash,
133            lsp,
134            parser_pool,
135        }
136    }
137
138    pub fn busy_subsystem_count(&self) -> usize {
139        self.estimates()
140            .iter()
141            .filter(|estimate| estimate.status == "busy")
142            .count()
143    }
144
145    pub fn not_estimated_subsystem_count(&self) -> usize {
146        self.estimates()
147            .iter()
148            .filter(|estimate| estimate.estimated_bytes.is_none())
149            .count()
150    }
151
152    fn estimates(&self) -> [&MemoryEstimate; 8] {
153        [
154            &self.semantic,
155            &self.trigram,
156            &self.symbols,
157            &self.callgraph,
158            &self.inspect,
159            &self.bash,
160            &self.lsp,
161            &self.parser_pool,
162        ]
163    }
164}
165
166#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
167pub struct SqliteMemorySnapshot {
168    pub status: &'static str,
169    pub memory_used_bytes: u64,
170    pub memory_highwater_bytes: u64,
171}
172
173impl SqliteMemorySnapshot {
174    fn measure() -> Self {
175        // SQLite's allocator counters are process-wide and internally synchronized.
176        // They intentionally replace per-connection guesses in root estimates.
177        let memory_used = unsafe { rusqlite::ffi::sqlite3_memory_used() };
178        let memory_highwater = unsafe { rusqlite::ffi::sqlite3_memory_highwater(0) };
179        Self {
180            status: "measured",
181            memory_used_bytes: nonnegative_i64_to_u64(memory_used),
182            memory_highwater_bytes: nonnegative_i64_to_u64(memory_highwater),
183        }
184    }
185}
186
187#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
188pub struct AllocatorMemorySnapshot {
189    pub status: &'static str,
190    pub bytes_in_use: Option<u64>,
191    pub size_allocated: Option<u64>,
192    pub retained_slack_bytes: Option<u64>,
193    #[serde(skip_serializing_if = "Option::is_none")]
194    pub not_estimated: Option<&'static str>,
195}
196
197impl AllocatorMemorySnapshot {
198    #[cfg(any(target_os = "macos", all(target_os = "linux", target_env = "gnu")))]
199    fn measured(bytes_in_use: u64, size_allocated: u64) -> Self {
200        Self {
201            status: "measured",
202            bytes_in_use: Some(bytes_in_use),
203            size_allocated: Some(size_allocated),
204            retained_slack_bytes: Some(size_allocated.saturating_sub(bytes_in_use)),
205            not_estimated: None,
206        }
207    }
208
209    // Not cfg-gated to the fallback platforms: linux-gnu also uses this at
210    // RUNTIME when the host glibc predates mallinfo2 (< 2.33), which only
211    // manifests on release binaries built against an old glibc floor.
212    #[cfg_attr(target_os = "macos", allow(dead_code))]
213    fn not_estimated(reason: &'static str) -> Self {
214        Self {
215            status: "not_estimated_on_this_platform",
216            bytes_in_use: None,
217            size_allocated: None,
218            retained_slack_bytes: None,
219            not_estimated: Some(reason),
220        }
221    }
222}
223
224#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
225pub struct ProcessMemorySnapshot {
226    pub rss_status: &'static str,
227    pub rss_bytes: Option<u64>,
228    #[serde(skip_serializing_if = "Option::is_none")]
229    pub rss_not_estimated: Option<&'static str>,
230    pub sqlite: SqliteMemorySnapshot,
231    /// Allocator bytes overlap the attributed subsystem totals and are an
232    /// allocation envelope, not another amount to subtract from RSS.
233    pub allocator: AllocatorMemorySnapshot,
234    pub total_attributed_bytes: u64,
235    pub unattributed_bytes: Option<i64>,
236    pub root_count: usize,
237    pub busy_subsystems: usize,
238    pub not_estimated_subsystems: usize,
239}
240
241#[derive(Debug, Clone, PartialEq, Eq)]
242pub struct AllocatorPressureRelief {
243    pub bytes_released: u64,
244    pub rss_before_bytes: Option<u64>,
245    pub rss_after_bytes: Option<u64>,
246    pub allocator_before: AllocatorMemorySnapshot,
247    pub allocator_after: AllocatorMemorySnapshot,
248}
249
250impl ProcessMemorySnapshot {
251    pub fn from_roots(
252        roots: &BTreeMap<String, RootMemorySnapshot>,
253        shared_semantic_bases: &MemoryEstimate,
254    ) -> Self {
255        let sqlite = SqliteMemorySnapshot::measure();
256        let allocator = allocator_memory_snapshot();
257        let total_attributed_bytes = roots
258            .values()
259            .map(|root| root.attributed_bytes)
260            .fold(0u64, u64::saturating_add)
261            .saturating_add(shared_semantic_bases.estimated_bytes.unwrap_or(0))
262            .saturating_add(sqlite.memory_used_bytes);
263        let busy_subsystems = roots
264            .values()
265            .map(RootMemorySnapshot::busy_subsystem_count)
266            .sum();
267        let not_estimated_subsystems = roots
268            .values()
269            .map(RootMemorySnapshot::not_estimated_subsystem_count)
270            .sum();
271        let rss_bytes = process_rss_bytes();
272        let unattributed_bytes =
273            rss_bytes.map(|rss| signed_difference(rss, total_attributed_bytes));
274        Self {
275            rss_status: if rss_bytes.is_some() {
276                "estimated"
277            } else {
278                "not_estimated_on_this_platform"
279            },
280            rss_bytes,
281            rss_not_estimated: rss_bytes
282                .is_none()
283                .then_some("platform_process_rss_unavailable"),
284            sqlite,
285            allocator,
286            total_attributed_bytes,
287            unattributed_bytes,
288            root_count: roots.len(),
289            busy_subsystems,
290            not_estimated_subsystems,
291        }
292    }
293}
294
295#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
296pub struct MemorySnapshot {
297    pub roots_status: &'static str,
298    pub roots: BTreeMap<String, RootMemorySnapshot>,
299    /// Immutable borrowed semantic snapshots, attributed once process-wide.
300    pub shared_semantic_bases: MemoryEstimate,
301    pub process: ProcessMemorySnapshot,
302}
303
304impl MemorySnapshot {
305    pub fn new(roots_status: &'static str, roots: BTreeMap<String, RootMemorySnapshot>) -> Self {
306        let shared_semantic_bases = crate::semantic_index::shared_semantic_bases_memory();
307        let process = ProcessMemorySnapshot::from_roots(&roots, &shared_semantic_bases);
308        Self {
309            roots_status,
310            roots,
311            shared_semantic_bases,
312            process,
313        }
314    }
315}
316
317pub fn path_bytes(path: &Path) -> u64 {
318    #[cfg(unix)]
319    {
320        use std::os::unix::ffi::OsStrExt;
321        usize_to_u64(path.as_os_str().as_bytes().len())
322    }
323    #[cfg(windows)]
324    {
325        use std::os::windows::ffi::OsStrExt;
326        usize_to_u64(path.as_os_str().encode_wide().count())
327            .saturating_mul(std::mem::size_of::<u16>() as u64)
328    }
329    #[cfg(not(any(unix, windows)))]
330    {
331        usize_to_u64(path.to_string_lossy().len())
332    }
333}
334
335pub fn usize_to_u64(value: usize) -> u64 {
336    u64::try_from(value).unwrap_or(u64::MAX)
337}
338
339pub fn estimated_json_bytes(value: &Value) -> u64 {
340    match value {
341        Value::Null => 0,
342        Value::Bool(_) => std::mem::size_of::<bool>() as u64,
343        Value::Number(_) => std::mem::size_of::<serde_json::Number>() as u64,
344        Value::String(value) => usize_to_u64(value.len()),
345        Value::Array(values) => values
346            .iter()
347            .map(estimated_json_bytes)
348            .fold(0u64, u64::saturating_add),
349        Value::Object(values) => values.iter().fold(0u64, |bytes, (key, value)| {
350            bytes
351                .saturating_add(usize_to_u64(key.len()))
352                .saturating_add(estimated_json_bytes(value))
353        }),
354    }
355}
356
357fn signed_difference(lhs: u64, rhs: u64) -> i64 {
358    let difference = i128::from(lhs) - i128::from(rhs);
359    difference.clamp(i128::from(i64::MIN), i128::from(i64::MAX)) as i64
360}
361
362fn nonnegative_i64_to_u64(value: i64) -> u64 {
363    u64::try_from(value).unwrap_or(0)
364}
365
366#[cfg(target_os = "macos")]
367fn allocator_memory_snapshot() -> AllocatorMemorySnapshot {
368    let mut statistics = std::mem::MaybeUninit::<libc::malloc_statistics_t>::zeroed();
369    unsafe {
370        libc::malloc_zone_statistics(libc::malloc_default_zone(), statistics.as_mut_ptr());
371    }
372    let statistics = unsafe { statistics.assume_init() };
373    AllocatorMemorySnapshot::measured(
374        usize_to_u64(statistics.size_in_use),
375        usize_to_u64(statistics.size_allocated),
376    )
377}
378
379#[cfg(all(target_os = "linux", target_env = "gnu"))]
380fn allocator_memory_snapshot() -> AllocatorMemorySnapshot {
381    // mallinfo2 exists only in glibc >= 2.33. Release Linux binaries link
382    // against an older glibc floor (cross gnu images, kept old so dlopen and
383    // wide distro compatibility hold), so a link-time reference to the symbol
384    // fails the release build even though native CI (glibc 2.35) links fine.
385    // Resolve it at runtime instead and report honestly when it is absent.
386    use std::sync::OnceLock;
387    type Mallinfo2Fn = unsafe extern "C" fn() -> libc::mallinfo2;
388    static MALLINFO2: OnceLock<Option<Mallinfo2Fn>> = OnceLock::new();
389    let resolved = MALLINFO2.get_or_init(|| {
390        let symbol = unsafe { libc::dlsym(libc::RTLD_DEFAULT, c"mallinfo2".as_ptr()) };
391        if symbol.is_null() {
392            None
393        } else {
394            // SAFETY: glibc declares mallinfo2 as `struct mallinfo2 (*)(void)`;
395            // the signature matches Mallinfo2Fn exactly.
396            Some(unsafe { std::mem::transmute::<*mut libc::c_void, Mallinfo2Fn>(symbol) })
397        }
398    });
399    let Some(mallinfo2) = resolved else {
400        return AllocatorMemorySnapshot::not_estimated("mallinfo2_requires_glibc_2_33");
401    };
402    let statistics = unsafe { mallinfo2() };
403    let mapped_bytes = statistics.hblkhd as u64;
404    let bytes_in_use = (statistics.uordblks as u64).saturating_add(mapped_bytes);
405    let size_allocated = (statistics.arena as u64).saturating_add(mapped_bytes);
406    AllocatorMemorySnapshot::measured(bytes_in_use, size_allocated)
407}
408
409#[cfg(not(any(target_os = "macos", all(target_os = "linux", target_env = "gnu"))))]
410fn allocator_memory_snapshot() -> AllocatorMemorySnapshot {
411    AllocatorMemorySnapshot::not_estimated("platform_allocator_statistics_unavailable")
412}
413
414#[cfg(target_os = "macos")]
415unsafe extern "C" {
416    fn malloc_zone_pressure_relief(zone: *mut libc::malloc_zone_t, goal: usize) -> usize;
417}
418
419/// Ask the macOS allocator to return unused pages after a process-wide idle gate.
420/// Callers own that gate because allocator pressure relief can add latency.
421#[cfg(target_os = "macos")]
422pub fn relieve_allocator_pressure() -> AllocatorPressureRelief {
423    let rss_before_bytes = process_rss_bytes();
424    let allocator_before = allocator_memory_snapshot();
425    let bytes_released = unsafe { malloc_zone_pressure_relief(std::ptr::null_mut(), 0) };
426    let allocator_after = allocator_memory_snapshot();
427    let rss_after_bytes = process_rss_bytes();
428    AllocatorPressureRelief {
429        bytes_released: usize_to_u64(bytes_released),
430        rss_before_bytes,
431        rss_after_bytes,
432        allocator_before,
433        allocator_after,
434    }
435}
436
437#[cfg(target_os = "macos")]
438fn process_rss_bytes() -> Option<u64> {
439    let mut info = std::mem::MaybeUninit::<libc::proc_taskinfo>::zeroed();
440    let size = std::mem::size_of::<libc::proc_taskinfo>();
441    let written = unsafe {
442        libc::proc_pidinfo(
443            libc::getpid(),
444            libc::PROC_PIDTASKINFO,
445            0,
446            info.as_mut_ptr().cast(),
447            i32::try_from(size).ok()?,
448        )
449    };
450    if written != i32::try_from(size).ok()? {
451        return None;
452    }
453    Some(unsafe { info.assume_init() }.pti_resident_size)
454}
455
456#[cfg(target_os = "linux")]
457fn process_rss_bytes() -> Option<u64> {
458    let statm = std::fs::read_to_string("/proc/self/statm").ok()?;
459    let resident_pages = statm.split_whitespace().nth(1)?.parse::<u64>().ok()?;
460    let page_size = unsafe { libc::sysconf(libc::_SC_PAGESIZE) };
461    if page_size <= 0 {
462        return None;
463    }
464    resident_pages.checked_mul(page_size as u64)
465}
466
467#[cfg(not(any(target_os = "macos", target_os = "linux")))]
468fn process_rss_bytes() -> Option<u64> {
469    None
470}
471
472#[cfg(test)]
473mod tests {
474    use super::*;
475
476    #[test]
477    fn process_snapshot_preserves_negative_residuals() {
478        assert_eq!(signed_difference(5, 8), -3);
479    }
480
481    #[test]
482    fn json_estimator_scales_with_payload_content() {
483        let empty = estimated_json_bytes(&serde_json::json!({}));
484        let populated = estimated_json_bytes(&serde_json::json!({"message": "hello"}));
485        assert_eq!(empty, 0);
486        assert!(populated >= 12);
487    }
488
489    #[test]
490    fn process_snapshot_exposes_sqlite_and_allocator_sections() {
491        let shared = MemoryEstimate::estimated(7);
492        let snapshot = ProcessMemorySnapshot::from_roots(&BTreeMap::new(), &shared);
493        assert_eq!(snapshot.sqlite.status, "measured");
494        assert!(snapshot.sqlite.memory_highwater_bytes >= snapshot.sqlite.memory_used_bytes);
495        assert_eq!(
496            snapshot.total_attributed_bytes,
497            snapshot.sqlite.memory_used_bytes.saturating_add(7)
498        );
499
500        let serialized = serde_json::to_value(&snapshot).expect("serialize process memory");
501        assert!(serialized["sqlite"]["memory_used_bytes"].is_u64());
502        assert!(serialized["allocator"].get("bytes_in_use").is_some());
503        assert!(serialized["allocator"].get("size_allocated").is_some());
504        assert!(serialized["allocator"]
505            .get("retained_slack_bytes")
506            .is_some());
507    }
508
509    #[cfg(any(target_os = "macos", all(target_os = "linux", target_env = "gnu")))]
510    #[test]
511    fn allocator_snapshot_reports_measured_slack() {
512        let allocator = allocator_memory_snapshot();
513        assert_eq!(allocator.status, "measured");
514        let in_use = allocator.bytes_in_use.expect("allocator bytes in use");
515        let allocated = allocator.size_allocated.expect("allocator size allocated");
516        assert_eq!(
517            allocator.retained_slack_bytes,
518            Some(allocated.saturating_sub(in_use))
519        );
520    }
521
522    #[cfg(not(any(target_os = "macos", all(target_os = "linux", target_env = "gnu"))))]
523    #[test]
524    fn allocator_snapshot_is_honest_when_platform_counters_are_unavailable() {
525        let allocator = allocator_memory_snapshot();
526        assert_eq!(allocator.status, "not_estimated_on_this_platform");
527        assert_eq!(allocator.bytes_in_use, None);
528        assert_eq!(allocator.size_allocated, None);
529        assert_eq!(allocator.retained_slack_bytes, None);
530        assert_eq!(
531            allocator.not_estimated,
532            Some("platform_allocator_statistics_unavailable")
533        );
534    }
535
536    #[cfg(target_os = "macos")]
537    #[test]
538    #[ignore = "bounded live RSS experiment; run explicitly after allocator changes"]
539    fn allocator_pressure_relief_warm_then_idle_measurement() {
540        let warm_pages = (0..16 * 1024)
541            .map(|seed| {
542                let mut page = Box::new([0u8; 4096]);
543                page[0] = seed as u8;
544                page
545            })
546            .collect::<Vec<_>>();
547        std::hint::black_box(&warm_pages);
548        drop(warm_pages);
549
550        let relief = relieve_allocator_pressure();
551        let sqlite = SqliteMemorySnapshot::measure();
552        eprintln!(
553            "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={}",
554            relief.rss_before_bytes,
555            relief.rss_after_bytes,
556            relief.allocator_before.bytes_in_use,
557            relief.allocator_after.bytes_in_use,
558            relief.allocator_before.size_allocated,
559            relief.allocator_after.size_allocated,
560            relief.allocator_before.retained_slack_bytes,
561            relief.allocator_after.retained_slack_bytes,
562            relief.bytes_released,
563            sqlite.memory_used_bytes,
564            sqlite.memory_highwater_bytes,
565        );
566        assert_eq!(relief.allocator_before.status, "measured");
567        assert_eq!(relief.allocator_after.status, "measured");
568    }
569}
aft/memory.rs

aft/
memory.rs
