disk sensors

src/disk/mod.rs

@@ -1,35 +1,34 @@
-use crate::{Error, Result, SensorData};
+use crate::{Error, MultiSensorSource, Result, SensorData};
 use ahash::{AHashSet, AHasher};
-use once_cell::sync::Lazy;
 use regex::Regex;
 use std::ffi::CString;
 use std::fmt::Write;
 use std::fs::File;
 use std::hash::{Hash, Hasher};
-use std::io::{BufRead, BufReader};
+use std::io::{Read, Seek};
 use std::mem::MaybeUninit;
 
 pub mod zfs;
 
 #[derive(Debug, Clone, Default)]
-pub struct IoStats {
+pub struct DiskStats {
     pub interface: String,
     pub bytes_sent: u64,
     pub bytes_received: u64,
 }
 
-impl SensorData for IoStats {
+impl SensorData for DiskStats {
     fn write<W: Write>(&self, mut w: W, hostname: &str) {
         if self.bytes_received > 0 || self.bytes_sent > 0 {
             writeln!(
                 &mut w,
-                "net_sent{{host=\"{}\", network=\"{}\"}} {}",
+                "disk_sent{{host=\"{}\", disk=\"{}\"}} {}",
                 hostname, self.interface, self.bytes_sent
             )
             .ok();
             writeln!(
                 &mut w,
-                "net_received{{host=\"{}\", network=\"{}\"}} {}",
+                "disk_received{{host=\"{}\", disk=\"{}\"}} {}",
                 hostname, self.interface, self.bytes_received
             )
             .ok();
@@ -37,6 +36,70 @@ impl SensorData for IoStats {
     }
 }
 
+pub struct DiskStatSource {
+    source: File,
+    buff: String,
+    regex: Regex,
+}
+
+impl DiskStatSource {
+    pub fn new() -> Result<DiskStatSource> {
+        Ok(DiskStatSource {
+            source: File::open("/proc/diskstats")?,
+            buff: String::new(),
+            regex: Regex::new(r" ([sv]d[a-z]+|nvme[0-9]n[0-9]|mmcblk[0-9]) ").unwrap(),
+        })
+    }
+}
+
+impl MultiSensorSource for DiskStatSource {
+    type Data = DiskStats;
+    type Iter<'a> = DiskStatParser<'a>;
+
+    fn read(&mut self) -> Result<Self::Iter<'_>> {
+        self.buff.clear();
+        self.source.rewind()?;
+        self.source.read_to_string(&mut self.buff)?;
+
+        Ok(DiskStatParser {
+            lines: self.buff.lines(),
+            regex: &self.regex,
+        })
+    }
+}
+
+pub struct DiskStatParser<'a> {
+    lines: std::str::Lines<'a>,
+    regex: &'a Regex,
+}
+
+impl Iterator for DiskStatParser<'_> {
+    type Item = Result<DiskStats>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let line = loop {
+            let line = self.lines.next()?;
+            if self.regex.is_match(line) {
+                break line;
+            }
+        };
+        let mut parts = line.split_whitespace().skip(2);
+        let name: String = parts.next()?.into();
+        let _read_count = parts.next();
+        let _read_merged_count = parts.next();
+        let read_sectors = parts.next()?.parse::<u64>().ok()?;
+        let mut parts = parts.skip(1);
+        let _write_count = parts.next();
+        let _write_merged_count = parts.next();
+        let write_sectors = parts.next()?.parse::<u64>().ok()?;
+        Some(Ok(DiskStats {
+            interface: name,
+            bytes_sent: write_sectors * 512,
+            bytes_received: read_sectors * 512,
+        }))
+    }
+}
+
 #[derive(Clone, Debug)]
 pub struct DiskUsage {
     pub name: String,
@@ -44,56 +107,87 @@ pub struct DiskUsage {
     pub free: u64,
 }
 
-pub fn disk_stats() -> Result<impl Iterator<Item = IoStats>> {
+impl SensorData for DiskUsage {
-    static DISK_REGEX: Lazy<Regex> =
+    fn write<W: Write>(&self, mut w: W, hostname: &str) {
-        Lazy::new(|| Regex::new(r" ([sv]d[a-z]+|nvme[0-9]n[0-9]|mmcblk[0-9]) ").unwrap());
+        if self.size > 0 {
-
+            writeln!(
-    let stat = BufReader::new(File::open("/proc/diskstats")?);
+                &mut w,
-    Ok(stat
+                "disk_size{{host=\"{}\", disk=\"{}\"}} {}",
-        .lines()
+                hostname, self.name, self.size
-        .filter_map(Result::ok)
+            )
-        .filter(|line| DISK_REGEX.is_match(line))
+            .ok();
-        .filter_map(|line: String| {
+            writeln!(
-            let mut parts = line.split_whitespace().skip(2);
+                &mut w,
-            let name: String = parts.next()?.into();
+                "disk_free{{host=\"{}\", disk=\"{}\"}} {}",
-            let _read_count = parts.next();
+                hostname, self.name, self.free
-            let _read_merged_count = parts.next();
+            )
-            let read_sectors = parts.next()?.parse::<u64>().ok()?;
+            .ok();
-            let mut parts = parts.skip(1);
+        }
-            let _write_count = parts.next();
+    }
-            let _write_merged_count = parts.next();
-            let write_sectors = parts.next()?.parse::<u64>().ok()?;
-            Some(IoStats {
-                interface: name,
-                bytes_sent: write_sectors * 512,
-                bytes_received: read_sectors * 512,
-            })
-        }))
 }
 
-pub fn disk_usage() -> Result<impl Iterator<Item = DiskUsage>> {
+pub struct DiskUsageSource {
-    let stat = BufReader::new(File::open("/proc/mounts")?);
+    source: File,
-    let mut found_disks = AHashSet::with_capacity(8);
+    buff: String,
-    Ok(stat
+}
-        .lines()
+
-        .filter_map(Result::ok)
+impl DiskUsageSource {
-        .filter(|line| line.starts_with('/'))
+    pub fn new() -> Result<DiskUsageSource> {
-        .filter(|line| !line.contains("/dev/loop"))
+        Ok(DiskUsageSource {
-        .filter(|line| !line.contains("fuse"))
+            source: File::open("/proc/mounts")?,
-        .filter_map(move |line: String| {
+            buff: String::new(),
-            let mut parts = line.split_ascii_whitespace();
+        })
-            let disk = parts.next()?;
+    }
-            if !found_disks.insert(hash_str(disk)) {
+}
-                return None;
+
+impl MultiSensorSource for DiskUsageSource {
+    type Data = DiskUsage;
+    type Iter<'a> = DiskUsageParser<'a>;
+
+    fn read(&mut self) -> Result<Self::Iter<'_>> {
+        self.buff.clear();
+        self.source.rewind()?;
+        self.source.read_to_string(&mut self.buff)?;
+
+        Ok(DiskUsageParser {
+            lines: self.buff.lines(),
+            found_disks: AHashSet::with_capacity(16),
+        })
+    }
+}
+
+pub struct DiskUsageParser<'a> {
+    lines: std::str::Lines<'a>,
+    found_disks: AHashSet<u64>,
+}
+
+impl Iterator for DiskUsageParser<'_> {
+    type Item = Result<DiskUsage>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let line = loop {
+            let line = self.lines.next()?;
+            if line.starts_with('/') && !line.contains("/dev/loop") && !line.contains("fuse") {
+                break line;
             }
-            let mount_point = parts.next()?;
+        };
-            let stat = statvfs(&mount_point).ok()?;
+
-            Some(DiskUsage {
+        let mut parts = line.split_ascii_whitespace();
-                name: mount_point.to_string(),
+        let disk = parts.next()?;
-                size: stat.f_blocks * stat.f_frsize as u64,
+        if !self.found_disks.insert(hash_str(disk)) {
-                free: stat.f_bavail * stat.f_frsize as u64,
+            return None;
-            })
+        }
+        let mount_point = parts.next()?;
+        let stat = match statvfs(&mount_point) {
+            Ok(stat) => stat,
+            Err(e) => return Some(Err(e)),
+        };
+        Some(Ok(DiskUsage {
+            name: mount_point.to_string(),
+            size: stat.f_blocks * stat.f_frsize as u64,
+            free: stat.f_bavail * stat.f_frsize as u64,
         }))
+    }
 }
 
 fn statvfs(path: &str) -> Result<libc::statvfs> {

src/lib.rs

@@ -5,7 +5,6 @@ pub mod hwmon;
 pub mod power;
 pub mod sensors;
 
-use crate::disk::disk_usage;
 use crate::disk::zfs::pools;
 use crate::disk::*;
 use crate::sensors::*;
@@ -51,6 +50,8 @@ pub struct Sensors {
     temp: Mutex<TemperatureSource>,
     net: Mutex<NetworkSource>,
     mem: Mutex<MemorySource>,
+    disk_stats: Mutex<DiskStatSource>,
+    disk_usage: Mutex<DiskUsageSource>,
 }
 
 impl Sensors {
@@ -61,14 +62,18 @@ impl Sensors {
             temp: Mutex::new(TemperatureSource::new()?),
             net: Mutex::new(NetworkSource::new()?),
             mem: Mutex::new(MemorySource::new()?),
+            disk_stats: Mutex::new(DiskStatSource::new()?),
+            disk_usage: Mutex::new(DiskUsageSource::new()?),
         })
     }
 }
 
 pub fn get_metrics(sensors: &Sensors) -> Result<String> {
     let hostname = &sensors.hostname;
-    let disk_usage = disk_usage()?;
+    let mut disk_source = sensors.disk_stats.lock().unwrap();
-    let disks = disk_stats()?;
+    let mut disk_usage_source = sensors.disk_usage.lock().unwrap();
+    let disks = disk_source.read()?;
+    let disk_usage = disk_usage_source.read()?;
     let cpu = sensors.cpu.lock().unwrap().read()?;
     let memory = sensors.mem.lock().unwrap().read()?;
     let temperatures = sensors.temp.lock().unwrap().read()?;
@@ -100,36 +105,14 @@ pub fn get_metrics(sensors: &Sensors) -> Result<String> {
         }
     }
     for disk in disks {
-        if disk.bytes_received > 0 && disk.bytes_sent > 0 {
+        if let Ok(disk) = disk {
-            writeln!(
+            disk.write(&mut result, hostname);
-                &mut result,
-                "disk_sent{{host=\"{}\", disk=\"{}\"}} {}",
-                hostname, disk.interface, disk.bytes_sent
-            )
-            .ok();
-            writeln!(
-                &mut result,
-                "disk_received{{host=\"{}\", disk=\"{}\"}} {}",
-                hostname, disk.interface, disk.bytes_received
-            )
-            .ok();
         }
     }
 
     for disk in disk_usage {
-        if disk.size > 0 {
+        if let Ok(disk) = disk {
-            writeln!(
+            disk.write(&mut result, hostname);
-                &mut result,
-                "disk_size{{host=\"{}\", disk=\"{}\"}} {}",
-                hostname, disk.name, disk.size
-            )
-            .ok();
-            writeln!(
-                &mut result,
-                "disk_free{{host=\"{}\", disk=\"{}\"}} {}",
-                hostname, disk.name, disk.free
-            )
-            .ok();
         }
     }
     for (label, temp) in temperatures {

src/sensors.rs

@@ -1,4 +1,3 @@
-use crate::disk::IoStats;
 use crate::hwmon::{Device, FileSource};
 use crate::{Error, MultiSensorSource, Result, SensorData, SensorSource};
 use std::array::IntoIter;
@@ -217,6 +216,32 @@ impl SensorSource for CpuTimeSource {
     }
 }
 
+#[derive(Debug, Clone, Default)]
+pub struct NetStats {
+    pub interface: String,
+    pub bytes_sent: u64,
+    pub bytes_received: u64,
+}
+
+impl SensorData for NetStats {
+    fn write<W: Write>(&self, mut w: W, hostname: &str) {
+        if self.bytes_received > 0 || self.bytes_sent > 0 {
+            writeln!(
+                &mut w,
+                "net_sent{{host=\"{}\", network=\"{}\"}} {}",
+                hostname, self.interface, self.bytes_sent
+            )
+            .ok();
+            writeln!(
+                &mut w,
+                "net_received{{host=\"{}\", network=\"{}\"}} {}",
+                hostname, self.interface, self.bytes_received
+            )
+            .ok();
+        }
+    }
+}
+
 pub struct NetworkSource {
     source: File,
     buff: String,
@@ -230,7 +255,7 @@ impl NetworkSource {
         })
     }
 
-    fn parse_line(line: &str) -> Result<IoStats> {
+    fn parse_line(line: &str) -> Result<NetStats> {
         let mut parts = line.trim_start().split_ascii_whitespace();
         if let (
             Some(interface),
@@ -255,7 +280,7 @@ impl NetworkSource {
             parts.next(),
             parts.next(),
         ) {
-            Ok(IoStats {
+            Ok(NetStats {
                 interface: interface.trim_end_matches(':').into(),
                 bytes_sent: bytes_sent.parse()?,
                 bytes_received: bytes_received.parse()?,
@@ -267,7 +292,7 @@ impl NetworkSource {
 }
 
 impl MultiSensorSource for NetworkSource {
-    type Data = IoStats;
+    type Data = NetStats;
     type Iter<'a> = NetworkStatParser<'a>;
 
     fn read(&mut self) -> Result<Self::Iter<'_>> {
@@ -286,7 +311,7 @@ pub struct NetworkStatParser<'a> {
 }
 
 impl<'a> Iterator for NetworkStatParser<'a> {
-    type Item = Result<IoStats>;
+    type Item = Result<NetStats>;
 
     fn next(&mut self) -> Option<Self::Item> {
         let line = loop {