Simplifications

.github/workflows/ci.yml

@@ -54,6 +54,7 @@ jobs:
       - uses: dtolnay/rust-toolchain@stable
         with:
           components: rustfmt, clippy
+      - uses: Swatinem/rust-cache@v2
       - run: cargo fmt --all -- --check
       - run: cargo clippy --no-deps
 

src/lib.rs

@@ -3,74 +3,47 @@ pub mod stats;
 pub mod store;
 pub mod utils;
 
+use crate::stats::cgroup_v1::CgroupV1Source;
+use crate::stats::cgroup_v2::CgroupV2Source;
+use crate::stats::proc::ProcSource;
 use crate::stats::{
-    CpuUsageValue, detect_cgroup_version, get_cgroup_v1_mount_points, get_cgroup_v2_mount_point,
+    SystemStatsSource, detect_cgroup_version, get_cgroup_v1_mount_points, get_cgroup_v2_mount_point,
 };
 use crate::store::StatsEntry;
+use std::path::PathBuf;
 use std::thread;
 use tracing::{debug, error};
 
 pub fn run_acolyte() {
     let stat_interval = env::get_stat_interval();
-    let cgroup_version = detect_cgroup_version("/proc/self/cgroup").ok();
-    let v2_mount_point = cgroup_version
-        .as_ref()
-        .filter(|v| v.has_v2())
-        .and_then(|_| get_cgroup_v2_mount_point("/proc/mounts").ok());
-    let v1_mount_points = cgroup_version
-        .as_ref()
-        .filter(|v| v.has_v1())
-        .and_then(|_| get_cgroup_v1_mount_points("/proc/mounts").ok());
+
+    let sources = get_sources();
 
     loop {
         let mut stats_entry = StatsEntry::new();
 
-        let maybe_num_cpus = stats::get_num_cpus(
-            cgroup_version.clone(),
-            v2_mount_point.clone(),
-            v1_mount_points.clone(),
-        );
-        if let Some(num_cpus) = maybe_num_cpus {
+        if let Some(num_cpus) = sources.iter().find_map(|source| source.get_num_cpus().ok()) {
             stats_entry.num_cpus = Some(num_cpus);
         }
 
-        if let Some(cpu_usage) = stats::get_cpu_usage(
-            cgroup_version.clone(),
-            v2_mount_point.clone(),
-            v1_mount_points.clone(),
-        ) {
-            // scale the cpu usage by the number of cpus
-            // so that 100% cpu usage on a 4 core machine is 4.0 etc.
-            let normalized_cpu_usage = match cpu_usage {
-                CpuUsageValue::FromCgroupV2(cgroup_usage) => Some(cgroup_usage),
-                CpuUsageValue::FromCgroupV1(cgroup_usage) => Some(cgroup_usage),
-                CpuUsageValue::FromProc(proc_usage) => {
-                    // for the `procfs` values to report the number in the right format,
-                    // we MUST know the number of cpus or the number will be misleading
-                    if let Some(num_cpus) = maybe_num_cpus {
-                        Some(proc_usage * num_cpus)
-                    } else {
-                        debug!("Failed to get number of CPUs, skipping procfs CPU usage");
-                        None
-                    }
-                }
-            };
-            stats_entry.cpu_usage = normalized_cpu_usage;
+        if let Some(cpu_usage) = sources
+            .iter()
+            .find_map(|source| source.get_cpu_usage().ok())
+        {
+            stats_entry.cpu_usage = cpu_usage.normalize(stats_entry.num_cpus);
         }
 
-        if let Some(mem_usage_kb) = stats::get_memory_usage_kb(
-            cgroup_version.clone(),
-            v2_mount_point.clone(),
-            v1_mount_points.clone(),
-        ) {
+        if let Some(mem_usage_kb) = sources
+            .iter()
+            .find_map(|source| source.get_memory_usage_kb().ok())
+        {
             stats_entry.memory_usage_kb = Some(mem_usage_kb);
         }
 
-        if let Some(mem_total_kb) = stats::get_memory_total_kb(
-            cgroup_version.clone(),
-            v2_mount_point.clone(),
-            v1_mount_points.clone(),
-        ) {
+        if let Some(mem_total_kb) = sources
+            .iter()
+            .find_map(|source| source.get_memory_total_kb().ok())
+        {
             stats_entry.memory_total_kb = Some(mem_total_kb);
         }
 
@@ -89,3 +62,31 @@ pub fn run_acolyte() {
         thread::sleep(stat_interval);
     }
 }
+
+fn get_sources() -> Vec<Box<dyn SystemStatsSource>> {
+    let mut sources: Vec<Box<dyn SystemStatsSource>> = vec![];
+    let cgroup_version = detect_cgroup_version("/proc/self/cgroup").ok();
+
+    if let Some(v2_mount_point) = cgroup_version
+        .as_ref()
+        .filter(|v| v.has_v2())
+        .and_then(|_| get_cgroup_v2_mount_point("/proc/mounts").ok())
+    {
+        sources.push(Box::new(CgroupV2Source::with_filesystem_reader_at(
+            v2_mount_point,
+        )));
+    }
+    if let Some(v1_mount_points) = cgroup_version
+        .as_ref()
+        .filter(|v| v.has_v1())
+        .and_then(|_| get_cgroup_v1_mount_points("/proc/mounts").ok())
+    {
+        sources.push(Box::new(CgroupV1Source::with_filesystem_reader_at(
+            v1_mount_points,
+        )));
+    }
+    sources.push(Box::new(ProcSource::with_filesystem_reader_at(
+        PathBuf::from("/proc"),
+    )));
+    sources
+}

src/stats/cgroup_v1/mod.rs

@@ -3,13 +3,11 @@ mod cpu_usage;
 mod memory_current;
 mod memory_max;
 mod num_cpus;
-use crate::utils::read_first_line;
-use std::fs::File;
-use std::io::{self, BufRead, BufReader};
-use std::path::PathBuf;
-
+use crate::utils::{get_path_or_croak, read_all_lines, read_first_line};
 #[cfg(test)]
 use mockall::automock;
+use std::io::{self};
+use std::path::PathBuf;
 
 #[derive(Default, Clone)]
 pub struct CgroupV1MountPoints {
@@ -65,30 +63,6 @@ impl CgroupV1MountPoints {
         self.memory_stat_path = memory.as_ref().map(|pb| pb.join("memory.stat"));
         self.memory = memory;
     }
-
-    pub fn cpu_quota_path(&self) -> Option<PathBuf> {
-        self.cpu_quota_path.clone()
-    }
-
-    pub fn cpu_period_path(&self) -> Option<PathBuf> {
-        self.cpu_period_path.clone()
-    }
-
-    pub fn cpu_usage_path(&self) -> Option<PathBuf> {
-        self.cpu_usage_path.clone()
-    }
-
-    pub fn memory_usage_path(&self) -> Option<PathBuf> {
-        self.memory_usage_path.clone()
-    }
-
-    pub fn memory_limit_path(&self) -> Option<PathBuf> {
-        self.memory_limit_path.clone()
-    }
-
-    pub fn memory_stat_path(&self) -> Option<PathBuf> {
-        self.memory_stat_path.clone()
-    }
 }
 
 pub struct CgroupV1Source<P: CgroupV1Provider> {
@@ -133,30 +107,6 @@ impl CgroupV1FilesystemReader {
     fn new(mount_points: CgroupV1MountPoints) -> Self {
         Self { mount_points }
     }
-
-    fn cpu_quota_path(&self) -> Option<PathBuf> {
-        self.mount_points.cpu_quota_path()
-    }
-
-    fn cpu_period_path(&self) -> Option<PathBuf> {
-        self.mount_points.cpu_period_path()
-    }
-
-    fn cpu_usage(&self) -> Option<PathBuf> {
-        self.mount_points.cpu_usage_path()
-    }
-
-    fn memory_usage_path(&self) -> Option<PathBuf> {
-        self.mount_points.memory_usage_path()
-    }
-
-    fn memory_limit_path(&self) -> Option<PathBuf> {
-        self.mount_points.memory_limit_path()
-    }
-
-    fn memory_stat_path(&self) -> Option<PathBuf> {
-        self.mount_points.memory_stat_path()
-    }
 }
 
 #[cfg_attr(test, automock)]
@@ -171,69 +121,44 @@ pub trait CgroupV1Provider {
 
 impl CgroupV1Provider for CgroupV1FilesystemReader {
     fn get_cgroup_v1_cpu_cfs_quota(&self) -> io::Result<String> {
-        let Some(file_path) = self.cpu_quota_path() else {
-            return Err(io::Error::new(
-                io::ErrorKind::NotFound,
-                "cpu.cfs_quota_us file not found",
-            ));
-        };
-
-        read_first_line(file_path)
+        read_first_line(get_path_or_croak(
+            &self.mount_points.cpu_quota_path,
+            "cpu.cfs_quota_us",
+        )?)
     }
 
     fn get_cgroup_v1_cpu_cfs_period(&self) -> io::Result<String> {
-        let Some(file_path) = self.cpu_period_path() else {
-            return Err(io::Error::new(
-                io::ErrorKind::NotFound,
-                "cpu.cfs_period_us file not found",
-            ));
-        };
-
-        read_first_line(file_path)
+        read_first_line(get_path_or_croak(
+            &self.mount_points.cpu_period_path,
+            "cpu.cfs_period_us",
+        )?)
     }
 
     fn get_cgroup_v1_cpuacct_usage(&self) -> io::Result<String> {
-        let Some(file_path) = self.cpu_usage() else {
-            return Err(io::Error::new(
-                io::ErrorKind::NotFound,
-                "cpuacct.usage file not found",
-            ));
-        };
-
-        read_first_line(file_path)
+        read_first_line(get_path_or_croak(
+            &self.mount_points.cpu_usage_path,
+            "cpuacct.usage",
+        )?)
     }
 
     fn get_cgroup_v1_memory_usage_in_bytes(&self) -> io::Result<String> {
-        let Some(file_path) = self.memory_usage_path() else {
-            return Err(io::Error::new(
-                io::ErrorKind::NotFound,
-                "memory.usage_in_bytes file not found",
-            ));
-        };
-
-        read_first_line(file_path)
+        read_first_line(get_path_or_croak(
+            &self.mount_points.memory_usage_path,
+            "memory.usage_in_bytes",
+        )?)
     }
 
     fn get_cgroup_v1_memory_limit_in_bytes(&self) -> io::Result<String> {
-        let Some(file_path) = self.memory_limit_path() else {
-            return Err(io::Error::new(
-                io::ErrorKind::NotFound,
-                "memory.limit_in_bytes file not found",
-            ));
-        };
-
-        read_first_line(file_path)
+        read_first_line(get_path_or_croak(
+            &self.mount_points.memory_limit_path,
+            "memory.limit_in_bytes",
+        )?)
     }
 
     fn get_cgroup_v1_memory_stat(&self) -> io::Result<Vec<String>> {
-        let Some(file_path) = self.memory_stat_path() else {
-            return Err(io::Error::new(
-                io::ErrorKind::NotFound,
-                "memory.stat file not found",
-            ));
-        };
-
-        let file = File::open(file_path)?;
-        BufReader::new(file).lines().collect()
+        read_all_lines(get_path_or_croak(
+            &self.mount_points.memory_stat_path,
+            "memory.stat",
+        )?)
     }
 }

src/stats/cgroup_v2/mod.rs

@@ -3,9 +3,8 @@ mod cpu_usage;
 mod memory_current;
 mod memory_max;
 mod num_cpus;
-use crate::utils::read_first_line;
-use std::fs::File;
-use std::io::{self, BufRead, BufReader};
+use crate::utils::{read_all_lines, read_first_line};
+use std::io::{self};
 use std::path::PathBuf;
 
 #[cfg(test)]
@@ -46,47 +45,38 @@ impl<P: CgroupV2Provider> SystemStatsSource for CgroupV2Source<P> {
 }
 
 pub struct CgroupV2FilesystemReader {
-    cgroup_v2_path: PathBuf,
+    cpu_max_path: PathBuf,
+    cpu_stat_path: PathBuf,
+    mem_current_path: PathBuf,
+    mem_max_path: PathBuf,
 }
 
 impl CgroupV2FilesystemReader {
     fn new(cgroup_v2_path: PathBuf) -> Self {
-        Self { cgroup_v2_path }
-    }
-
-    fn cpu_max_path(&self) -> PathBuf {
-        self.cgroup_v2_path.join("cpu.max")
-    }
-
-    fn cpu_stat_path(&self) -> PathBuf {
-        self.cgroup_v2_path.join("cpu.stat")
-    }
-
-    fn mem_current_path(&self) -> PathBuf {
-        self.cgroup_v2_path.join("memory.current")
-    }
-
-    fn mem_max_path(&self) -> PathBuf {
-        self.cgroup_v2_path.join("memory.max")
+        Self {
+            cpu_max_path: cgroup_v2_path.join("cpu.max"),
+            cpu_stat_path: cgroup_v2_path.join("cpu.stat"),
+            mem_current_path: cgroup_v2_path.join("memory.current"),
+            mem_max_path: cgroup_v2_path.join("memory.max"),
+        }
     }
 }
 
 impl CgroupV2Provider for CgroupV2FilesystemReader {
     fn get_cgroup_v2_cpu_stat(&self) -> io::Result<Vec<String>> {
-        let file = File::open(self.cpu_stat_path())?;
-        BufReader::new(file).lines().collect()
+        read_all_lines(&self.cpu_stat_path)
     }
 
     fn get_cgroup_v2_cpu_max(&self) -> io::Result<String> {
-        read_first_line(self.cpu_max_path())
+        read_first_line(&self.cpu_max_path)
     }
 
     fn get_cgroup_v2_memory_current(&self) -> io::Result<String> {
-        read_first_line(self.mem_current_path())
+        read_first_line(&self.mem_current_path)
     }
 
     fn get_cgroup_v2_memory_max(&self) -> io::Result<String> {
-        read_first_line(self.mem_max_path())
+        read_first_line(&self.mem_max_path)
     }
 }