feat: add disk stats collector

This allows collection of I/O stats (bytes read/written) for all disks

Author: Leandros

examples/config.kdl

@@ -12,4 +12,7 @@ metrics {
     mountpoints "/"
   }
   pressure enabled=#true
+  disk_stats enabled=#true {
+      mountpoints "/"
+  }
 }

src/collector.rs

@@ -12,8 +12,7 @@ use smol::lock::RwLock;
 
 use crate::config::UserConfig;
 use crate::metrics::collector::{
-    CpuStatsCollector, FilesystemStatsCollector, MemoryStatsCollector, NetworkStatsCollector,
-    NodeInfoCollector, PressureCollector, SystemdUnitStateCollector,
+    CpuStatsCollector, DiskStatsCollector, FilesystemStatsCollector, MemoryStatsCollector, NetworkStatsCollector, NodeInfoCollector, PressureCollector, SystemdUnitStateCollector
 };
 use crate::metrics::Metric;
 
@@ -99,6 +98,13 @@ impl Collector {
             inner.metrics.push(collector);
         }
 
+        if metrics.disk_stats.enabled {
+            let mountpoints = metrics.disk_stats.mountpoints.join(",");
+            let options = HashMap::from([("mountpoints".to_owned(), mountpoints)]);
+            let collector = Box::new(DiskStatsCollector::new(options)?);
+            inner.metrics.push(collector);
+        }
+
         Ok(())
     }
 

src/config.rs

@@ -26,6 +26,7 @@ pub struct MetricsConfig {
     pub network_throughput: NetworkThroughputConfig,
     pub disk_usage: DiskUsageConfig,
     pub pressure: PressureConfig,
+    pub disk_stats: DiskStatConfig,
 }
 
 #[derive(Debug)]
@@ -67,6 +68,12 @@ pub struct PressureConfig {
     pub enabled: bool,
 }
 
+#[derive(Debug)]
+pub struct DiskStatConfig {
+    pub enabled: bool,
+    pub mountpoints: Vec<String>,
+}
+
 impl Default for MetricsConfig {
     fn default() -> Self {
         Self {
@@ -89,6 +96,10 @@ impl Default for MetricsConfig {
                 mountpoints: vec![],
             },
             pressure: PressureConfig { enabled: true },
+            disk_stats: DiskStatConfig {
+                enabled: false,
+                mountpoints: vec![],
+            },
         }
     }
 }
@@ -205,6 +216,28 @@ impl UserConfig {
                         .unwrap_or_default();
                     ret.pressure = PressureConfig { enabled };
                 }
+
+                if let Some(node) = children.get("disk_stats") {
+                    let enabled = node
+                        .get("enabled")
+                        .and_then(|el| el.as_bool())
+                        .unwrap_or_default();
+                    let mountpoints = node
+                        .children()
+                        .and_then(|el| el.get("mountpoints"))
+                        .map(|el| el.entries())
+                        .map(|it| {
+                            it.iter()
+                                .filter_map(|el| el.value().as_string())
+                                .map(ToOwned::to_owned)
+                                .collect::<Vec<_>>()
+                        })
+                        .unwrap_or_default();
+                    ret.disk_stats = DiskStatConfig {
+                        enabled,
+                        mountpoints,
+                    };
+                }
             }
 
             ret

src/metrics/collector.rs

@@ -11,6 +11,7 @@ use prometheus_client::metrics::gauge::Gauge;
 use smol::lock::Mutex;
 
 use super::cpu::{CpuUsage, LoadAverages};
+use super::disk::IOMetrics;
 use super::fs::FilesystemUsage;
 use super::info::NodeInfo;
 use super::memory::MemoryStats;
@@ -493,3 +494,80 @@ impl Metric for PressureCollector {
         })
     }
 }
