clusterd: expose heap usage and limit

This commit extends the usage metrics provided by clusterd's
`/api/usage-metrics` endpoint with the memory usage, as well as the
effective heap (memory+swap) limit.

As for the existing usage collection, we mainly focus on Linux support,
for macOS we only provide the required stubs to make sure the code
compiles and run.

Author: teskje

src/clusterd/src/usage_metrics.rs

@@ -15,7 +15,7 @@
 use std::path::PathBuf;
 
 use serde::Serialize;
-use tracing::error;
+use tracing::{debug, error};
 
 /// A system usage metrics collector.
 pub(crate) struct Collector {
@@ -25,9 +25,15 @@ pub(crate) struct Collector {
 impl Collector {
     /// Collect current system usage metrics.
     pub fn collect(&self) -> Usage {
+        let disk_bytes = self.collect_disk_usage();
+        let (memory_bytes, swap_bytes) = collect_heap_usage();
+        let heap_limit = collect_heap_limit();
+
         Usage {
-            disk_bytes: self.collect_disk_usage(),
-            swap_bytes: self.collect_swap_usage(),
+            disk_bytes,
+            memory_bytes,
+            swap_bytes,
+            heap_limit,
         }
     }
 
@@ -49,35 +55,164 @@ impl Collector {
         let used_blocks = u64::from(stat.blocks() - stat.blocks_available());
         let used_bytes = used_blocks * stat.fragment_size();
 
+        debug!("disk usage: {used_bytes}");
+
         Some(used_bytes)
     }
+}
+
+/// A system usage measurement.
+#[derive(Serialize)]
+pub(crate) struct Usage {
+    disk_bytes: Option<u64>,
+    memory_bytes: Option<u64>,
+    swap_bytes: Option<u64>,
+    heap_limit: Option<u64>,
+}
+
+#[cfg(target_os = "linux")]
+mod linux {
+    use std::fs;
+    use std::path::Path;
 
-    #[cfg(target_os = "linux")]
-    fn collect_swap_usage(&self) -> Option<u64> {
-        use mz_compute::memory_limiter::ProcStatus;
+    use mz_compute::memory_limiter;
+    use mz_ore::cast::CastInto;
+    use tracing::{debug, error};
+
+    /// Collect memory and swap usage.
+    pub fn collect_heap_usage() -> (Option<u64>, Option<u64>) {
         use mz_ore::cast::CastInto;
 
-        match ProcStatus::from_proc() {
+        match memory_limiter::ProcStatus::from_proc() {
             Ok(status) => {
-                let bytes = status.vm_swap.cast_into();
-                Some(bytes)
+                let memory_bytes = status.vm_rss.cast_into();
+                let swap_bytes = status.vm_swap.cast_into();
+
+                debug!("memory usage: {memory_bytes}");
+                debug!("swap usage: {swap_bytes}");
+
+                (Some(memory_bytes), Some(swap_bytes))
             }
             Err(err) => {
                 error!("error reading /proc/self/status: {err}");
-                None
+                (None, None)
             }
         }
     }
 
-    #[cfg(not(target_os = "linux"))]
-    fn collect_swap_usage(&self) -> Option<u64> {
-        None
+    /// Collect the heap limit, i.e. memory + swap limit.
+    pub fn collect_heap_limit() -> Option<u64> {
+        // If we don't know the physical limits, we can't know the heap limit.
+        let (phys_mem_limit, phys_swap_limit) = get_physical_limits()?;
+
+        // Limits might be reduced by the cgroup.
+        let (cgroup_mem_limit, cgroup_swap_limit) = get_cgroup_limits();
+        let mem_limit = cgroup_mem_limit.unwrap_or(u64::MAX).min(phys_mem_limit);
+        let swap_limit = cgroup_swap_limit.unwrap_or(u64::MAX).min(phys_swap_limit);
+
+        let heap_limit = mem_limit + swap_limit;
+
+        // Heap limit might be reduced by the memory limiter.
+        let limiter_limit = memory_limiter::get_memory_limit().map(CastInto::cast_into);
+        let heap_limit = limiter_limit.unwrap_or(u64::MAX).min(heap_limit);
+
+        debug!("memory limit: {mem_limit} (phys={phys_mem_limit}, cgroup={cgroup_mem_limit:?})");
+        debug!("swap limit: {swap_limit} (phys={phys_swap_limit}, cgroup={cgroup_swap_limit:?})");
+        debug!("heap limit: {heap_limit} (limiter={limiter_limit:?})");
+
+        Some(heap_limit)
+    }
+
+    /// Collect the physical memory and swap limits.
+    fn get_physical_limits() -> Option<(u64, u64)> {
+        let meminfo = match fs::read_to_string("/proc/meminfo") {
+            Ok(meminfo) => meminfo,
+            Err(error) => {
+                error!("reading `/proc/meminfo`: {error}");
+                return None;
+            }
+        };
+
+        fn parse_kib_line(line: &str) -> Option<u64> {
+            let result = line
+                .split_whitespace()
+                .nth(1)
+                .and_then(|x| x.parse::<u64>().ok())
+                .map(|kib| kib * 1024);
+            if result.is_none() {
+                error!("invalid `/proc/meminfo` line: {line}");
+            }
+            result
+        }
+
+        let mut memory = None;
+        let mut swap = None;
+        for line in meminfo.lines() {
+            if line.starts_with("MemTotal:") {
+                memory = parse_kib_line(line);
+            } else if line.starts_with("SwapTotal:") {
+                swap = parse_kib_line(line);
+            }
+        }
+
+        let Some(memory) = memory else {
+            error!("failed to collect physical memory limit");
+            return None;
+        };
+        let Some(swap) = swap else {
+            error!("failed to collect physical swap limit");
+            return None;
+        };
+
+        Some((memory, swap))
+    }
+
+    /// Collect the memory and swap limits enforced by the current cgroup.
+    ///
+    /// We make the following simplifying assumptions that hold for a standard Kubernetes
+    /// environment:
+    //  * The current process is a member of exactly one cgroups v2 hierarchy.
+    //  * The cgroups hierarchy is mounted at `/sys/fs/cgroup`.
+    //  * The limits are applied to the current cgroup directly (and not one of its ancestors).
+    fn get_cgroup_limits() -> (Option<u64>, Option<u64>) {
+        let Ok(proc_cgroup) = fs::read_to_string("/proc/self/cgroup") else {
+            return (None, None);
+        };
+        let Some(cgroup_path) = proc_cgroup.split(':').nth(2) else {
+            error!("invalid `/proc/self/cgroup` format: {proc_cgroup}");
+            return (None, None);
+        };
+
+        let root = Path::new("/sys/fs/cgroup").join(cgroup_path);
+        let memory_file = root.join("memory.max");
+        let swap_file = root.join("memory.swap.max");
+
+        let memory = fs::read_to_string(memory_file)
+            .ok()
+            .and_then(|s| s.parse().ok());
+        let swap = fs::read_to_string(swap_file)
+            .ok()
+            .and_then(|s| s.parse().ok());
+
+        (memory, swap)
     }
 }
 
-/// A system usage measurement.
-#[derive(Serialize)]
-pub(crate) struct Usage {
-    disk_bytes: Option<u64>,
-    swap_bytes: Option<u64>,
+#[cfg(not(target_os = "linux"))]
+mod macos {
+    use mz_compute::memory_limiter;
+    use mz_ore::cast::CastInto;
+
+    pub fn collect_heap_usage() -> (Option<u64>, Option<u64>) {
+        (None, None)
+    }
+
+    pub fn collect_heap_limit() -> Option<u64> {
+        memory_limiter::get_memory_limit().map(CastInto::cast_into)
+    }
 }
+
+#[cfg(target_os = "linux")]
+use linux::*;
+#[cfg(not(target_os = "linux"))]
+use macos::*;

src/compute/src/memory_limiter.rs

@@ -51,7 +51,8 @@ pub fn start_limiter(memory_limit: usize, metrics_registry: &MetricsRegistry) {
     mz_ore::task::spawn(|| "memory-limiter", LimiterTask::run(config_rx, metrics));
 
     *limiter = Some(Limiter {
-        memory_limit,
+        base_memory_limit: memory_limit,
+        effective_memory_limit: memory_limit,
         config_tx,
     });
 }
@@ -63,17 +64,25 @@ pub fn apply_limiter_config(config: &ConfigSet) {
     }
 }
 
+/// Get the current effective memory limit.
+pub fn get_memory_limit() -> Option<usize> {
+    let limiter = LIMITER.lock().expect("poisoned");
+    limiter.as_ref().map(|l| l.effective_memory_limit)
+}
+
 /// A handle to a running memory limiter task.
 struct Limiter {
-    /// The process memory limit.
-    memory_limit: usize,
+    /// The base process memory limit.
+    base_memory_limit: usize,
+    /// The effective memory limit, obtained by applying dyncfgs to the base limit.
+    effective_memory_limit: usize,
     /// A sender for limiter configuration updates.
     config_tx: UnboundedSender<LimiterConfig>,
 }
 
 impl Limiter {
     /// Apply the given configuration to the limiter.
-    fn apply_config(&self, config: &ConfigSet) {
+    fn apply_config(&mut self, config: &ConfigSet) {
         let mut interval = MEMORY_LIMITER_INTERVAL.get(config);
         // A zero duration means the limiter is disabled. Translate that into an ~infinite duration
         // so the limiter doesn't have to worry about the special case.
@@ -82,12 +91,14 @@ impl Limiter {
         }
 
         let memory_limit =
-            f64::cast_lossy(self.memory_limit) * MEMORY_LIMITER_USAGE_BIAS.get(config);
+            f64::cast_lossy(self.base_memory_limit) * MEMORY_LIMITER_USAGE_BIAS.get(config);
         let memory_limit = usize::cast_lossy(memory_limit);
 
         let burst_budget = f64::cast_lossy(memory_limit) * MEMORY_LIMITER_BURST_FACTOR.get(config);
         let burst_budget = usize::cast_lossy(burst_budget);
 
+        self.effective_memory_limit = memory_limit;
+
         self.config_tx
             .send(LimiterConfig {
                 interval,