Fix get available memory to work with cgroup v1

Author: preinlein

.github/workflows/container.yml

@@ -77,7 +77,7 @@ jobs:
   copy-to-ecr:
     name: Copy ${{ matrix.arch }} to ECR
     needs: build
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     strategy:
       matrix:
         arch: [amd64, arm64]
@@ -137,7 +137,7 @@ jobs:
   manifest:
     name: Create ${{ matrix.registry }} manifest
     needs: [build, copy-to-ecr]
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     strategy:
       matrix:
         registry: [ghcr, ecr]

lading/src/lib.rs

@@ -22,6 +22,10 @@ pub(crate) mod proto;
 pub mod target;
 pub mod target_metrics;
 
+use std::fs::read_to_string;
+use std::io::Error;
+use std::path::PathBuf;
+
 use byte_unit::Byte;
 use sysinfo::System;
 
@@ -34,21 +38,256 @@ pub(crate) fn full<T: Into<bytes::Bytes>>(
         .boxed()
 }
 
-/// Get available memory for the process, checking cgroup v2 limits first,
+/// Get available memory for the process, checking cgroup limits first,
 /// then falling back to system memory.
+///
+/// Resolves the process's own cgroup path from `/proc/self/cgroup`, then
+/// checks the memory limit at that path. Falls back to the cgroup root, then
+/// cgroup v1, and finally total system memory.
 #[must_use]
 pub fn get_available_memory() -> Byte {
-    if let Ok(content) = std::fs::read_to_string("/sys/fs/cgroup/memory.max") {
-        let content = content.trim();
-        if content == "max" {
-            return Byte::from_u64(u64::MAX);
+    get_available_memory_with(read_to_string)
+}
+
+/// The cgroup v1 "no limit" sentinel: page-aligned `i64::MAX`.
+const CGROUP_V1_NO_LIMIT: u64 = 0x7FFF_FFFF_FFFF_F000;
+
+fn parse_cgroup_limit(content: &str) -> Option<Byte> {
+    let content = content.trim();
+    if content == "max" {
+        return None;
+    }
+    let bytes: u64 = content.parse().ok()?;
+    if bytes >= CGROUP_V1_NO_LIMIT {
+        return None;
+    }
+    Some(Byte::from_u64(bytes))
+}
+
+/// Parse `/proc/self/cgroup` to find the cgroup v2 path for this process.
+///
+/// In cgroup v2 unified mode, the file contains a single line like `0::/path`.
+/// Returns the path portion (e.g., `/path`), or `None` if not found.
+fn resolve_cgroup_v2_path(proc_cgroup_content: &str) -> Option<&str> {
+    for line in proc_cgroup_content.lines() {
+        // cgroup v2 unified entries start with "0::"
+        if let Some(path) = line.strip_prefix("0::") {
+            let path = path.trim();
+            if !path.is_empty() && path != "/" {
+                return Some(path);
+            }
         }
-        let ignore_case = true;
-        if let Ok(limit) = Byte::parse_str(content.trim(), ignore_case) {
+    }
+    None
+}
+
+fn get_available_memory_with<F>(read_file: F) -> Byte
+where
+    F: Fn(PathBuf) -> Result<String, Error>,
+{
+    // Try process-specific cgroup v2 path (handles host cgroup namespace)
+    if let Ok(proc_cgroup) = read_file(PathBuf::from("/proc/self/cgroup"))
+        && let Some(cgroup_path) = resolve_cgroup_v2_path(&proc_cgroup)
+    {
+        let memory_max = PathBuf::from(format!("/sys/fs/cgroup{cgroup_path}/memory.max"));
+        if let Ok(content) = read_file(memory_max)
+            && let Some(limit) = parse_cgroup_limit(&content)
+        {
             return limit;
         }
     }
 
+    // Try cgroup v2 root (works when container has its own cgroup namespace)
+    if let Ok(content) = read_file(PathBuf::from("/sys/fs/cgroup/memory.max"))
+        && let Some(limit) = parse_cgroup_limit(&content)
+    {
+        return limit;
+    }
+
+    // Try cgroup v1
+    if let Ok(content) = read_file(PathBuf::from("/sys/fs/cgroup/memory/memory.limit_in_bytes"))
+        && let Some(limit) = parse_cgroup_limit(&content)
+    {
+        return limit;
+    }
+
     let sys = System::new_all();
     Byte::from_u64(sys.total_memory())
 }
+
+#[cfg(test)]
+mod tests {
+    use std::io::ErrorKind;
+
+    use super::*;
+
+    fn mock_reader(
+        files: &'static [(&'static str, &'static str)],
+    ) -> impl Fn(PathBuf) -> Result<String, Error> {
+        move |path| {
+            for (p, content) in files {
+                if PathBuf::from(p) == path {
+                    return Ok((*content).to_string());
+                }
+            }
+            Err(Error::new(ErrorKind::NotFound, "not found"))
+        }
+    }
+
+    #[test]
+    fn cgroup_v2_max_falls_back_to_system_memory() {
+        let read = mock_reader(&[("/sys/fs/cgroup/memory.max", "max\n")]);
+        let mem = get_available_memory_with(read);
+        let sys = System::new_all();
+        assert_eq!(mem, Byte::from_u64(sys.total_memory()));
+    }
+
+    #[test]
+    fn cgroup_v2_numeric_limit() {
+        let read = mock_reader(&[("/sys/fs/cgroup/memory.max", "12582912\n")]);
+        let mem = get_available_memory_with(read);
+        assert_eq!(mem, Byte::from_u64(12_582_912));
+    }
+
+    #[test]
+    fn cgroup_v2_garbage_falls_back_to_system_memory() {
+        let read = mock_reader(&[("/sys/fs/cgroup/memory.max", "not-a-number\n")]);
+        let mem = get_available_memory_with(read);
+        let sys = System::new_all();
+        assert_eq!(mem, Byte::from_u64(sys.total_memory()));
+    }
+
+    #[test]
+    fn no_cgroup_file_falls_back_to_system_memory() {
+        let read = mock_reader(&[]);
+        let mem = get_available_memory_with(read);
+        let sys = System::new_all();
+        assert_eq!(mem, Byte::from_u64(sys.total_memory()));
+    }
+
+    #[test]
+    fn cgroup_v1_numeric_limit() {
+        let read = mock_reader(&[("/sys/fs/cgroup/memory/memory.limit_in_bytes", "12582912\n")]);
+        let mem = get_available_memory_with(read);
+        assert_eq!(mem, Byte::from_u64(12_582_912));
+    }
+
+    #[test]
+    fn cgroup_v1_no_limit_sentinel_falls_back_to_system_memory() {
+        // cgroup v1 uses page-aligned i64::MAX as "no limit"
+        let read = mock_reader(&[(
+            "/sys/fs/cgroup/memory/memory.limit_in_bytes",
+            "9223372036854771712\n",
+        )]);
+        let mem = get_available_memory_with(read);
+        let sys = System::new_all();
+        assert_eq!(mem, Byte::from_u64(sys.total_memory()));
+    }
+
+    #[test]
+    fn cgroup_v2_takes_precedence_over_v1() {
+        let read = mock_reader(&[
+            ("/sys/fs/cgroup/memory.max", "8388608\n"),
+            ("/sys/fs/cgroup/memory/memory.limit_in_bytes", "16777216\n"),
+        ]);
+        let mem = get_available_memory_with(read);
+        assert_eq!(mem, Byte::from_u64(8_388_608));
+    }
+
+    #[test]
+    fn cgroup_v2_max_falls_through_to_v1() {
+        let read = mock_reader(&[
+            ("/sys/fs/cgroup/memory.max", "max\n"),
+            ("/sys/fs/cgroup/memory/memory.limit_in_bytes", "16777216\n"),
+        ]);
+        let mem = get_available_memory_with(read);
+        assert_eq!(mem, Byte::from_u64(16_777_216));
+    }
+
+    #[test]
+    fn process_specific_cgroup_v2_path() {
+        // Simulates host cgroup namespace: /proc/self/cgroup points to a
+        // nested cgroup, and the memory limit is at that nested path.
+        let read = mock_reader(&[
+            (
+                "/proc/self/cgroup",
+                "0::/system.slice/docker-abc123.scope\n",
+            ),
+            (
+                "/sys/fs/cgroup/system.slice/docker-abc123.scope/memory.max",
+                "12582912\n",
+            ),
+            // Root cgroup has no limit
+            ("/sys/fs/cgroup/memory.max", "max\n"),
+        ]);
+        let mem = get_available_memory_with(read);
+        assert_eq!(mem, Byte::from_u64(12_582_912));
+    }
+
+    #[test]
+    fn process_specific_path_takes_precedence_over_root() {
+        let read = mock_reader(&[
+            (
+                "/proc/self/cgroup",
+                "0::/system.slice/docker-abc123.scope\n",
+            ),
+            (
+                "/sys/fs/cgroup/system.slice/docker-abc123.scope/memory.max",
+                "8388608\n",
+            ),
+            ("/sys/fs/cgroup/memory.max", "16777216\n"),
+        ]);
+        let mem = get_available_memory_with(read);
+        assert_eq!(mem, Byte::from_u64(8_388_608));
+    }
+
+    #[test]
+    fn process_specific_path_max_falls_through_to_root() {
+        let read = mock_reader(&[
+            (
+                "/proc/self/cgroup",
+                "0::/system.slice/docker-abc123.scope\n",
+            ),
+            (
+                "/sys/fs/cgroup/system.slice/docker-abc123.scope/memory.max",
+                "max\n",
+            ),
+            ("/sys/fs/cgroup/memory.max", "16777216\n"),
+        ]);
+        let mem = get_available_memory_with(read);
+        assert_eq!(mem, Byte::from_u64(16_777_216));
+    }
+
+    #[test]
+    fn proc_cgroup_root_path_skipped() {
+        // When /proc/self/cgroup shows root path "/", skip it and fall
+        // through to the root cgroup check.
+        let read = mock_reader(&[
+            ("/proc/self/cgroup", "0::/\n"),
+            ("/sys/fs/cgroup/memory.max", "12582912\n"),
+        ]);
+        let mem = get_available_memory_with(read);
+        assert_eq!(mem, Byte::from_u64(12_582_912));
+    }
+
+    #[test]
+    fn resolve_cgroup_v2_path_parses_correctly() {
+        assert_eq!(
+            resolve_cgroup_v2_path("0::/system.slice/docker-abc.scope\n"),
+            Some("/system.slice/docker-abc.scope")
+        );
+        assert_eq!(resolve_cgroup_v2_path("0::/\n"), None);
+        assert_eq!(resolve_cgroup_v2_path(""), None);
+        // cgroup v1 lines (non-zero hierarchy ID) are ignored
+        assert_eq!(
+            resolve_cgroup_v2_path("1:memory:/docker/abc\n0::/system.slice\n"),
+            Some("/system.slice")
+        );
+    }
+
+    #[test]
+    fn get_available_memory_returns_nonzero() {
+        let mem = get_available_memory();
+        assert!(mem.as_u64() > 0);
+    }
+}