1use std::collections::BTreeMap;
2use std::path::Path;
3
4use serde::Serialize;
5use serde_json::Value;
6
7#[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 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 #[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 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
295pub const MEMORY_SNAPSHOT_ROOT_DETAIL_CAP: usize = 8;
302
303#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
304pub struct MemorySnapshot {
305 pub roots_status: &'static str,
306 pub roots: BTreeMap<String, RootMemorySnapshot>,
310 pub roots_total: usize,
312 pub roots_omitted: usize,
314 pub roots_omitted_bytes: u64,
317 pub shared_semantic_bases: MemoryEstimate,
319 pub process: ProcessMemorySnapshot,
320}
321
322impl MemorySnapshot {
323 pub fn new(roots_status: &'static str, roots: BTreeMap<String, RootMemorySnapshot>) -> Self {
324 let shared_semantic_bases = crate::semantic_index::shared_semantic_bases_memory();
325 let process = ProcessMemorySnapshot::from_roots(&roots, &shared_semantic_bases);
327 let roots_total = roots.len();
328 let (roots, roots_omitted, roots_omitted_bytes) =
329 cap_root_detail(roots, MEMORY_SNAPSHOT_ROOT_DETAIL_CAP);
330 Self {
331 roots_status,
332 roots,
333 roots_total,
334 roots_omitted,
335 roots_omitted_bytes,
336 shared_semantic_bases,
337 process,
338 }
339 }
340}
341
342fn cap_root_detail(
345 roots: BTreeMap<String, RootMemorySnapshot>,
346 cap: usize,
347) -> (BTreeMap<String, RootMemorySnapshot>, usize, u64) {
348 if roots.len() <= cap {
349 return (roots, 0, 0);
350 }
351 let mut entries: Vec<(String, RootMemorySnapshot)> = roots.into_iter().collect();
352 entries.sort_by(|a, b| {
354 b.1.attributed_bytes
355 .cmp(&a.1.attributed_bytes)
356 .then_with(|| a.0.cmp(&b.0))
357 });
358 let omitted: Vec<(String, RootMemorySnapshot)> = entries.split_off(cap);
359 let omitted_bytes = omitted
360 .iter()
361 .map(|(_, snapshot)| snapshot.attributed_bytes)
362 .fold(0u64, u64::saturating_add);
363 (entries.into_iter().collect(), omitted.len(), omitted_bytes)
364}
365
366#[cfg(test)]
367mod snapshot_cap_tests {
368 use super::*;
369
370 fn root_with_bytes(bytes: u64) -> RootMemorySnapshot {
371 let estimate = MemoryEstimate::estimated;
372 RootMemorySnapshot::new(
373 estimate(bytes),
374 estimate(0),
375 estimate(0),
376 estimate(0),
377 estimate(0),
378 estimate(0),
379 estimate(0),
380 estimate(0),
381 )
382 }
383
384 #[test]
385 fn detail_map_capped_but_totals_cover_all_roots() {
386 let mut roots = BTreeMap::new();
387 for i in 0..(MEMORY_SNAPSHOT_ROOT_DETAIL_CAP + 4) {
388 roots.insert(
390 format!("/root/{i:02}"),
391 root_with_bytes((i as u64 + 1) * 1000),
392 );
393 }
394 let snapshot = MemorySnapshot::new("ready", roots);
395
396 assert_eq!(snapshot.roots.len(), MEMORY_SNAPSHOT_ROOT_DETAIL_CAP);
397 assert_eq!(snapshot.roots_total, MEMORY_SNAPSHOT_ROOT_DETAIL_CAP + 4);
398 assert_eq!(snapshot.roots_omitted, 4);
399 assert_eq!(snapshot.roots_omitted_bytes, 1000 + 2000 + 3000 + 4000);
401 assert!(snapshot
403 .roots
404 .values()
405 .all(|root| root.attributed_bytes > 4000));
406 let expected_total: u64 = (1..=(MEMORY_SNAPSHOT_ROOT_DETAIL_CAP as u64 + 4))
408 .map(|i| i * 1000)
409 .sum();
410 assert!(snapshot.process.total_attributed_bytes >= expected_total);
411 }
412
413 #[test]
414 fn under_cap_keeps_everything_with_zero_omitted() {
415 let mut roots = BTreeMap::new();
416 roots.insert("/a".to_string(), root_with_bytes(10));
417 roots.insert("/b".to_string(), root_with_bytes(20));
418 let snapshot = MemorySnapshot::new("ready", roots);
419 assert_eq!(snapshot.roots.len(), 2);
420 assert_eq!(snapshot.roots_total, 2);
421 assert_eq!(snapshot.roots_omitted, 0);
422 assert_eq!(snapshot.roots_omitted_bytes, 0);
423 }
424}
425
426pub fn path_bytes(path: &Path) -> u64 {
427 #[cfg(unix)]
428 {
429 use std::os::unix::ffi::OsStrExt;
430 usize_to_u64(path.as_os_str().as_bytes().len())
431 }
432 #[cfg(windows)]
433 {
434 use std::os::windows::ffi::OsStrExt;
435 usize_to_u64(path.as_os_str().encode_wide().count())
436 .saturating_mul(std::mem::size_of::<u16>() as u64)
437 }
438 #[cfg(not(any(unix, windows)))]
439 {
440 usize_to_u64(path.to_string_lossy().len())
441 }
442}
443
444pub fn usize_to_u64(value: usize) -> u64 {
445 u64::try_from(value).unwrap_or(u64::MAX)
446}
447
448pub fn estimated_json_bytes(value: &Value) -> u64 {
449 match value {
450 Value::Null => 0,
451 Value::Bool(_) => std::mem::size_of::<bool>() as u64,
452 Value::Number(_) => std::mem::size_of::<serde_json::Number>() as u64,
453 Value::String(value) => usize_to_u64(value.len()),
454 Value::Array(values) => values
455 .iter()
456 .map(estimated_json_bytes)
457 .fold(0u64, u64::saturating_add),
458 Value::Object(values) => values.iter().fold(0u64, |bytes, (key, value)| {
459 bytes
460 .saturating_add(usize_to_u64(key.len()))
461 .saturating_add(estimated_json_bytes(value))
462 }),
463 }
464}
465
466fn signed_difference(lhs: u64, rhs: u64) -> i64 {
467 let difference = i128::from(lhs) - i128::from(rhs);
468 difference.clamp(i128::from(i64::MIN), i128::from(i64::MAX)) as i64
469}
470
471fn nonnegative_i64_to_u64(value: i64) -> u64 {
472 u64::try_from(value).unwrap_or(0)
473}
474
475#[cfg(target_os = "macos")]
476fn allocator_memory_snapshot() -> AllocatorMemorySnapshot {
477 let mut statistics = std::mem::MaybeUninit::<libc::malloc_statistics_t>::zeroed();
478 unsafe {
479 libc::malloc_zone_statistics(libc::malloc_default_zone(), statistics.as_mut_ptr());
480 }
481 let statistics = unsafe { statistics.assume_init() };
482 AllocatorMemorySnapshot::measured(
483 usize_to_u64(statistics.size_in_use),
484 usize_to_u64(statistics.size_allocated),
485 )
486}
487
488#[cfg(all(target_os = "linux", target_env = "gnu"))]
489fn allocator_memory_snapshot() -> AllocatorMemorySnapshot {
490 use std::sync::OnceLock;
496 type Mallinfo2Fn = unsafe extern "C" fn() -> libc::mallinfo2;
497 static MALLINFO2: OnceLock<Option<Mallinfo2Fn>> = OnceLock::new();
498 let resolved = MALLINFO2.get_or_init(|| {
499 let symbol = unsafe { libc::dlsym(libc::RTLD_DEFAULT, c"mallinfo2".as_ptr()) };
500 if symbol.is_null() {
501 None
502 } else {
503 Some(unsafe { std::mem::transmute::<*mut libc::c_void, Mallinfo2Fn>(symbol) })
506 }
507 });
508 let Some(mallinfo2) = resolved else {
509 return AllocatorMemorySnapshot::not_estimated("mallinfo2_requires_glibc_2_33");
510 };
511 let statistics = unsafe { mallinfo2() };
512 let mapped_bytes = statistics.hblkhd as u64;
513 let bytes_in_use = (statistics.uordblks as u64).saturating_add(mapped_bytes);
514 let size_allocated = (statistics.arena as u64).saturating_add(mapped_bytes);
515 AllocatorMemorySnapshot::measured(bytes_in_use, size_allocated)
516}
517
518#[cfg(not(any(target_os = "macos", all(target_os = "linux", target_env = "gnu"))))]
519fn allocator_memory_snapshot() -> AllocatorMemorySnapshot {
520 AllocatorMemorySnapshot::not_estimated("platform_allocator_statistics_unavailable")
521}
522
523#[cfg(target_os = "macos")]
524unsafe extern "C" {
525 fn malloc_zone_pressure_relief(zone: *mut libc::malloc_zone_t, goal: usize) -> usize;
526}
527
528#[cfg(target_os = "macos")]
531pub fn relieve_allocator_pressure() -> AllocatorPressureRelief {
532 let rss_before_bytes = process_rss_bytes();
533 let allocator_before = allocator_memory_snapshot();
534 let bytes_released = unsafe { malloc_zone_pressure_relief(std::ptr::null_mut(), 0) };
535 let allocator_after = allocator_memory_snapshot();
536 let rss_after_bytes = process_rss_bytes();
537 AllocatorPressureRelief {
538 bytes_released: usize_to_u64(bytes_released),
539 rss_before_bytes,
540 rss_after_bytes,
541 allocator_before,
542 allocator_after,
543 }
544}
545
546#[cfg(target_os = "macos")]
547fn process_rss_bytes() -> Option<u64> {
548 let mut info = std::mem::MaybeUninit::<libc::proc_taskinfo>::zeroed();
549 let size = std::mem::size_of::<libc::proc_taskinfo>();
550 let written = unsafe {
551 libc::proc_pidinfo(
552 libc::getpid(),
553 libc::PROC_PIDTASKINFO,
554 0,
555 info.as_mut_ptr().cast(),
556 i32::try_from(size).ok()?,
557 )
558 };
559 if written != i32::try_from(size).ok()? {
560 return None;
561 }
562 Some(unsafe { info.assume_init() }.pti_resident_size)
563}
564
565#[cfg(target_os = "linux")]
566fn process_rss_bytes() -> Option<u64> {
567 let statm = std::fs::read_to_string("/proc/self/statm").ok()?;
568 let resident_pages = statm.split_whitespace().nth(1)?.parse::<u64>().ok()?;
569 let page_size = unsafe { libc::sysconf(libc::_SC_PAGESIZE) };
570 if page_size <= 0 {
571 return None;
572 }
573 resident_pages.checked_mul(page_size as u64)
574}
575
576#[cfg(not(any(target_os = "macos", target_os = "linux")))]
577fn process_rss_bytes() -> Option<u64> {
578 None
579}
580
581#[cfg(test)]
582mod tests {
583 use super::*;
584
585 #[test]
586 fn process_snapshot_preserves_negative_residuals() {
587 assert_eq!(signed_difference(5, 8), -3);
588 }
589
590 #[test]
591 fn json_estimator_scales_with_payload_content() {
592 let empty = estimated_json_bytes(&serde_json::json!({}));
593 let populated = estimated_json_bytes(&serde_json::json!({"message": "hello"}));
594 assert_eq!(empty, 0);
595 assert!(populated >= 12);
596 }
597
598 #[test]
599 fn process_snapshot_exposes_sqlite_and_allocator_sections() {
600 let shared = MemoryEstimate::estimated(7);
601 let snapshot = ProcessMemorySnapshot::from_roots(&BTreeMap::new(), &shared);
602 assert_eq!(snapshot.sqlite.status, "measured");
603 assert!(snapshot.sqlite.memory_highwater_bytes >= snapshot.sqlite.memory_used_bytes);
604 assert_eq!(
605 snapshot.total_attributed_bytes,
606 snapshot.sqlite.memory_used_bytes.saturating_add(7)
607 );
608
609 let serialized = serde_json::to_value(&snapshot).expect("serialize process memory");
610 assert!(serialized["sqlite"]["memory_used_bytes"].is_u64());
611 assert!(serialized["allocator"].get("bytes_in_use").is_some());
612 assert!(serialized["allocator"].get("size_allocated").is_some());
613 assert!(serialized["allocator"]
614 .get("retained_slack_bytes")
615 .is_some());
616 }
617
618 #[cfg(any(target_os = "macos", all(target_os = "linux", target_env = "gnu")))]
619 #[test]
620 fn allocator_snapshot_reports_measured_slack() {
621 let allocator = allocator_memory_snapshot();
622 assert_eq!(allocator.status, "measured");
623 let in_use = allocator.bytes_in_use.expect("allocator bytes in use");
624 let allocated = allocator.size_allocated.expect("allocator size allocated");
625 assert_eq!(
626 allocator.retained_slack_bytes,
627 Some(allocated.saturating_sub(in_use))
628 );
629 }
630
631 #[cfg(not(any(target_os = "macos", all(target_os = "linux", target_env = "gnu"))))]
632 #[test]
633 fn allocator_snapshot_is_honest_when_platform_counters_are_unavailable() {
634 let allocator = allocator_memory_snapshot();
635 assert_eq!(allocator.status, "not_estimated_on_this_platform");
636 assert_eq!(allocator.bytes_in_use, None);
637 assert_eq!(allocator.size_allocated, None);
638 assert_eq!(allocator.retained_slack_bytes, None);
639 assert_eq!(
640 allocator.not_estimated,
641 Some("platform_allocator_statistics_unavailable")
642 );
643 }
644
645 #[cfg(target_os = "macos")]
646 #[test]
647 #[ignore = "bounded live RSS experiment; run explicitly after allocator changes"]
648 fn allocator_pressure_relief_warm_then_idle_measurement() {
649 let warm_pages = (0..16 * 1024)
650 .map(|seed| {
651 let mut page = Box::new([0u8; 4096]);
652 page[0] = seed as u8;
653 page
654 })
655 .collect::<Vec<_>>();
656 std::hint::black_box(&warm_pages);
657 drop(warm_pages);
658
659 let relief = relieve_allocator_pressure();
660 let sqlite = SqliteMemorySnapshot::measure();
661 eprintln!(
662 "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={}",
663 relief.rss_before_bytes,
664 relief.rss_after_bytes,
665 relief.allocator_before.bytes_in_use,
666 relief.allocator_after.bytes_in_use,
667 relief.allocator_before.size_allocated,
668 relief.allocator_after.size_allocated,
669 relief.allocator_before.retained_slack_bytes,
670 relief.allocator_after.retained_slack_bytes,
671 relief.bytes_released,
672 sqlite.memory_used_bytes,
673 sqlite.memory_highwater_bytes,
674 );
675 assert_eq!(relief.allocator_before.status, "measured");
676 assert_eq!(relief.allocator_after.status, "measured");
677 }
678}