feature: take cgroups into account for memory usage (#2061)

This PR adds support for taking cgroup usage/total into account when calculating memory and swap usage.

CHANGELOG.md

@@ -44,6 +44,7 @@ That said, these are more guidelines rather than hard rules, though the project
 
 - [#2031](https://github.com/ClementTsang/bottom/pull/2031): Tweak display/hiding logic for a graph widget's legend.
 - [#2039](https://github.com/ClementTsang/bottom/pull/2039): Replace `hide_table_gap` with `table_gap`.
+- [#2061](https://github.com/ClementTsang/bottom/pull/2061): Take cgroup into account for RAM/swap usage.
 
 ### Other
 

src/collection.rs

@@ -10,6 +10,7 @@ pub mod amd;
 
 #[cfg(target_os = "linux")]
 mod linux {
+    pub mod cgroups;
     pub mod utils;
 }
 
@@ -32,6 +33,8 @@ use starship_battery::{Battery, Manager};
 
 use super::DataFilters;
 use crate::app::layout_manager::UsedWidgets;
+#[cfg(target_os = "linux")]
+use crate::collection::linux::cgroups::CgroupMemCollector;
 #[cfg(any(target_os = "linux", feature = "gpu"))]
 use crate::utils::int_hash::IntHashMap;
 
@@ -188,6 +191,9 @@ pub struct DataCollector {
     gpus_total_mem: Option<u64>,
     #[cfg(feature = "zfs")]
     free_arc_mem: bool,
+
+    #[cfg(target_os = "linux")]
+    cgroup_memory_data: CgroupMemCollector,
 }
 
 const LESS_ROUTINE_TASK_TIME: Duration = Duration::from_secs(60);
@@ -232,6 +238,8 @@ impl DataCollector {
             free_arc_mem: false,
             last_list_collection_time: last_collection_time,
             should_run_less_routine_tasks: true,
+            #[cfg(target_os = "linux")]
+            cgroup_memory_data: CgroupMemCollector::default(),
         }
     }
 
@@ -361,6 +369,9 @@ impl DataCollector {
 
         self.refresh_sysinfo_data();
 
+        #[cfg(target_os = "linux")]
+        self.cgroup_memory_data.refresh();
+
         self.update_cpu_usage();
         self.update_memory_usage();
         self.update_temps();
@@ -483,7 +494,7 @@ impl DataCollector {
     #[inline]
     fn update_memory_usage(&mut self) {
         if self.widgets_to_harvest.use_mem {
-            self.data.memory = memory::get_ram_usage(&self.sys.system);
+            self.data.memory = memory::get_ram_usage(self);
 
             #[cfg(feature = "zfs")]
             {
@@ -523,7 +534,7 @@ impl DataCollector {
                 self.data.cache = memory::get_cache_usage(&self.sys.system);
             }
 
-            self.data.swap = memory::get_swap_usage(&self.sys.system);
+            self.data.swap = memory::get_swap_usage(self);
         }
     }
 

src/collection/linux/cgroups.rs

@@ -0,0 +1,144 @@
+//! cgroup-related code for Linux.
+//!
+//! For info about cgroups, see things like [the kernel docs](https://www.kernel.org/doc/html/latest/admin-guide/cgroup-v2.html)
+//! and [Kubernetes docs](https://kubernetes.io/docs/concepts/architecture/cgroups/#deprecation-of-cgroup-v1).
+
+use std::{fs, io::BufRead};
+
+/// cgroup memory limits.
+#[derive(Debug)]
+pub(crate) enum CgroupMemLimit {
+    Bytes(u64),
+    Max,
+}
+
+/// cgroup memory usage data.
+#[derive(Debug)]
+pub(crate) struct CgroupMemData {
+    pub used_bytes: u64,
+    pub limit: Option<CgroupMemLimit>,
+}
+
+/// overall cgroup memory data.
+#[derive(Default, Debug)]
+pub(crate) struct CgroupMemCollector {
+    pub ram: Option<CgroupMemData>,
+    pub swap: Option<CgroupMemData>,
+}
+
+fn read_u64(path: &str) -> Option<u64> {
+    fs::read_to_string(path).ok()?.trim().parse().ok()
+}
+
+fn read_stat_key(path: &str, key: &str) -> Option<u64> {
+    // TODO: Maybe check if this is worth it for the files we read.
+    let file = fs::File::open(path).ok()?;
+    for line in std::io::BufReader::new(file).lines().map_while(Result::ok) {
+        if let Some(rest) = line.strip_prefix(key) {
+            if let Some(val) = rest.strip_prefix(' ') {
+                return val.trim().parse().ok();
+            }
+        }
+    }
+
+    None
+}
+
+impl CgroupMemCollector {
+    /// Refresh the cgroup memory data.
+    ///
+    /// Based on [docker's CLI](https://github.com/docker/cli/blob/master/cli/command/container/stats_helpers.go#L254).
+    pub(crate) fn refresh(&mut self) {
+        if !self.try_update_memory_cgroup_v1() && !self.try_update_memory_cgroup_v2() {
+            self.ram = None;
+            self.swap = None;
+        }
+    }
+
+    /// Try and update the memory using cgroup v1 semantics. If successful, returns `true`.
+    fn try_update_memory_cgroup_v1(&mut self) -> bool {
+        if let Some(mem_usage) = read_u64("/sys/fs/cgroup/memory/memory.usage_in_bytes") {
+            // --- Memory ---
+            let inactive =
+                read_stat_key("/sys/fs/cgroup/memory/memory.stat", "total_inactive_file");
+            let used_bytes = match inactive {
+                Some(inactive) if inactive < mem_usage => mem_usage - inactive,
+                _ => mem_usage,
+            };
+
+            // Technically if it's like, some insanely high value (https://unix.stackexchange.com/a/421182)
+            // then it's "unlimited" but we can just make it so we take the max of the main and this anyway.
+            let mem_limit_raw = read_u64("/sys/fs/cgroup/memory/memory.limit_in_bytes");
+            let mem_limit = mem_limit_raw.map(CgroupMemLimit::Bytes);
+
+            self.ram = Some(CgroupMemData {
+                used_bytes,
+                limit: mem_limit,
+            });
+
+            // --- Swap ---
+            // Since swap is dependent on the normal memory usage, we couple it together.
+            if let Some(memsw) = read_u64("/sys/fs/cgroup/memory/memory.memsw.usage_in_bytes") {
+                let used_bytes = memsw.saturating_sub(mem_usage);
+
+                // Same idea for here.
+                let swap_limit = read_u64("/sys/fs/cgroup/memory/memory.memsw.limit_in_bytes")
+                    .map(|memsw_limit| memsw_limit.saturating_sub(mem_limit_raw.unwrap_or(0)))
+                    .map(CgroupMemLimit::Bytes);
+
+                self.swap = Some(CgroupMemData {
+                    used_bytes,
+                    limit: swap_limit,
+                });
+            }
+
+            true
+        } else {
+            false
+        }
+    }
+
+    /// Try and update the memory using cgroup v2 semantics. If successful, returns `true`.
+    fn try_update_memory_cgroup_v2(&mut self) -> bool {
+        let mut could_update = false;
+
+        if let Some(mem_current) = read_u64("/sys/fs/cgroup/memory.current") {
+            // --- Memory ---
+            let inactive = read_stat_key("/sys/fs/cgroup/memory.stat", "inactive_file");
+            let used_bytes = match inactive {
+                Some(inactive) if inactive < mem_current => mem_current - inactive,
+                _ => mem_current,
+            };
+
+            let limit = fs::read_to_string("/sys/fs/cgroup/memory.max")
+                .ok()
+                .and_then(|s| match s.trim() {
+                    "max" => Some(CgroupMemLimit::Max),
+                    v => v.parse::<u64>().map(CgroupMemLimit::Bytes).ok(),
+                });
+
+            self.ram = Some(CgroupMemData { used_bytes, limit });
+
+            could_update = true;
+        }
+
+        // --- Swap ---
+        if let Some(swap_current) = read_u64("/sys/fs/cgroup/memory.swap.current") {
+            let limit = fs::read_to_string("/sys/fs/cgroup/memory.swap.max")
+                .ok()
+                .and_then(|s| match s.trim() {
+                    "max" => Some(CgroupMemLimit::Max),
+                    v => v.parse::<u64>().map(CgroupMemLimit::Bytes).ok(),
+                });
+
+            self.swap = Some(CgroupMemData {
+                used_bytes: swap_current,
+                limit,
+            });
+
+            could_update = true;
+        }
+
+        could_update
+    }
+}

src/collection/memory/sysinfo.rs

@@ -2,9 +2,7 @@
 
 use std::num::NonZeroU64;
 
-use sysinfo::System;
-
-use crate::collection::memory::MemData;
+use crate::collection::{DataCollector, memory::MemData};
 
 #[inline]
 fn get_usage(used: u64, total: u64) -> Option<MemData> {
@@ -14,26 +12,88 @@ fn get_usage(used: u64, total: u64) -> Option<MemData> {
     })
 }
 
-/// Returns RAM usage.
-pub(crate) fn get_ram_usage(sys: &System) -> Option<MemData> {
-    get_usage(sys.used_memory(), sys.total_memory())
+/// Returns memory (RAM) usage using sysinfo.
+///
+/// On Linux, this will take cgroup usage/limits into account.
+pub(crate) fn get_ram_usage(collector: &DataCollector) -> Option<MemData> {
+    let sys = &collector.sys.system;
+
+    cfg_if::cfg_if! {
+        if #[cfg(target_os = "linux")] {
+            use crate::collection::linux::cgroups;
+
+            let base_used = sys.used_memory();
+            let base_total = sys.total_memory();
+
+            let (used, total) = match &collector.cgroup_memory_data.ram {
+                Some(cgroup_data) => {
+                    let used = cgroup_data.used_bytes;
+                    let total = match cgroup_data.limit {
+                        Some(cgroups::CgroupMemLimit::Bytes(bytes)) => bytes,
+                        Some(cgroups::CgroupMemLimit::Max) => base_total,
+                        None => base_total,
+                    };
+
+                    (used, total)
+                }
+                None => (base_used, base_total),
+            };
+
+            get_usage(used, total)
+        } else {
+            get_usage(sys.used_memory(), sys.total_memory())
+        }
+    }
 }
 
-/// Returns SWAP usage.
+/// Returns SWAP usage using sysinfo.
+///
+/// On Linux, this will take cgroup usage/limits into account.
 #[cfg(not(target_os = "windows"))]
-pub(crate) fn get_swap_usage(sys: &System) -> Option<MemData> {
-    get_usage(sys.used_swap(), sys.total_swap())
+pub(crate) fn get_swap_usage(collector: &DataCollector) -> Option<MemData> {
+    let sys = &collector.sys.system;
+
+    cfg_if::cfg_if! {
+        if #[cfg(target_os = "linux")] {
+            use crate::collection::linux::cgroups;
+
+            let base_used = sys.used_swap();
+            let base_total = sys.total_swap();
+
+            let (used, total) = match &collector.cgroup_memory_data.swap {
+                Some(cgroup_data) => {
+                    let used = cgroup_data.used_bytes;
+                    let total = match cgroup_data.limit {
+                        Some(cgroups::CgroupMemLimit::Bytes(bytes)) => bytes,
+                        Some(cgroups::CgroupMemLimit::Max) => base_total,
+                        None => base_total,
+                    };
+
+                    (used, total)
+                }
+                None => (base_used, base_total),
+            };
+
+            get_usage(used, total)
+        } else {
+            get_usage(sys.used_swap(), sys.total_swap())
+        }
+    }
 }
 
-/// Returns cache usage. sysinfo has no way to do this directly but it should
-/// equal the difference between the available and free memory. Free memory is
-/// defined as memory not containing any data, which means cache and buffer
-/// memory are not "free". Available memory is defined as memory able
-/// to be allocated by processes, which includes cache and buffer memory. On
-/// Windows, this will always be 0. For more information, see [docs](https://docs.rs/sysinfo/latest/sysinfo/struct.System.html#method.available_memory)
-/// and [memory explanation](https://askubuntu.com/questions/867068/what-is-available-memory-while-using-free-command)
+/// Returns cache usage using sysinfo.
+///
+/// sysinfo has no way to do this directly but it should equal the difference
+/// between the available and free memory. Free memory is defined as memory
+/// not containing any data, which means cache and buffer memory are not
+/// "free". Available memory is defined as memory able to be allocated by
+/// processes, which includes cache and buffer memory. On Windows, this will
+/// always be 0.
+///
+/// For more information, see [sysinfo docs](https://docs.rs/sysinfo/latest/sysinfo/struct.System.html#method.available_memory)
+/// and [this explanation on memory](https://askubuntu.com/questions/867068/what-is-available-memory-while-using-free-command)
 #[cfg(not(target_os = "windows"))]
-pub(crate) fn get_cache_usage(sys: &System) -> Option<MemData> {
+pub(crate) fn get_cache_usage(sys: &sysinfo::System) -> Option<MemData> {
     let mem_used = sys.available_memory().saturating_sub(sys.free_memory());
     let mem_total = sys.total_memory();
 

src/collection/memory/windows.rs

@@ -13,7 +13,7 @@ use windows::{
     core::w,
 };
 
-use crate::collection::memory::MemData;
+use crate::collection::{DataCollector, memory::MemData};
 
 /// Get swap memory usage on Windows. This does it by using checking Windows'
 /// performance counters. This is based on the technique done by psutil [here](https://github.com/giampaolo/psutil/pull/2160).
@@ -25,7 +25,12 @@ use crate::collection::memory::MemData;
 /// - <https://github.com/giampaolo/psutil/issues/2431>
 /// - <https://github.com/oshi/oshi/issues/1175>
 /// - <https://github.com/oshi/oshi/issues/1182>
-pub(crate) fn get_swap_usage(sys: &System) -> Option<MemData> {
+pub(crate) fn get_swap_usage(collector: &DataCollector) -> Option<MemData> {
+    let sys = &collector.sys.system;
+    get_swap_usage_inner(sys)
+}
+
+fn get_swap_usage_inner(sys: &System) -> Option<MemData> {
     let total_bytes = NonZeroU64::new(sys.total_swap())?;
 
     // See https://kennykerr.ca/rust-getting-started/string-tutorial.html
@@ -85,7 +90,7 @@ mod tests {
             RefreshKind::nothing().with_memory(MemoryRefreshKind::nothing().with_swap()),
         );
 
-        let swap_usage = get_swap_usage(&sys);
+        let swap_usage = get_swap_usage_inner(&sys);
         if sys.total_swap() > 0 {
             // Not sure if we can guarantee this to always pass on a machine, so I'll just
             // print out.