use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, OnceLock, RwLock};
use std::time::{Duration, Instant};
use anyhow::Result;
pub const DEFAULT_BUDGET_PCT: u64 = 20;
pub const POLL: Duration = Duration::from_millis(250);
pub const CHECK_EVERY_OBJECTS: u64 = 4096;
pub const REFLUSH_EVERY: Duration = Duration::from_secs(10);
pub const RESUME_MARGIN_BYTES: u64 = 64 << 20;
pub const PAUSE_DEADLINE: Duration = Duration::from_secs(300);
pub const GIVE_UP_AFTER_EMPTY_FLUSHES: u32 = 3;
pub fn parse_meminfo(text: &str) -> Option<(u64, u64)> {
let mut total = None;
let mut avail = None;
for line in text.lines() {
let Some((key, rest)) = line.split_once(':') else { continue };
let Some(kib) = rest.split_ascii_whitespace().next().and_then(|v| v.parse::<u64>().ok())
else {
continue;
};
match key {
"MemTotal" => total = Some(kib * 1024),
"MemAvailable" => avail = Some(kib * 1024),
_ => {}
}
if total.is_some() && avail.is_some() {
break;
}
}
Some((total?, avail?))
}
pub fn fmt_bytes(n: u64) -> String {
const K: u64 = 1024;
if n >= K * K * K {
format!("{:.1}G", n as f64 / (K * K * K) as f64)
} else if n >= K * K {
format!("{}M", n / (K * K))
} else if n >= K {
format!("{}K", n / K)
} else {
format!("{n}B")
}
}
pub struct MemWatchdog {
meminfo: PathBuf,
total: u64,
budget: u64,
poll: Duration,
pauses: AtomicU64,
paused_ms: AtomicU64,
paused_now: AtomicBool,
resume_margin: u64,
deadline: Duration,
reflush: Duration,
give_ups: AtomicU64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PauseEnd {
Recovered,
Cancelled,
Deadline,
NothingToGive,
}
impl PauseEnd {
pub fn as_str(self) -> &'static str {
match self {
PauseEnd::Recovered => "recovered",
PauseEnd::Cancelled => "cancelled",
PauseEnd::Deadline => "deadline",
PauseEnd::NothingToGive => "nothing-to-give",
}
}
pub fn gave_up(self) -> bool {
matches!(self, PauseEnd::Deadline | PauseEnd::NothingToGive)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Pause {
pub ms: u64,
pub end: PauseEnd,
pub available: u64,
pub budget: u64,
pub empty_flushes: u32,
}
impl MemWatchdog {
pub fn from_env() -> Self {
let meminfo = crate::arms::read_env(crate::arms::ENV_MEMINFO)
.map(PathBuf::from)
.unwrap_or_else(|| PathBuf::from("/proc/meminfo"));
let total = std::fs::read_to_string(&meminfo)
.ok()
.and_then(|t| parse_meminfo(&t))
.map(|(t, _)| t)
.unwrap_or(0);
let budget = match crate::arms::read_env(crate::arms::ENV_MEM_BUDGET_BYTES)
.and_then(|v| v.parse::<u64>().ok())
{
Some(b) => b,
None => {
let pct = crate::arms::read_env(crate::arms::ENV_MEM_BUDGET_PCT)
.and_then(|v| v.parse::<u64>().ok())
.unwrap_or(DEFAULT_BUDGET_PCT)
.min(100);
total / 100 * pct
}
};
let margin = crate::arms::read_env(crate::arms::ENV_MEM_RESUME_MARGIN_BYTES)
.and_then(|v| v.parse::<u64>().ok())
.unwrap_or(RESUME_MARGIN_BYTES);
let deadline = crate::arms::read_env(crate::arms::ENV_MEM_PAUSE_MAX_MS)
.and_then(|v| v.parse::<u64>().ok())
.map(Duration::from_millis)
.unwrap_or(PAUSE_DEADLINE);
Self::with(&meminfo, budget, POLL).bounded(margin, deadline)
}
pub fn with(meminfo: &Path, budget: u64, poll: Duration) -> Self {
let total = std::fs::read_to_string(meminfo)
.ok()
.and_then(|t| parse_meminfo(&t))
.map(|(t, _)| t)
.unwrap_or(0);
Self {
meminfo: meminfo.to_path_buf(),
total,
budget,
poll,
pauses: AtomicU64::new(0),
paused_ms: AtomicU64::new(0),
paused_now: AtomicBool::new(false),
resume_margin: RESUME_MARGIN_BYTES,
deadline: PAUSE_DEADLINE,
reflush: REFLUSH_EVERY,
give_ups: AtomicU64::new(0),
}
}
pub fn bounded(mut self, resume_margin: u64, deadline: Duration) -> Self {
self.resume_margin = resume_margin;
self.deadline = deadline;
self
}
pub fn with_reflush(mut self, every: Duration) -> Self {
self.reflush = every;
self
}
pub fn disabled() -> Self {
Self::with(Path::new("/nonexistent"), 0, POLL)
}
pub fn budget(&self) -> u64 {
self.budget
}
pub fn total(&self) -> u64 {
self.total
}
pub fn available(&self) -> Option<u64> {
std::fs::read_to_string(&self.meminfo)
.ok()
.and_then(|t| parse_meminfo(&t))
.map(|(_, a)| a)
}
pub fn breached(&self) -> Option<(u64, u64)> {
if self.budget == 0 {
return None;
}
let avail = self.available()?;
(avail < self.budget).then_some((avail, self.budget))
}
pub fn resume_margin(&self) -> u64 {
self.resume_margin.min(self.budget.max(1))
}
pub fn pause_deadline(&self) -> Duration {
self.deadline
}
fn recovered(&self) -> bool {
match self.available() {
None => true,
Some(avail) => avail >= self.budget.saturating_add(self.resume_margin()),
}
}
pub fn pauses(&self) -> u64 {
self.pauses.load(Ordering::Relaxed)
}
pub fn paused_ms(&self) -> u64 {
self.paused_ms.load(Ordering::Relaxed)
}
pub fn is_paused(&self) -> bool {
self.paused_now.load(Ordering::Relaxed)
}
pub fn give_ups(&self) -> u64 {
self.give_ups.load(Ordering::Relaxed)
}
pub fn pause_until_ok(
&self,
label: &str,
cancel: Option<&AtomicBool>,
flush: &mut dyn FnMut() -> Result<u64>,
) -> Result<Pause> {
let Some((avail, budget)) = self.breached() else {
return Ok(Pause {
ms: 0,
end: PauseEnd::Recovered,
available: self.available().unwrap_or(0),
budget: self.budget,
empty_flushes: 0,
});
};
self.pauses.fetch_add(1, Ordering::Relaxed);
self.paused_now.store(true, Ordering::Relaxed);
let margin = self.resume_margin();
tracing::warn!(
target: "gunnar::index",
event = "index.paused",
repo = label,
available = avail,
budget = budget,
resume_at = budget.saturating_add(margin),
deadline_ms = self.deadline.as_millis() as u64,
"index paused: {} available < {} budget; resumes at {}, gives up after {} s",
fmt_bytes(avail),
fmt_bytes(budget),
fmt_bytes(budget.saturating_add(margin)),
self.deadline.as_secs()
);
let started = Instant::now();
let mut empty_flushes = 0u32;
let mut flush_err: Option<anyhow::Error> = None;
match flush() {
Ok(0) => empty_flushes = 1,
Ok(_) => {}
Err(e) => flush_err = Some(e),
}
let mut last_flush = Instant::now();
let end = loop {
if cancel.is_some_and(|c| c.load(Ordering::Relaxed)) {
break PauseEnd::Cancelled;
}
if self.recovered() {
break PauseEnd::Recovered;
}
if !self.deadline.is_zero() && started.elapsed() >= self.deadline {
break PauseEnd::Deadline;
}
if empty_flushes >= GIVE_UP_AFTER_EMPTY_FLUSHES
&& self.available().is_some_and(|a| a <= avail)
{
break PauseEnd::NothingToGive;
}
if last_flush.elapsed() >= self.reflush {
match flush() {
Ok(0) => empty_flushes += 1,
Ok(_) => empty_flushes = 0,
Err(e) => {
if flush_err.is_none() {
flush_err = Some(e);
}
}
}
last_flush = Instant::now();
}
std::thread::sleep(self.poll);
};
let ms = started.elapsed().as_millis() as u64;
self.paused_ms.fetch_add(ms, Ordering::Relaxed);
self.paused_now.store(false, Ordering::Relaxed);
if end.gave_up() {
self.give_ups.fetch_add(1, Ordering::Relaxed);
}
let now = self.available().unwrap_or(0);
if end.gave_up() {
tracing::warn!(
target: "gunnar::index",
event = "index.gave_up",
repo = label,
available = now,
budget = budget,
paused_ms = ms,
empty_flushes = empty_flushes,
reason = end.as_str(),
"index gave up waiting for memory after {ms} ms ({}): {} available, {} needed. \
The caller decides — proceed, shed, or refuse the pack by name — but nobody \
waits here again.",
end.as_str(),
fmt_bytes(now),
fmt_bytes(budget.saturating_add(margin))
);
} else {
tracing::info!(
target: "gunnar::index",
event = "index.resumed",
repo = label,
available = now,
budget = budget,
paused_ms = ms,
reason = end.as_str(),
"index resumed after {ms} ms ({}): {} available",
end.as_str(),
fmt_bytes(now)
);
}
if let Some(e) = flush_err {
return Err(e);
}
Ok(Pause { ms, end, available: now, budget, empty_flushes })
}
}
fn global_slot() -> &'static RwLock<Arc<MemWatchdog>> {
static G: OnceLock<RwLock<Arc<MemWatchdog>>> = OnceLock::new();
G.get_or_init(|| RwLock::new(Arc::new(MemWatchdog::from_env())))
}
pub fn global() -> Arc<MemWatchdog> {
global_slot().read().unwrap_or_else(|e| e.into_inner()).clone()
}
pub fn set_global(dog: Arc<MemWatchdog>) -> Arc<MemWatchdog> {
std::mem::replace(&mut *global_slot().write().unwrap_or_else(|e| e.into_inner()), dog)
}
pub fn peak_rss_bytes() -> u64 {
std::fs::read_to_string("/proc/self/status")
.ok()
.and_then(|t| {
t.lines().find_map(|l| {
l.strip_prefix("VmHWM:")
.and_then(|r| r.split_ascii_whitespace().next())
.and_then(|v| v.parse::<u64>().ok())
.map(|kib| kib * 1024)
})
})
.unwrap_or(0)
}
pub fn mem_total() -> u64 {
global().total()
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Arc;
use std::sync::atomic::AtomicU64;
fn fake_meminfo(dir: &Path, total_kib: u64, avail_kib: u64) -> PathBuf {
let p = dir.join("meminfo");
std::fs::write(
&p,
format!(
"MemTotal: {total_kib} kB\nMemFree: {avail_kib} kB\nMemAvailable: {avail_kib} kB\nBuffers: 0 kB\n"
),
)
.unwrap();
p
}
fn scratch(name: &str) -> PathBuf {
let d = std::env::temp_dir().join(format!("memguard-{}-{name}", std::process::id()));
let _ = std::fs::remove_dir_all(&d);
std::fs::create_dir_all(&d).unwrap();
d
}
#[test]
fn parses_the_two_fields_and_ignores_the_rest() {
let (t, a) = parse_meminfo(
"MemTotal: 16000000 kB\nMemFree: 1 kB\nMemAvailable: 3000000 kB\nSwapTotal: 0 kB\n",
)
.unwrap();
assert_eq!(t, 16_000_000 * 1024);
assert_eq!(a, 3_000_000 * 1024);
assert!(parse_meminfo("MemTotal: 1 kB\n").is_none(), "no MemAvailable = no answer");
assert!(parse_meminfo("garbage").is_none());
}
#[test]
fn above_the_floor_costs_nothing_and_flushes_nothing() {
let d = scratch("above");
let p = fake_meminfo(&d, 16_000_000, 8_000_000);
let w = MemWatchdog::with(&p, 3_200_000 * 1024, Duration::from_millis(1));
let mut flushed = 0;
let p = w.pause_until_ok("t", None, &mut || {
flushed += 1;
Ok(1)
})
.unwrap();
assert_eq!(p.ms, 0);
assert_eq!(p.end, PauseEnd::Recovered);
assert_eq!(flushed, 0);
assert_eq!(w.pauses(), 0);
assert!(w.breached().is_none());
}
#[test]
fn below_the_floor_flushes_once_then_waits_for_memory_to_come_back() {
let d = scratch("below");
let p = fake_meminfo(&d, 16_000_000, 2_900_000);
let w = Arc::new(MemWatchdog::with(&p, 3_200_000 * 1024, Duration::from_millis(5)));
assert_eq!(w.breached(), Some((2_900_000 * 1024, 3_200_000 * 1024)));
let raise_at = Arc::new(AtomicU64::new(0));
let (p2, r2) = (p.clone(), raise_at.clone());
let t = std::thread::spawn(move || {
std::thread::sleep(Duration::from_millis(120));
r2.store(1, Ordering::SeqCst);
fake_meminfo(p2.parent().unwrap(), 16_000_000, 9_000_000);
});
let mut flushed = 0;
let paused = w
.pause_until_ok("t", None, &mut || {
flushed += 1;
Ok(1)
})
.unwrap();
let ms = paused.ms;
t.join().unwrap();
assert_eq!(paused.end, PauseEnd::Recovered);
assert_eq!(flushed, 1, "flush runs exactly once per pause");
assert_eq!(raise_at.load(Ordering::SeqCst), 1, "the wait outlasted the raise");
assert!(ms >= 100, "waited {ms} ms; must have waited for the raise (~120 ms)");
assert_eq!(w.pauses(), 1);
assert!(w.paused_ms() >= 100);
assert!(!w.is_paused());
}
#[test]
fn a_cancel_ends_the_wait_without_memory_coming_back() {
let d = scratch("cancel");
let p = fake_meminfo(&d, 16_000_000, 100);
let w = Arc::new(MemWatchdog::with(&p, 3_200_000 * 1024, Duration::from_millis(5)));
let cancel = Arc::new(AtomicBool::new(false));
let c2 = cancel.clone();
let t = std::thread::spawn(move || {
std::thread::sleep(Duration::from_millis(60));
c2.store(true, Ordering::SeqCst);
});
let paused = w.pause_until_ok("t", Some(&cancel), &mut || Ok(1)).unwrap();
t.join().unwrap();
assert_eq!(paused.end, PauseEnd::Cancelled);
assert!(
(50..2000).contains(&paused.ms),
"ended by the cancel, not by memory: {} ms",
paused.ms
);
}
#[test]
fn a_long_pause_reflushes_and_resumes_only_past_the_margin() {
let d = scratch("reflush");
let p = fake_meminfo(&d, 16_000_000, 1_000_000);
let budget = 3_200_000 * 1024;
let w = Arc::new(MemWatchdog::with(&p, budget, Duration::from_millis(5)));
let p2 = p.clone();
let t = std::thread::spawn(move || {
std::thread::sleep(Duration::from_millis(60));
fake_meminfo(p2.parent().unwrap(), 16_000_000, 3_232_000);
std::thread::sleep(Duration::from_millis(120));
fake_meminfo(p2.parent().unwrap(), 16_000_000, 9_000_000);
});
let mut flushes = 0;
let paused = w
.pause_until_ok("t", None, &mut || {
flushes += 1;
Ok(1)
})
.unwrap();
t.join().unwrap();
let ms = paused.ms;
assert_eq!(paused.end, PauseEnd::Recovered);
assert!(
ms >= 170,
"resumed 32 MB over the floor after {ms} ms; the margin is {}",
fmt_bytes(w.resume_margin())
);
assert!(flushes >= 1);
assert_eq!(flushes, 1);
}
#[test]
fn a_two_percent_shortfall_does_not_demand_twelve_percent_of_headroom() {
let d = scratch("livelock");
let budget = 3_361_581_460u64; let p = fake_meminfo(&d, 16_000_000, (budget - 75_000_000) / 1024);
let w = Arc::new(MemWatchdog::with(&p, budget, Duration::from_millis(5)));
assert!(w.breached().is_some());
assert_eq!(w.resume_margin(), RESUME_MARGIN_BYTES);
assert!(
budget / 8 > 400_000_000,
"the old margin really was 410 MB on this box: {}",
fmt_bytes(budget / 8)
);
let p2 = p.clone();
let t = std::thread::spawn(move || {
std::thread::sleep(Duration::from_millis(40));
fake_meminfo(p2.parent().unwrap(), 16_000_000, (budget + 100_000_000) / 1024);
});
let paused = w.pause_until_ok("t", None, &mut || Ok(0)).unwrap();
t.join().unwrap();
assert_eq!(
paused.end,
PauseEnd::Recovered,
"100 MB over a 3.1G floor is a resume; an eighth of the floor would still be waiting"
);
assert_eq!(w.give_ups(), 0);
}
#[test]
fn a_pause_that_never_recovers_ends_on_its_deadline() {
let d = scratch("deadline");
let p = fake_meminfo(&d, 16_000_000, 100);
let w = MemWatchdog::with(&p, 3_200_000 * 1024, Duration::from_millis(5))
.bounded(RESUME_MARGIN_BYTES, Duration::from_millis(120));
let started = Instant::now();
let paused = w.pause_until_ok("t", None, &mut || Ok(4096)).unwrap();
assert_eq!(paused.end, PauseEnd::Deadline);
assert!(paused.end.gave_up());
assert!(started.elapsed() >= Duration::from_millis(100));
assert!(paused.ms < 5_000, "the deadline is the bound, not a suggestion");
assert_eq!(w.give_ups(), 1);
}
#[test]
fn three_flushes_that_free_nothing_end_the_pause_by_name() {
let d = scratch("nothing");
let p = fake_meminfo(&d, 16_000_000, 100);
let w = MemWatchdog::with(&p, 3_200_000 * 1024, Duration::from_millis(1))
.bounded(RESUME_MARGIN_BYTES, Duration::from_secs(600))
.with_reflush(Duration::from_millis(5));
let mut flushes = 0u32;
let started = Instant::now();
let paused = w
.pause_until_ok("t", None, &mut || {
flushes += 1;
Ok(0)
})
.unwrap();
assert_eq!(paused.end, PauseEnd::NothingToGive);
assert!(paused.end.gave_up());
assert_eq!(paused.empty_flushes, GIVE_UP_AFTER_EMPTY_FLUSHES);
assert_eq!(flushes, GIVE_UP_AFTER_EMPTY_FLUSHES, "one at the top, then the re-flushes");
assert!(started.elapsed() < Duration::from_secs(30), "the deadline was ten minutes");
assert_eq!(w.give_ups(), 1);
}
#[test]
fn a_flush_that_frees_something_is_not_a_give_up_by_name() {
let d = scratch("something");
let p = fake_meminfo(&d, 16_000_000, 100);
let w = MemWatchdog::with(&p, 3_200_000 * 1024, Duration::from_millis(1))
.bounded(RESUME_MARGIN_BYTES, Duration::from_millis(200))
.with_reflush(Duration::from_millis(5));
let paused = w.pause_until_ok("t", None, &mut || Ok(1)).unwrap();
assert_eq!(paused.end, PauseEnd::Deadline);
assert_eq!(paused.empty_flushes, 0);
}
#[test]
fn a_zero_budget_never_pauses_and_an_unreadable_file_never_pauses() {
let w = MemWatchdog::disabled();
assert!(w.breached().is_none());
let d = scratch("unreadable");
let w = MemWatchdog::with(&d.join("missing"), 1 << 40, Duration::from_millis(1));
assert!(w.breached().is_none(), "no /proc = no opinion, never a stall");
}
#[test]
fn the_floor_is_a_share_of_memtotal_by_default() {
let d = scratch("pct");
let p = fake_meminfo(&d, 16_000_000, 8_000_000);
let total = parse_meminfo(&std::fs::read_to_string(&p).unwrap()).unwrap().0;
assert_eq!(total / 100 * DEFAULT_BUDGET_PCT, 3_276_800_000);
assert_eq!(fmt_bytes(3_276_800_000), "3.1G");
assert_eq!(fmt_bytes(2_900_000 * 1024), "2.8G");
assert_eq!(fmt_bytes(512 << 20), "512M");
}
}