improve cgroup memory usage calculation to avoid counting purgeable cache

crates/kumo-server-memory/src/lib.rs

@@ -16,6 +16,8 @@ use cgroups_rs::{Hierarchy, MaxValue};
 use kumo_prometheus::declare_metric;
 use nix::sys::resource::{rlim_t, RLIM_INFINITY};
 use nix::unistd::{sysconf, SysconfVar};
+#[cfg(any(target_os = "linux", test))]
+use std::collections::HashMap;
 use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
 use std::sync::{LazyLock, Mutex};
 use std::time::Duration;
@@ -40,7 +42,10 @@ static OVER_LIMIT_COUNT: IntCounter("memory_over_limit_count");
 }
 
 declare_metric! {
-/// number of bytes of used memory (Resident Set Size)
+/// number of bytes of used memory that drives memory limit decisions.
+/// When sourced from a cgroup this is the working set
+/// (memory.current - inactive_file, floored by anonymous memory);
+/// otherwise it is the Resident Set Size from /proc/self/statm.
 static MEM_USAGE: Gauge("memory_usage");
 }
 
@@ -64,6 +69,11 @@ static SUBSCRIBER: LazyLock<Mutex<Option<Receiver<()>>>> = LazyLock::new(|| Mute
 static OVER_LIMIT: AtomicBool = AtomicBool::new(false);
 static LOW_MEM: AtomicBool = AtomicBool::new(false);
 
+/// Set true once we have warned about an unreadable cgroup memory.stat,
+/// so the fail-safe path warns at most once.
+#[cfg(target_os = "linux")]
+static WARNED_EMPTY_MEMORY_STAT: AtomicBool = AtomicBool::new(false);
+
 // Default this to a reasonable non-zero value, as it is possible
 // in the test harness on the CI system to inject concurrent with
 // early startup.
@@ -72,7 +82,11 @@ static LOW_MEM: AtomicBool = AtomicBool::new(false);
 // have to deal with this small window on startup.
 static HEAD_ROOM: AtomicUsize = AtomicUsize::new(u32::MAX as usize);
 
-/// Represents the current memory usage of this process
+/// Represents the current memory usage of this process.
+///
+/// `bytes` is the value used for all memory-limit decisions. When sourced from
+/// a cgroup it is the working set (`current - inactive_file`, floored by
+/// anonymous memory); otherwise it is the Resident Set Size.
 #[derive(Debug, Clone, Copy)]
 pub struct MemoryUsage {
     pub bytes: u64,
@@ -105,9 +119,25 @@ impl MemoryUsage {
             .controller_of()
             .ok_or_else(|| anyhow::anyhow!("no memory controller?"))?;
         let stat = mem.memory_stat();
-        Ok(Self {
-            bytes: stat.usage_in_bytes,
-        })
+        let current = stat.usage_in_bytes;
+
+        match parse_working_set_inputs(v2, &stat.stat.raw) {
+            Some((inactive_file, anon)) => Ok(Self {
+                bytes: working_set_usage(current, inactive_file, anon),
+            }),
+            None => {
+                // memory.stat was empty/unreadable: fall back to the raw
+                // usage rather than erroring, and warn at most once.
+                if !WARNED_EMPTY_MEMORY_STAT.swap(true, Ordering::Relaxed) {
+                    tracing::warn!(
+                        "cgroup memory.stat was empty or unreadable; \
+                         reporting raw memory.current ({}) as usage",
+                        human(current)
+                    );
+                }
+                Ok(Self { bytes: current })
+            }
+        }
     }
 
     pub fn get_linux_statm() -> anyhow::Result<Self> {
@@ -120,6 +150,40 @@ impl MemoryUsage {
     }
 }
 
+/// Computes the working set from cgroup memory statistics.
+///
+/// `current - inactive_file` removes cold, kernel-reclaimable page cache;
+/// the `max(anon)` floor ensures a fast anonymous burst between samples
+/// cannot hide behind a transiently low working set. Saturating subtraction
+/// guards against `inactive_file` exceeding `current`.
+#[cfg(any(target_os = "linux", test))]
+fn working_set_usage(current: u64, inactive_file: u64, anon: u64) -> u64 {
+    current.saturating_sub(inactive_file).max(anon)
+}
+
+/// Extracts `(inactive_file, anon)` from a parsed `memory.stat` for the
+/// working-set computation, or `None` when `memory.stat` could not be read
+/// (the crate returns an empty `raw` map on failure).
+///
+/// cgroup v2 uses the `inactive_file` / `anon` keys. cgroup v1 uses the
+/// hierarchical `total_inactive_file` key, and `anon` is
+/// `total_rss + total_rss_huge`; the hierarchical (`total_*`) values are used
+/// so the anon floor is consistent with the hierarchical `memory.usage_in_bytes`
+/// it is compared against.
+#[cfg(any(target_os = "linux", test))]
+fn parse_working_set_inputs(v2: bool, raw: &HashMap<String, u64>) -> Option<(u64, u64)> {
+    if raw.is_empty() {
+        return None;
+    }
+    let get = |key: &str| raw.get(key).copied().unwrap_or(0);
+    if v2 {
+        Some((get("inactive_file"), get("anon")))
+    } else {
+        let anon = get("total_rss").saturating_add(get("total_rss_huge"));
+        Some((get("total_inactive_file"), anon))
+    }
+}
+
 fn human(n: u64) -> String {
     humansize::format_size(n, humansize::DECIMAL)
 }
@@ -679,3 +743,95 @@ impl JemallocStats {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    const GB: u64 = 1024 * 1024 * 1024;
+    const MB: u64 = 1024 * 1024;
+
+    fn raw(pairs: &[(&str, u64)]) -> HashMap<String, u64> {
+        pairs
+            .iter()
+            .map(|(key, value)| ((*key).to_string(), *value))
+            .collect()
+    }
+
+    #[test]
+    fn working_set_normal() {
+        // The motivating pod: 51GB current, 46GB cold cache, 1.8GB real
+        // footprint. Working set should drop the cold cache.
+        let current = 51 * GB;
+        let inactive_file = 46 * GB;
+        let anon = 1_800 * MB;
+        assert_eq!(
+            working_set_usage(current, inactive_file, anon),
+            current - inactive_file
+        );
+    }
+
+    #[test]
+    fn anon_floor_wins() {
+        // A fast anon burst exceeds the working set; the floor must win so a
+        // transiently low working set can't hide it.
+        let current = 5 * GB;
+        let inactive_file = 4 * GB;
+        let anon = 4 * GB;
+        assert_eq!(working_set_usage(current, inactive_file, anon), 4 * GB);
+    }
+
+    #[test]
+    fn underflow_no_panic() {
+        // inactive_file >= current must saturate to the anon floor (0 here),
+        // never panic.
+        assert_eq!(working_set_usage(10 * GB, 10 * GB, 0), 0);
+        assert_eq!(working_set_usage(10 * GB, 12 * GB, 0), 0);
+    }
+
+    #[test]
+    fn parse_v2() {
+        // active_file / file are present but irrelevant; only inactive_file
+        // and anon feed the working set.
+        let raw = raw(&[
+            ("anon", 2 * GB),
+            ("inactive_file", 3 * GB),
+            ("active_file", GB),
+            ("file", 4 * GB),
+        ]);
+        assert_eq!(parse_working_set_inputs(true, &raw), Some((3 * GB, 2 * GB)));
+    }
+
+    #[test]
+    fn parse_v1_hierarchical() {
+        // v1 uses the hierarchical total_* keys, and anon is
+        // total_rss + total_rss_huge. The non-hierarchical rss must be ignored.
+        let raw = raw(&[
+            ("total_rss", 2 * GB),
+            ("total_rss_huge", 512 * MB),
+            ("total_inactive_file", 3 * GB),
+            ("rss", GB),
+        ]);
+        assert_eq!(
+            parse_working_set_inputs(false, &raw),
+            Some((3 * GB, 2 * GB + 512 * MB))
+        );
+    }
+
+    #[test]
+    fn empty_raw_is_none() {
+        // An empty raw map is how the crate signals an unreadable memory.stat;
+        // both versions must report None so the caller can fail safe.
+        let raw = HashMap::new();
+        assert_eq!(parse_working_set_inputs(true, &raw), None);
+        assert_eq!(parse_working_set_inputs(false, &raw), None);
+    }
+
+    #[test]
+    fn missing_keys_default_zero() {
+        // A non-empty map missing our keys yields zeros (still Some), so the
+        // working set degrades to current rather than failing safe.
+        let raw = raw(&[("cache", GB)]);
+        assert_eq!(parse_working_set_inputs(true, &raw), Some((0, 0)));
+    }
+}

docs/changelog/main.md

@@ -32,6 +32,17 @@
 
 ## Other Changes and Enhancements
 
+ * When running under a cgroup memory limit, `memory_usage` is now the
+   cgroup *working set* (`memory.current - inactive_file`, floored by
+   anonymous memory) rather than the raw `memory.current`. This excludes
+   cold, kernel-reclaimable file-backed page cache, which the kernel drops
+   before it would OOM-kill the container, and matches what kubelet/cAdvisor
+   report as `container_memory_working_set_bytes`. It removes premature load
+   shedding caused by phantom cache pressure for workloads that write a lot of
+   logs or spool to disk. **This lowers the observed `memory_usage` for any
+   deployment under a cgroup limit; alerts keyed on `memory_usage` may need
+   re-baselining.** See
+   [Memory Management](../reference/memory.md#working-set-under-a-cgroup).
  * [enable_dane](../reference/kumo/make_egress_path/enable_dane.md) can now be
    used with the Hickory resolver (with DNSSEC validation enabled); it no
    longer requires the unbound resolver.

docs/reference/memory.md

@@ -29,16 +29,48 @@ as `memory_limit`, and is reported in bytes.
 A background task is started to monitor the memory usage of the system
 which is derived from:
 
-* If running in a cgroup, the usage reported by that cgroup. Note that
+* If running under a cgroup memory limit, the cgroup *working set* (see
+  [Working Set under a cgroup](#working-set-under-a-cgroup) below). Note that
   this may include memory used by other processes in that same cgroup.
-* The Resident Set Size (RSS) as reported in `/proc/self/statm`
+* Otherwise, the Resident Set Size (RSS) as reported in `/proc/self/statm`.
 * The value configured via
   [kumo.set_memory_soft_limit](kumo/set_memory_soft_limit.md), if any,
   will always take precedence over the above.
 
 The current usage is published via the prometheus metrics endpoint
 as `memory_usage` and is reported in bytes.
 
+### Working Set under a cgroup
+
+When KumoMTA runs under a cgroup memory limit, the value that drives
+`memory_usage` (and therefore headroom and load shedding) depends on the
+KumoMTA version:
+
+|KumoMTA Version|`memory_usage` under a cgroup|
+|---------------|-----------------------------|
+|Earlier versions|raw `memory.current` (cgroup v2) / `memory.usage_in_bytes` (cgroup v1)|
+|{{since('dev', inline=True)}}|working set: `max(current - inactive_file, anon)`|
+
+The earlier raw counter includes cold, kernel-reclaimable file-backed page
+cache (`inactive_file`), which the kernel drops under memory pressure before it
+would OOM-kill the container.  For workloads that write a lot of logs or spool
+to disk that cache can dominate, driving headroom to 0 and triggering load
+shedding against memory that was never at risk.
+
+The working set subtracts that reclaimable cache.  It is the same quantity that
+kubelet/cAdvisor reports as `container_memory_working_set_bytes`, and it aligns
+with the kernel's own OOM reclaimability accounting.  It still counts everything
+that can actually drive an OOM kill: anonymous memory, dirty page cache not yet
+written back, warm file-backed cache (`active_file`), tmpfs/shm, and slab.
+
+The raw cgroup counters that feed this (`memory.current`, `inactive_file`,
+`active_file`) are not republished by KumoMTA; query them from the kernel
+directly or from tooling such as cAdvisor/kubelet, which expose
+`container_memory_working_set_bytes` and the underlying cache figures.
+
+If the cgroup `memory.stat` cannot be read, KumoMTA falls back to reporting the
+raw `memory.current` as `memory_usage` and logs a single warning.
+
 ### Headroom and Load Shedding
 
 The *memory headroom* of the system is defined as the current value of