+
+#[derive(Debug)]
+pub struct DiskStatsCollector {
+    bytes_written: Family<DiskStatsLabels, Gauge<u64, AtomicU64>>,
+    bytes_read: Family<DiskStatsLabels, Gauge<u64, AtomicU64>>,
+    mountpoints: Vec<String>,
+}
+
+#[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)]
+struct DiskStatsLabels {
+    /// Device name
+    device: String,
+    /// Mountpoint
+    mountpoint: String,
+}
+
+impl DiskStatsCollector {
+    pub fn new(options: hashbrown::HashMap<String, String>) -> Result<Self> {
+        let mountpoints = if let Some(mountpoints_str) = options.get("mountpoints") {
+            mountpoints_str
+                .split(',')
+                .map(|s| s.trim().to_string())
+                .collect()
+        } else {
+            vec!["/".to_string()] // Default to root filesystem
+        };
+
+        Ok(Self {
+            bytes_written: Default::default(),
+            bytes_read: Default::default(),
+            mountpoints,
+        })
+    }
+}
+
+impl Metric for DiskStatsCollector {
+    fn register(&self, registry: &mut prometheus_client::registry::Registry) {
+        registry.register(
+            "litemon_disk_bytes_written_total",
+            "Number of bytes written to disk since boot",
+            self.bytes_written.clone(),
+        );
+        registry.register(
+            "litemon_disk_bytes_read_total",
+            "Number of bytes read from disk since boot",
+            self.bytes_read.clone(),
+        );
+    }
+
+    fn collect(&self) -> DynFuture<'_, Result<()>> {
+        Box::pin(async move {
+            let stats = IOMetrics::all().await?;
+            for (device, stats) in &stats.disks {
+                if !self.mountpoints.iter().any(|el| el == &stats.mountpoint) {
+                    continue;
+                }
+
+                let labels = DiskStatsLabels {
+                    device: device.clone(),
+                    mountpoint: stats.mountpoint.clone(),
+                };
+
+                let bytes_read_prev = self.bytes_read.get_or_create(&labels).get();
+                self.bytes_read
+                    .get_or_create(&labels)
+                    .inc_by(stats.bytes_read_total - bytes_read_prev);
+
+                let bytes_written_prev = self.bytes_written.get_or_create(&labels).get();
+                self.bytes_written
+                    .get_or_create(&labels)
+                    .inc_by(stats.bytes_written_total - bytes_written_prev);
+            }
+
+            Ok(())
+        })
+    }
+}

src/metrics/disk.rs

@@ -0,0 +1,82 @@
+//! Metrics about bytes read/written to disks.
+
+use anyhow::{Context, Result};
+use hashbrown::HashMap;
+
+/// Metrics about the bytes read/written to each attached disk.
+#[derive(Debug)]
+pub struct IOMetrics {
+    /// Each disk is one entry.
+    pub disks: HashMap<String, DiskMetrics>,
+}
+
+/// Metrics about a single disk.
+#[derive(Debug)]
+pub struct DiskMetrics {
+    /// Bytes read since boot.
+    pub bytes_read_total: u64,
+    /// Bytes written since boot.
+    pub bytes_written_total: u64,
+    /// Mountpoint
+    pub mountpoint: String,
+}
+
+impl IOMetrics {
+    /// Retrieve disk metrics for all attached disks.
+    pub async fn all() -> Result<Self> {
+        let stats =
+            smol::unblock(|| procfs::diskstats().context("reading /proc/diskstats")).await?;
+        // We try to only collect statistics about disks which are mounted.
+        let mounts = smol::unblock(|| procfs::mounts().context("reading /proc/mounts")).await?;
+
+        let mut ret = IOMetrics {
+            disks: HashMap::new(),
+        };
+
+        for mount in &mounts {
+            let device = &mount.fs_spec;
+            let mount_point = &mount.fs_file;
+
+            // Skip non-physical mounts.
+            let Some(device_name) = device.strip_prefix("/dev/") else {
+                continue;
+            };
+
+            // Skip pseudo filesystems
+            if mount_point.starts_with("/proc")
+                || mount_point.starts_with("/sys")
+                || mount_point.starts_with("/dev")
+                || mount_point.starts_with("/run")
+            {
+                continue;
+            }
+
+            let Some(stat) = stats.iter().find(|el| el.name == device_name) else {
+                continue;
+            };
+
+            let Some(root_device) = stats.iter().find(|el| device_name.starts_with(&el.name))
+            else {
+                continue;
+            };
+
+            let sector_size_str =
+                smol::fs::read_to_string(format!("/sys/block/{}/queue/hw_sector_size", root_device.name))
+                    .await
+                    .context("reading sector size")?;
+            let sector_size = sector_size_str
+                .trim()
+                .parse::<u64>()
+                .context("parsing sector size")?;
+            let metrics = DiskMetrics {
+                bytes_read_total: stat.sectors_read * sector_size,
+                bytes_written_total: stat.sectors_written * sector_size,
+                mountpoint: mount_point.clone(),
+            };
+
+            ret.disks.insert(stat.name.clone(), metrics);
+        }
+
+        Ok(ret)
+    }
+}

src/metrics/mod.rs

@@ -4,6 +4,7 @@ use anyhow::Result;
 
 pub mod collector;
 pub mod cpu;
+pub mod disk;
 pub mod fs;
 pub mod info;
 pub mod memory;