refactor(memory_usage): replace sys-info with systemstat (#4028)

src/modules/memory_usage.rs

@@ -1,23 +1,44 @@
-use byte_unit::{Byte, ByteUnit};
+use systemstat::{
+    data::{saturating_sub_bytes, ByteSize},
+    Platform, System,
+};
 
 use super::{Context, Module, ModuleConfig};
 
 use crate::configs::memory_usage::MemoryConfig;
 use crate::formatter::StringFormatter;
 
-fn format_kib(n_kib: u64) -> String {
-    let byte = Byte::from_unit(n_kib as f64, ByteUnit::KiB).unwrap_or_else(|_| Byte::from_bytes(0));
-    let mut display_bytes = byte.get_appropriate_unit(true).format(0);
-    display_bytes.retain(|c| c != ' ');
+// Display a `ByteSize` in a human readable format.
+fn display_bs(bs: ByteSize) -> String {
+    let mut display_bytes = bs.to_string_as(true);
+    let mut keep = true;
+    // Skip decimals and the space before the byte unit.
+    display_bytes.retain(|c| match c {
+        ' ' => {
+            keep = true;
+            false
+        }
+        '.' => {
+            keep = false;
+            false
+        }
+        _ => keep,
+    });
     display_bytes
 }
 
-fn format_pct(pct_number: f64) -> String {
-    format!("{:.0}%", pct_number)
+// Calculate the memory usage from total and free memory
+fn pct(total: ByteSize, free: ByteSize) -> f64 {
+    100.0 * saturating_sub_bytes(total, free).0 as f64 / total.0 as f64
 }
 
-fn format_usage_total(usage: u64, total: u64) -> String {
-    format!("{}/{}", format_kib(usage), format_kib(total))
+// Print usage string used/total
+fn format_usage_total(total: ByteSize, free: ByteSize) -> String {
+    format!(
+        "{}/{}",
+        display_bs(saturating_sub_bytes(total, free)),
+        display_bs(total)
+    )
 }
 
 /// Creates a module with system memory usage information
@@ -31,36 +52,37 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
         return None;
     }
 
-    let system = match sys_info::mem_info() {
-        Ok(info) => info,
-        Err(err) => {
-            log::warn!("Unable to access memory usage information:\n{}", err);
-            return None;
+    let system = System::new();
+
+    // `memory_and_swap` only works on platforms that have an implementation for swap memory
+    // But getting both together is faster on some platforms (Windows/Linux)
+    let (memory, swap) = match system.memory_and_swap() {
+        // Ignore swap if total is 0
+        Ok((mem, swap)) if swap.total.0 > 0 => (mem, Some(swap)),
+        Ok((mem, _)) => (mem, None),
+        Err(e) => {
+            log::debug!(
+                "Failed to retrieve both memory and swap, falling back to memory only: {}",
+                e
+            );
+            let mem = match system.memory() {
+                Ok(mem) => mem,
+                Err(e) => {
+                    log::warn!("Failed to retrieve memory: {}", e);
+                    return None;
+                }
+            };
+
+            (mem, None)
         }
     };
 
-    // avail includes reclaimable memory, but isn't supported on all platforms
-    let avail_memory_kib = match system.avail {
-        0 => system.free,
-        _ => system.avail,
-    };
-    let used_memory_kib = system.total.saturating_sub(avail_memory_kib);
-    let total_memory_kib = system.total;
-    let ram_used = (used_memory_kib as f64 / total_memory_kib as f64) * 100.;
-    let ram_pct = format_pct(ram_used);
+    let used_pct = pct(memory.total, memory.free);
 
-    let threshold = config.threshold;
-    if ram_used.round() < threshold as f64 {
+    if (used_pct.round() as i64) < config.threshold {
         return None;
     }
 
-    let ram = format_usage_total(used_memory_kib, total_memory_kib);
-    let total_swap_kib = system.swap_total;
-    let used_swap_kib = system.swap_total.saturating_sub(system.swap_free);
-    let percent_swap_used = (used_swap_kib as f64 / total_swap_kib as f64) * 100.;
-    let swap_pct = format_pct(percent_swap_used);
-    let swap = format_usage_total(used_swap_kib, total_swap_kib);
-
     let parsed = StringFormatter::new(config.format).and_then(|formatter| {
         formatter
             .map_meta(|var, _| match var {
@@ -72,11 +94,16 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
                 _ => None,
             })
             .map(|variable| match variable {
-                "ram" => Some(Ok(&ram)),
-                "ram_pct" => Some(Ok(&ram_pct)),
-                // swap only shown if there is swap on the system
-                "swap" if total_swap_kib > 0 => Some(Ok(&swap)),
-                "swap_pct" if total_swap_kib > 0 => Some(Ok(&swap_pct)),
+                "ram" => Some(Ok(format_usage_total(memory.total, memory.free))),
+                "ram_pct" => Some(Ok(format!("{:.0}%", used_pct))),
+                "swap" => Some(Ok(format_usage_total(
+                    swap.as_ref()?.total,
+                    swap.as_ref()?.free,
+                ))),
+                "swap_pct" => Some(Ok(format!(
+                    "{:.0}%",
+                    pct(swap.as_ref()?.total, swap.as_ref()?.free)
+                ))),
                 _ => None,
             })
             .parse(None, Some(context))
@@ -92,3 +119,71 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
 
     Some(module)
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    use crate::test::ModuleRenderer;
+
+    #[test]
+    fn test_format_usage_total() {
+        assert_eq!(
+            format_usage_total(ByteSize(1024 * 1024 * 1024), ByteSize(1024 * 1024 * 1024)),
+            "0B/1GiB"
+        );
+        assert_eq!(
+            format_usage_total(
+                ByteSize(1024 * 1024 * 1024),
+                ByteSize(1024 * 1024 * 1024 / 2)
+            ),
+            "512MiB/1GiB"
+        );
+        assert_eq!(
+            format_usage_total(ByteSize(1024 * 1024 * 1024), ByteSize(0)),
+            "1GiB/1GiB"
+        );
+    }
+
+    #[test]
+    fn test_pct() {
+        assert_eq!(
+            pct(ByteSize(1024 * 1024 * 1024), ByteSize(1024 * 1024 * 1024)),
+            0.0
+        );
+        assert_eq!(
+            pct(
+                ByteSize(1024 * 1024 * 1024),
+                ByteSize(1024 * 1024 * 1024 / 2)
+            ),
+            50.0
+        );
+        assert_eq!(pct(ByteSize(1024 * 1024 * 1024), ByteSize(0)), 100.0);
+    }
+
+    #[test]
+    fn zero_threshold() {
+        let output = ModuleRenderer::new("memory_usage")
+            .config(toml::toml! {
+                [memory_usage]
+                disabled = false
+                threshold = 0
+            })
+            .collect();
+
+        assert!(output.is_some())
+    }
+
+    #[test]
+    fn impossible_threshold() {
+        let output = ModuleRenderer::new("memory_usage")
+            .config(toml::toml! {
+                [memory_usage]
+                disabled = false
+                threshold = 9999
+            })
+            .collect();
+
+        assert!(output.is_none())
+    }
+}