disks

Cargo.lock

@@ -718,6 +718,7 @@ checksum = "b8a653442b9bdd11a77d3753a60443c60c4437d3acac8e6c3d4a6a9acd7cceed"
 dependencies = [
  "heim-common",
  "heim-cpu",
+ "heim-disk",
  "heim-host",
  "heim-memory",
  "heim-net",
@@ -763,6 +764,23 @@ dependencies = [
  "winapi",
 ]
 
+[[package]]
+name = "heim-disk"
+version = "0.1.0-rc.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "75603ff3868851c04954ee86bf610a6bd45be2732a0e81c35fd72b2b90fa4718"
+dependencies = [
+ "bitflags",
+ "cfg-if 1.0.0",
+ "core-foundation",
+ "heim-common",
+ "heim-runtime",
+ "libc",
+ "mach",
+ "widestring",
+ "winapi",
+]
+
 [[package]]
 name = "heim-host"
 version = "0.1.0-rc.1"
@@ -1245,7 +1263,9 @@ dependencies = [
  "futures-util",
  "heim",
  "libzetta",
+ "once_cell",
  "parse-display",
+ "regex",
  "tokio",
  "warp",
 ]
@@ -2160,6 +2180,12 @@ dependencies = [
  "cc",
 ]
 
+[[package]]
+name = "widestring"
+version = "0.4.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c168940144dd21fd8046987c16a46a33d5fc84eec29ef9dcddc2ac9e31526b7c"
+
 [[package]]
 name = "winapi"
 version = "0.3.9"

Cargo.toml

@@ -7,11 +7,13 @@ edition = "2018"
 [dependencies]
 libzetta = "0.2"
 color-eyre = "0.5"
-heim = { version = "0.1.0-rc.1", features = ["sensors", "cpu", "memory", "net", "host"] }
+heim = { version = "0.1.0-rc.1", features = ["sensors", "cpu", "memory", "net", "host", "disk"] }
 warp = "0.3"
 tokio = { version = "1", features = ["macros", "rt-multi-thread"] }
 ctrlc = { version = "3", features = ["termination"] }
 dotenv = "0.15"
 futures-util = "0.3"
 futures-lite = "1"
-parse-display = "0.4"
+parse-display = "0.4"
+regex = "1"
+once_cell = "1"

src/heim.rs

@@ -4,7 +4,9 @@ use futures_util::future;
 use futures_util::stream::{Stream, StreamExt};
 use heim::sensors::TemperatureSensor;
 use heim::units::{information, ratio, thermodynamic_temperature};
+use once_cell::sync::{Lazy, OnceCell};
 use parse_display::Display;
+use regex::Regex;
 use std::collections::HashMap;
 use std::time::Duration;
 use tokio::time::sleep;
@@ -23,7 +25,7 @@ pub struct Memory {
 }
 
 #[derive(Debug, Clone, Default)]
-pub struct NetworkStats {
+pub struct IOStats {
     pub interface: String,
     pub bytes_sent: u64,
     pub bytes_received: u64,
@@ -74,12 +76,12 @@ impl Heim {
         Ok((measurement_2 - measurement_1).get::<ratio::percent>() / cores as f32)
     }
 
-    pub async fn network_stats(&self) -> Result<impl Stream<Item = NetworkStats>> {
+    pub async fn network_stats(&self) -> Result<impl Stream<Item = IOStats>> {
         let networks = heim::net::io_counters().await?;
         Ok(networks
             .filter_map(|network| future::ready(network.ok()))
             .filter(|network| future::ready(network.interface().starts_with("enp")))
-            .map(|network| NetworkStats {
+            .map(|network| IOStats {
                 interface: network.interface().into(),
                 bytes_sent: network.bytes_sent().get::<information::byte>(),
                 bytes_received: network.bytes_recv().get::<information::byte>(),
@@ -89,4 +91,26 @@ impl Heim {
     pub async fn hostname(&self) -> Result<String> {
         Ok(heim::host::platform().await?.hostname().to_string())
     }
+
+    pub async fn disk_stats(&self) -> Result<impl Stream<Item = IOStats>> {
+        static DISK_REGEX: Lazy<Regex> =
+            Lazy::new(|| Regex::new(r"^([sv]d[a-z]+|nvme\dn\d)$").unwrap());
+        let disks = heim::disk::io_counters().await?;
+        Ok(disks
+            .filter_map(|disk| future::ready(disk.ok()))
+            .filter_map(|disk| {
+                future::ready(
+                    disk.device_name()
+                        .to_str()
+                        .map(str::to_string)
+                        .map(|name| (disk, name)),
+                )
+            })
+            .filter(|(_disk, name)| future::ready(DISK_REGEX.is_match(&name)))
+            .map(|(disk, name)| IOStats {
+                interface: name,
+                bytes_sent: disk.write_bytes().get::<information::byte>(),
+                bytes_received: disk.read_bytes().get::<information::byte>(),
+            }))
+    }
 }

src/main.rs

@@ -1,7 +1,7 @@
 mod heim;
 mod zfs;
 
-use crate::heim::{Heim, Memory, NetworkStats, TemperatureLabel};
+use crate::heim::{Heim, IOStats, Memory, TemperatureLabel};
 use crate::zfs::{ZfsPool, ZFS};
 use color_eyre::{Report, Result};
 use futures_util::stream::StreamExt;
@@ -24,13 +24,14 @@ impl From<Report> for ReportRejection {
 impl Reject for ReportRejection {}
 
 async fn get_metrics(heim: Heim, zfs: ZFS) -> Result<String, ReportRejection> {
-    let (hostname, pools, cpu, memory, network, temperatures): (
+    let (hostname, pools, cpu, memory, network, temperatures, disks): (
         String,
         Vec<ZfsPool>,
         f32,
         Memory,
         _,
         HashMap<TemperatureLabel, f32>,
+        _,
     ) = try_join! {
         heim.hostname(),
         zfs.pools(),
@@ -38,8 +39,10 @@ async fn get_metrics(heim: Heim, zfs: ZFS) -> Result<String, ReportRejection> {
         heim.memory(),
         heim.network_stats(),
         heim.temperatures(),
+        heim.disk_stats(),
     }?;
     pin_mut!(network);
+    pin_mut!(disks);
     let mut result = String::with_capacity(256);
     writeln!(&mut result, "cpu_usage{{host=\"{}\"}} {:.1}", hostname, cpu).ok();
     writeln!(
@@ -75,8 +78,8 @@ async fn get_metrics(heim: Heim, zfs: ZFS) -> Result<String, ReportRejection> {
         .ok();
     }
     while let Some(network) = network.next().await {
+        let network: IOStats = network;
         if network.bytes_received > 0 || network.bytes_sent > 0 {
-            let network: NetworkStats = network;
             writeln!(
                 &mut result,
                 "net_sent{{host=\"{}\", network=\"{}\"}} {}",
@@ -91,6 +94,23 @@ async fn get_metrics(heim: Heim, zfs: ZFS) -> Result<String, ReportRejection> {
             .ok();
         }
     }
+    while let Some(disk) = disks.next().await {
+        let disk: IOStats = disk;
+        if disk.bytes_received > 0 && disk.bytes_sent > 0 {
+            writeln!(
+                &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 (label, temp) in temperatures {
         writeln!(
             &mut result,