1use std::path::{Path, PathBuf};
36use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
37use std::sync::{Arc, OnceLock, RwLock};
38use std::time::{Duration, Instant};
39
40use anyhow::Result;
41
42pub const DEFAULT_BUDGET_PCT: u64 = 20;
44pub const POLL: Duration = Duration::from_millis(250);
46pub const CHECK_EVERY_OBJECTS: u64 = 4096;
49pub const REFLUSH_EVERY: Duration = Duration::from_secs(10);
55pub const RESUME_MARGIN_BYTES: u64 = 64 << 20;
69
70pub const PAUSE_DEADLINE: Duration = Duration::from_secs(300);
75
76pub const GIVE_UP_AFTER_EMPTY_FLUSHES: u32 = 3;
82
83pub fn parse_meminfo(text: &str) -> Option<(u64, u64)> {
85 let mut total = None;
86 let mut avail = None;
87 for line in text.lines() {
88 let Some((key, rest)) = line.split_once(':') else { continue };
89 let Some(kib) = rest.split_ascii_whitespace().next().and_then(|v| v.parse::<u64>().ok())
90 else {
91 continue;
92 };
93 match key {
94 "MemTotal" => total = Some(kib * 1024),
95 "MemAvailable" => avail = Some(kib * 1024),
96 _ => {}
97 }
98 if total.is_some() && avail.is_some() {
99 break;
100 }
101 }
102 Some((total?, avail?))
103}
104
105pub fn fmt_bytes(n: u64) -> String {
108 const K: u64 = 1024;
109 if n >= K * K * K {
110 format!("{:.1}G", n as f64 / (K * K * K) as f64)
111 } else if n >= K * K {
112 format!("{}M", n / (K * K))
113 } else if n >= K {
114 format!("{}K", n / K)
115 } else {
116 format!("{n}B")
117 }
118}
119
120pub struct MemWatchdog {
122 meminfo: PathBuf,
123 total: u64,
124 budget: u64,
125 poll: Duration,
126 pauses: AtomicU64,
128 paused_ms: AtomicU64,
130 paused_now: AtomicBool,
132 resume_margin: u64,
134 deadline: Duration,
136 reflush: Duration,
138 give_ups: AtomicU64,
140}
141
142#[derive(Debug, Clone, Copy, PartialEq, Eq)]
148pub enum PauseEnd {
149 Recovered,
151 Cancelled,
153 Deadline,
155 NothingToGive,
158}
159
160impl PauseEnd {
161 pub fn as_str(self) -> &'static str {
162 match self {
163 PauseEnd::Recovered => "recovered",
164 PauseEnd::Cancelled => "cancelled",
165 PauseEnd::Deadline => "deadline",
166 PauseEnd::NothingToGive => "nothing-to-give",
167 }
168 }
169
170 pub fn gave_up(self) -> bool {
173 matches!(self, PauseEnd::Deadline | PauseEnd::NothingToGive)
174 }
175}
176
177#[derive(Debug, Clone, Copy, PartialEq, Eq)]
179pub struct Pause {
180 pub ms: u64,
181 pub end: PauseEnd,
182 pub available: u64,
183 pub budget: u64,
184 pub empty_flushes: u32,
186}
187
188impl MemWatchdog {
189 pub fn from_env() -> Self {
191 let meminfo = crate::arms::read_env(crate::arms::ENV_MEMINFO)
192 .map(PathBuf::from)
193 .unwrap_or_else(|| PathBuf::from("/proc/meminfo"));
194 let total = std::fs::read_to_string(&meminfo)
195 .ok()
196 .and_then(|t| parse_meminfo(&t))
197 .map(|(t, _)| t)
198 .unwrap_or(0);
199 let budget = match crate::arms::read_env(crate::arms::ENV_MEM_BUDGET_BYTES)
200 .and_then(|v| v.parse::<u64>().ok())
201 {
202 Some(b) => b,
203 None => {
204 let pct = crate::arms::read_env(crate::arms::ENV_MEM_BUDGET_PCT)
205 .and_then(|v| v.parse::<u64>().ok())
206 .unwrap_or(DEFAULT_BUDGET_PCT)
207 .min(100);
208 total / 100 * pct
209 }
210 };
211 let margin = crate::arms::read_env(crate::arms::ENV_MEM_RESUME_MARGIN_BYTES)
212 .and_then(|v| v.parse::<u64>().ok())
213 .unwrap_or(RESUME_MARGIN_BYTES);
214 let deadline = crate::arms::read_env(crate::arms::ENV_MEM_PAUSE_MAX_MS)
215 .and_then(|v| v.parse::<u64>().ok())
216 .map(Duration::from_millis)
217 .unwrap_or(PAUSE_DEADLINE);
218 Self::with(&meminfo, budget, POLL).bounded(margin, deadline)
219 }
220
221 pub fn with(meminfo: &Path, budget: u64, poll: Duration) -> Self {
224 let total = std::fs::read_to_string(meminfo)
225 .ok()
226 .and_then(|t| parse_meminfo(&t))
227 .map(|(t, _)| t)
228 .unwrap_or(0);
229 Self {
230 meminfo: meminfo.to_path_buf(),
231 total,
232 budget,
233 poll,
234 pauses: AtomicU64::new(0),
235 paused_ms: AtomicU64::new(0),
236 paused_now: AtomicBool::new(false),
237 resume_margin: RESUME_MARGIN_BYTES,
238 deadline: PAUSE_DEADLINE,
239 reflush: REFLUSH_EVERY,
240 give_ups: AtomicU64::new(0),
241 }
242 }
243
244 pub fn bounded(mut self, resume_margin: u64, deadline: Duration) -> Self {
248 self.resume_margin = resume_margin;
249 self.deadline = deadline;
250 self
251 }
252
253 pub fn with_reflush(mut self, every: Duration) -> Self {
257 self.reflush = every;
258 self
259 }
260
261 pub fn disabled() -> Self {
263 Self::with(Path::new("/nonexistent"), 0, POLL)
264 }
265
266 pub fn budget(&self) -> u64 {
267 self.budget
268 }
269
270 pub fn total(&self) -> u64 {
271 self.total
272 }
273
274 pub fn available(&self) -> Option<u64> {
277 std::fs::read_to_string(&self.meminfo)
278 .ok()
279 .and_then(|t| parse_meminfo(&t))
280 .map(|(_, a)| a)
281 }
282
283 pub fn breached(&self) -> Option<(u64, u64)> {
285 if self.budget == 0 {
286 return None;
287 }
288 let avail = self.available()?;
289 (avail < self.budget).then_some((avail, self.budget))
290 }
291
292 pub fn resume_margin(&self) -> u64 {
297 self.resume_margin.min(self.budget.max(1))
298 }
299
300 pub fn pause_deadline(&self) -> Duration {
304 self.deadline
305 }
306
307 fn recovered(&self) -> bool {
309 match self.available() {
310 None => true,
311 Some(avail) => avail >= self.budget.saturating_add(self.resume_margin()),
312 }
313 }
314
315 pub fn pauses(&self) -> u64 {
317 self.pauses.load(Ordering::Relaxed)
318 }
319
320 pub fn paused_ms(&self) -> u64 {
322 self.paused_ms.load(Ordering::Relaxed)
323 }
324
325 pub fn is_paused(&self) -> bool {
327 self.paused_now.load(Ordering::Relaxed)
328 }
329
330 pub fn give_ups(&self) -> u64 {
332 self.give_ups.load(Ordering::Relaxed)
333 }
334
335 pub fn pause_until_ok(
356 &self,
357 label: &str,
358 cancel: Option<&AtomicBool>,
359 flush: &mut dyn FnMut() -> Result<u64>,
360 ) -> Result<Pause> {
361 let Some((avail, budget)) = self.breached() else {
362 return Ok(Pause {
363 ms: 0,
364 end: PauseEnd::Recovered,
365 available: self.available().unwrap_or(0),
366 budget: self.budget,
367 empty_flushes: 0,
368 });
369 };
370 self.pauses.fetch_add(1, Ordering::Relaxed);
371 self.paused_now.store(true, Ordering::Relaxed);
372 let margin = self.resume_margin();
373 tracing::warn!(
374 target: "gunnar::index",
375 event = "index.paused",
376 repo = label,
377 available = avail,
378 budget = budget,
379 resume_at = budget.saturating_add(margin),
380 deadline_ms = self.deadline.as_millis() as u64,
381 "index paused: {} available < {} budget; resumes at {}, gives up after {} s",
382 fmt_bytes(avail),
383 fmt_bytes(budget),
384 fmt_bytes(budget.saturating_add(margin)),
385 self.deadline.as_secs()
386 );
387 let started = Instant::now();
388 let mut empty_flushes = 0u32;
393 let mut flush_err: Option<anyhow::Error> = None;
394 match flush() {
395 Ok(0) => empty_flushes = 1,
396 Ok(_) => {}
397 Err(e) => flush_err = Some(e),
398 }
399 let mut last_flush = Instant::now();
400 let end = loop {
401 if cancel.is_some_and(|c| c.load(Ordering::Relaxed)) {
402 break PauseEnd::Cancelled;
403 }
404 if self.recovered() {
405 break PauseEnd::Recovered;
406 }
407 if !self.deadline.is_zero() && started.elapsed() >= self.deadline {
408 break PauseEnd::Deadline;
409 }
410 if empty_flushes >= GIVE_UP_AFTER_EMPTY_FLUSHES
411 && self.available().is_some_and(|a| a <= avail)
412 {
413 break PauseEnd::NothingToGive;
414 }
415 if last_flush.elapsed() >= self.reflush {
416 match flush() {
417 Ok(0) => empty_flushes += 1,
418 Ok(_) => empty_flushes = 0,
419 Err(e) => {
420 if flush_err.is_none() {
421 flush_err = Some(e);
422 }
423 }
424 }
425 last_flush = Instant::now();
426 }
427 std::thread::sleep(self.poll);
428 };
429 let ms = started.elapsed().as_millis() as u64;
430 self.paused_ms.fetch_add(ms, Ordering::Relaxed);
431 self.paused_now.store(false, Ordering::Relaxed);
432 if end.gave_up() {
433 self.give_ups.fetch_add(1, Ordering::Relaxed);
434 }
435 let now = self.available().unwrap_or(0);
436 if end.gave_up() {
437 tracing::warn!(
438 target: "gunnar::index",
439 event = "index.gave_up",
440 repo = label,
441 available = now,
442 budget = budget,
443 paused_ms = ms,
444 empty_flushes = empty_flushes,
445 reason = end.as_str(),
446 "index gave up waiting for memory after {ms} ms ({}): {} available, {} needed. \
447 The caller decides — proceed, shed, or refuse the pack by name — but nobody \
448 waits here again.",
449 end.as_str(),
450 fmt_bytes(now),
451 fmt_bytes(budget.saturating_add(margin))
452 );
453 } else {
454 tracing::info!(
455 target: "gunnar::index",
456 event = "index.resumed",
457 repo = label,
458 available = now,
459 budget = budget,
460 paused_ms = ms,
461 reason = end.as_str(),
462 "index resumed after {ms} ms ({}): {} available",
463 end.as_str(),
464 fmt_bytes(now)
465 );
466 }
467 if let Some(e) = flush_err {
468 return Err(e);
469 }
470 Ok(Pause { ms, end, available: now, budget, empty_flushes })
471 }
472}
473
474fn global_slot() -> &'static RwLock<Arc<MemWatchdog>> {
475 static G: OnceLock<RwLock<Arc<MemWatchdog>>> = OnceLock::new();
476 G.get_or_init(|| RwLock::new(Arc::new(MemWatchdog::from_env())))
477}
478
479pub fn global() -> Arc<MemWatchdog> {
484 global_slot().read().unwrap_or_else(|e| e.into_inner()).clone()
485}
486
487pub fn set_global(dog: Arc<MemWatchdog>) -> Arc<MemWatchdog> {
491 std::mem::replace(&mut *global_slot().write().unwrap_or_else(|e| e.into_inner()), dog)
492}
493
494pub fn peak_rss_bytes() -> u64 {
498 std::fs::read_to_string("/proc/self/status")
499 .ok()
500 .and_then(|t| {
501 t.lines().find_map(|l| {
502 l.strip_prefix("VmHWM:")
503 .and_then(|r| r.split_ascii_whitespace().next())
504 .and_then(|v| v.parse::<u64>().ok())
505 .map(|kib| kib * 1024)
506 })
507 })
508 .unwrap_or(0)
509}
510
511pub fn mem_total() -> u64 {
514 global().total()
515}
516
517#[cfg(test)]
518mod tests {
519 use super::*;
520 use std::sync::Arc;
521 use std::sync::atomic::AtomicU64;
522
523 fn fake_meminfo(dir: &Path, total_kib: u64, avail_kib: u64) -> PathBuf {
524 let p = dir.join("meminfo");
525 std::fs::write(
526 &p,
527 format!(
528 "MemTotal: {total_kib} kB\nMemFree: {avail_kib} kB\nMemAvailable: {avail_kib} kB\nBuffers: 0 kB\n"
529 ),
530 )
531 .unwrap();
532 p
533 }
534
535 fn scratch(name: &str) -> PathBuf {
536 let d = std::env::temp_dir().join(format!("memguard-{}-{name}", std::process::id()));
537 let _ = std::fs::remove_dir_all(&d);
538 std::fs::create_dir_all(&d).unwrap();
539 d
540 }
541
542 #[test]
543 fn parses_the_two_fields_and_ignores_the_rest() {
544 let (t, a) = parse_meminfo(
545 "MemTotal: 16000000 kB\nMemFree: 1 kB\nMemAvailable: 3000000 kB\nSwapTotal: 0 kB\n",
546 )
547 .unwrap();
548 assert_eq!(t, 16_000_000 * 1024);
549 assert_eq!(a, 3_000_000 * 1024);
550 assert!(parse_meminfo("MemTotal: 1 kB\n").is_none(), "no MemAvailable = no answer");
551 assert!(parse_meminfo("garbage").is_none());
552 }
553
554 #[test]
555 fn above_the_floor_costs_nothing_and_flushes_nothing() {
556 let d = scratch("above");
557 let p = fake_meminfo(&d, 16_000_000, 8_000_000);
558 let w = MemWatchdog::with(&p, 3_200_000 * 1024, Duration::from_millis(1));
559 let mut flushed = 0;
560 let p = w.pause_until_ok("t", None, &mut || {
561 flushed += 1;
562 Ok(1)
563 })
564 .unwrap();
565 assert_eq!(p.ms, 0);
566 assert_eq!(p.end, PauseEnd::Recovered);
567 assert_eq!(flushed, 0);
568 assert_eq!(w.pauses(), 0);
569 assert!(w.breached().is_none());
570 }
571
572 #[test]
577 fn below_the_floor_flushes_once_then_waits_for_memory_to_come_back() {
578 let d = scratch("below");
579 let p = fake_meminfo(&d, 16_000_000, 2_900_000);
580 let w = Arc::new(MemWatchdog::with(&p, 3_200_000 * 1024, Duration::from_millis(5)));
581 assert_eq!(w.breached(), Some((2_900_000 * 1024, 3_200_000 * 1024)));
582 let raise_at = Arc::new(AtomicU64::new(0));
583 let (p2, r2) = (p.clone(), raise_at.clone());
584 let t = std::thread::spawn(move || {
585 std::thread::sleep(Duration::from_millis(120));
586 r2.store(1, Ordering::SeqCst);
587 fake_meminfo(p2.parent().unwrap(), 16_000_000, 9_000_000);
588 });
589 let mut flushed = 0;
590 let paused = w
591 .pause_until_ok("t", None, &mut || {
592 flushed += 1;
593 Ok(1)
594 })
595 .unwrap();
596 let ms = paused.ms;
597 t.join().unwrap();
598 assert_eq!(paused.end, PauseEnd::Recovered);
599 assert_eq!(flushed, 1, "flush runs exactly once per pause");
600 assert_eq!(raise_at.load(Ordering::SeqCst), 1, "the wait outlasted the raise");
601 assert!(ms >= 100, "waited {ms} ms; must have waited for the raise (~120 ms)");
602 assert_eq!(w.pauses(), 1);
603 assert!(w.paused_ms() >= 100);
604 assert!(!w.is_paused());
605 }
606
607 #[test]
608 fn a_cancel_ends_the_wait_without_memory_coming_back() {
609 let d = scratch("cancel");
610 let p = fake_meminfo(&d, 16_000_000, 100);
611 let w = Arc::new(MemWatchdog::with(&p, 3_200_000 * 1024, Duration::from_millis(5)));
612 let cancel = Arc::new(AtomicBool::new(false));
613 let c2 = cancel.clone();
614 let t = std::thread::spawn(move || {
615 std::thread::sleep(Duration::from_millis(60));
616 c2.store(true, Ordering::SeqCst);
617 });
618 let paused = w.pause_until_ok("t", Some(&cancel), &mut || Ok(1)).unwrap();
619 t.join().unwrap();
620 assert_eq!(paused.end, PauseEnd::Cancelled);
621 assert!(
622 (50..2000).contains(&paused.ms),
623 "ended by the cancel, not by memory: {} ms",
624 paused.ms
625 );
626 }
627
628 #[test]
632 fn a_long_pause_reflushes_and_resumes_only_past_the_margin() {
633 let d = scratch("reflush");
634 let p = fake_meminfo(&d, 16_000_000, 1_000_000);
635 let budget = 3_200_000 * 1024;
636 let w = Arc::new(MemWatchdog::with(&p, budget, Duration::from_millis(5)));
637 let p2 = p.clone();
638 let t = std::thread::spawn(move || {
639 std::thread::sleep(Duration::from_millis(60));
640 fake_meminfo(p2.parent().unwrap(), 16_000_000, 3_232_000);
642 std::thread::sleep(Duration::from_millis(120));
643 fake_meminfo(p2.parent().unwrap(), 16_000_000, 9_000_000);
644 });
645 let mut flushes = 0;
646 let paused = w
647 .pause_until_ok("t", None, &mut || {
648 flushes += 1;
649 Ok(1)
650 })
651 .unwrap();
652 t.join().unwrap();
653 let ms = paused.ms;
654 assert_eq!(paused.end, PauseEnd::Recovered);
655 assert!(
656 ms >= 170,
657 "resumed 32 MB over the floor after {ms} ms; the margin is {}",
658 fmt_bytes(w.resume_margin())
659 );
660 assert!(flushes >= 1);
661 assert_eq!(flushes, 1);
664 }
665
666 #[test]
671 fn a_two_percent_shortfall_does_not_demand_twelve_percent_of_headroom() {
672 let d = scratch("livelock");
673 let budget = 3_361_581_460u64; let p = fake_meminfo(&d, 16_000_000, (budget - 75_000_000) / 1024);
676 let w = Arc::new(MemWatchdog::with(&p, budget, Duration::from_millis(5)));
677 assert!(w.breached().is_some());
678 assert_eq!(w.resume_margin(), RESUME_MARGIN_BYTES);
679 assert!(
680 budget / 8 > 400_000_000,
681 "the old margin really was 410 MB on this box: {}",
682 fmt_bytes(budget / 8)
683 );
684 let p2 = p.clone();
685 let t = std::thread::spawn(move || {
686 std::thread::sleep(Duration::from_millis(40));
687 fake_meminfo(p2.parent().unwrap(), 16_000_000, (budget + 100_000_000) / 1024);
689 });
690 let paused = w.pause_until_ok("t", None, &mut || Ok(0)).unwrap();
691 t.join().unwrap();
692 assert_eq!(
693 paused.end,
694 PauseEnd::Recovered,
695 "100 MB over a 3.1G floor is a resume; an eighth of the floor would still be waiting"
696 );
697 assert_eq!(w.give_ups(), 0);
698 }
699
700 #[test]
704 fn a_pause_that_never_recovers_ends_on_its_deadline() {
705 let d = scratch("deadline");
706 let p = fake_meminfo(&d, 16_000_000, 100);
707 let w = MemWatchdog::with(&p, 3_200_000 * 1024, Duration::from_millis(5))
708 .bounded(RESUME_MARGIN_BYTES, Duration::from_millis(120));
709 let started = Instant::now();
710 let paused = w.pause_until_ok("t", None, &mut || Ok(4096)).unwrap();
711 assert_eq!(paused.end, PauseEnd::Deadline);
712 assert!(paused.end.gave_up());
713 assert!(started.elapsed() >= Duration::from_millis(100));
714 assert!(paused.ms < 5_000, "the deadline is the bound, not a suggestion");
715 assert_eq!(w.give_ups(), 1);
716 }
717
718 #[test]
723 fn three_flushes_that_free_nothing_end_the_pause_by_name() {
724 let d = scratch("nothing");
725 let p = fake_meminfo(&d, 16_000_000, 100);
726 let w = MemWatchdog::with(&p, 3_200_000 * 1024, Duration::from_millis(1))
727 .bounded(RESUME_MARGIN_BYTES, Duration::from_secs(600))
728 .with_reflush(Duration::from_millis(5));
729 let mut flushes = 0u32;
730 let started = Instant::now();
731 let paused = w
732 .pause_until_ok("t", None, &mut || {
733 flushes += 1;
734 Ok(0)
735 })
736 .unwrap();
737 assert_eq!(paused.end, PauseEnd::NothingToGive);
738 assert!(paused.end.gave_up());
739 assert_eq!(paused.empty_flushes, GIVE_UP_AFTER_EMPTY_FLUSHES);
740 assert_eq!(flushes, GIVE_UP_AFTER_EMPTY_FLUSHES, "one at the top, then the re-flushes");
741 assert!(started.elapsed() < Duration::from_secs(30), "the deadline was ten minutes");
742 assert_eq!(w.give_ups(), 1);
743 }
744
745 #[test]
749 fn a_flush_that_frees_something_is_not_a_give_up_by_name() {
750 let d = scratch("something");
751 let p = fake_meminfo(&d, 16_000_000, 100);
752 let w = MemWatchdog::with(&p, 3_200_000 * 1024, Duration::from_millis(1))
753 .bounded(RESUME_MARGIN_BYTES, Duration::from_millis(200))
754 .with_reflush(Duration::from_millis(5));
755 let paused = w.pause_until_ok("t", None, &mut || Ok(1)).unwrap();
756 assert_eq!(paused.end, PauseEnd::Deadline);
757 assert_eq!(paused.empty_flushes, 0);
758 }
759
760 #[test]
761 fn a_zero_budget_never_pauses_and_an_unreadable_file_never_pauses() {
762 let w = MemWatchdog::disabled();
763 assert!(w.breached().is_none());
764 let d = scratch("unreadable");
765 let w = MemWatchdog::with(&d.join("missing"), 1 << 40, Duration::from_millis(1));
766 assert!(w.breached().is_none(), "no /proc = no opinion, never a stall");
767 }
768
769 #[test]
770 fn the_floor_is_a_share_of_memtotal_by_default() {
771 let d = scratch("pct");
772 let p = fake_meminfo(&d, 16_000_000, 8_000_000);
773 let total = parse_meminfo(&std::fs::read_to_string(&p).unwrap()).unwrap().0;
775 assert_eq!(total / 100 * DEFAULT_BUDGET_PCT, 3_276_800_000);
776 assert_eq!(fmt_bytes(3_276_800_000), "3.1G");
777 assert_eq!(fmt_bytes(2_900_000 * 1024), "2.8G");
778 assert_eq!(fmt_bytes(512 << 20), "512M");
779 }
780}