Add disk R/W stats

src/bin/home.rs

@@ -3,6 +3,7 @@ use i3monkit::{ColorRGB, Header, I3Protocol, WidgetCollection};
 use i3xs::widgets::{
     cpu::CpuWidget,
     datetime::{DateTimeWidget, TimeColor},
+    disk::AllDiskSpeedWidget,
     network::AllNetworkSpeedWidget,
     power::PowerSupply,
 };
@@ -20,9 +21,12 @@ fn main() {
     bar.push(PowerSupply::default());
     bar.push(CpuWidget::default());
 
-    // Realtime upload/download rate for a interface
+    // Realtime upload/download rate for a any active NICs
     bar.push(AllNetworkSpeedWidget::new(6));
 
+    // Realtime read/write rate for a any active disks
+    bar.push(AllDiskSpeedWidget::new(6));
+
     let mut dt = DateTimeWidget::new("%m/%d %H:%M");
     dt.set_colors(vec![
         TimeColor {

src/bin/work.rs

@@ -23,6 +23,9 @@ fn main() {
     // Realtime upload/download rate for a interface
     bar.push(AllNetworkSpeedWidget::new(6));
 
+    // Realtime read/write rate for a any active disks
+    bar.push(AllDiskSpeedWidget::new(6));
+
     let mut dt = DateTimeWidget::tz("%H:%M %Z", chrono_tz::Europe::London);
     dt.set_colors(vec![
         TimeColor {

src/widgets/disk.rs

@@ -0,0 +1,206 @@
+use std::{
+    collections::vec_deque::VecDeque,
+    fs::File,
+    io::{BufRead, BufReader, Error, ErrorKind, Result},
+    path::PathBuf,
+    time::SystemTime,
+};
+
+use i3monkit::{Block, Widget, WidgetUpdate};
+
+use crate::spark;
+
+const DISK_STATS_PAT: &str = "/sys/class/block/*/*/stat";
+
+struct TransferStat {
+    rx: u64,
+    tx: u64,
+    ts: SystemTime,
+}
+
+impl TransferStat {
+    fn read_stat(stat_path: &str) -> Result<Self> {
+        // File format documented at https://www.kernel.org/doc/Documentation/block/stat.txt
+        let s = std::fs::read_to_string(stat_path)?;
+        let parts: Vec<_> = s.split(" ").filter(|s| !s.is_empty()).collect();
+
+        let rx = parts[2].parse().unwrap_or(0) * 512;
+        let tx = parts[6].parse().unwrap_or(0) * 512;
+        let ts = SystemTime::now();
+
+        Ok(Self { rx, tx, ts })
+    }
+    fn duration(&self, earlier: &Self) -> f64 {
+        let duration = self.ts.duration_since(earlier.ts).unwrap();
+        duration.as_secs() as f64 + duration.subsec_nanos() as f64 / 1_000_000_000.0
+    }
+
+    fn rx_rate(&self, earlier: &Self) -> f64 {
+        let duration = self.duration(earlier);
+        if duration < 1e-5 {
+            return std::f64::NAN;
+        }
+
+        (self.rx - earlier.rx) as f64 / duration
+    }
+
+    fn tx_rate(&self, earlier: &Self) -> f64 {
+        let duration = self.duration(earlier);
+        if duration < 1e-5 {
+            return std::f64::NAN;
+        }
+
+        (self.tx - earlier.tx) as f64 / duration
+    }
+}
+
+/// A widget that shows the disk R/W realtimely
+pub struct DiskSpeedWidget {
+    device: String,
+    last_stat: TransferStat,
+    rx_history: VecDeque<f32>,
+    tx_history: VecDeque<f32>,
+}
+
+impl DiskSpeedWidget {
+    /// Create the widget, for given device.
+    ///
+    /// **device** The interface to monitor
+    pub fn new(device: &str, num_samples: usize) -> Self {
+        let last_stat = TransferStat::read_stat(device).unwrap();
+        let device = device.to_string();
+        let rx_history: VecDeque<f32> = (0..num_samples).map(|_| 1.).collect();
+        let tx_history: VecDeque<f32> = (0..num_samples).map(|_| 1.).collect();
+        Self {
+            last_stat,
+            device,
+            rx_history,
+            tx_history,
+        }
+    }
+
+    fn is_active(&self) -> bool {
+        (self.rx_history.iter().sum::<f32>() + self.tx_history.iter().sum::<f32>()) > 0.
+    }
+
+    fn format_rate(rate: f64) -> String {
+        if rate.is_nan() {
+            return "N/A".to_string();
+        }
+
+        const UNIT_NAME: [&str; 6] = [" B/s", "KB/s", "MB/s", "GB/s", "TB/s", "PB/s"];
+
+        let mut best_unit = UNIT_NAME[0];
+        let mut best_multiplier = 1.0;
+
+        for unit in UNIT_NAME[1..].iter() {
+            if best_multiplier > rate / 1024.0 {
+                break;
+            }
+            best_unit = unit;
+            best_multiplier *= 1024.0
+        }
+
+        format!("{:6.1}{}", rate / best_multiplier, best_unit)
+    }
+
+    fn get_human_readable_stat(&mut self) -> Result<(String, String)> {
+        let cur_stat = TransferStat::read_stat(&self.device)?;
+
+        let rx_rate = cur_stat.rx_rate(&self.last_stat);
+        let tx_rate = cur_stat.tx_rate(&self.last_stat);
+
+        self.rx_history.push_back(rx_rate as f32);
+        self.rx_history.pop_front();
+        self.tx_history.push_back(tx_rate as f32);
+        self.tx_history.pop_front();
+        self.last_stat = cur_stat;
+
+        Ok((Self::format_rate(rx_rate), Self::format_rate(tx_rate)))
+    }
+
+    fn render(&mut self) -> Option<String> {
+        if let Ok((rx, tx)) = self.get_human_readable_stat() {
+            let (rx_history, tx_history) = if self.rx_history.len() > 1 {
+                let g = spark::Graph {
+                    min: None,
+                    max: None,
+                };
+                (
+                    g.render(&self.rx_history.iter().cloned().collect::<Vec<f32>>()),
+                    g.render(&self.tx_history.iter().cloned().collect::<Vec<f32>>()),
+                )
+            } else {
+                ("".to_string(), "".to_string())
+            };
+            return Some(format!(
+                "R:<tt>{} {}</tt> W:<tt>{} {}</tt> ",
+                rx, rx_history, tx, tx_history
+            ));
+        }
+        None
+    }
+}
+
+impl Widget for DiskSpeedWidget {
+    fn update(&mut self) -> Option<WidgetUpdate> {
+        if let Some(text) = self.render() {
+            let mut data = Block::new();
+            data.use_pango();
+            data.append_full_text(&text);
+            return Some(WidgetUpdate {
+                refresh_interval: std::time::Duration::new(1, 0),
+                data: Some(data),
+            });
+        }
+        None
+    }
+}
+
+pub struct AllDiskSpeedWidget {
+    num_samples: usize,
+    disks: Vec<DiskSpeedWidget>,
+}
+impl AllDiskSpeedWidget {
+    pub fn new(num_samples: usize) -> Self {
+        let disks = glob::glob(DISK_STATS_PAT)
+            .expect("couldn't find glob for disks")
+            .map(|path| {
+                let p = path.unwrap().display().to_string();
+                DiskSpeedWidget::new(&p, num_samples)
+            })
+            .collect();
+
+        AllDiskSpeedWidget { num_samples, disks }
+    }
+}
+
+impl Widget for AllDiskSpeedWidget {
+    fn update(&mut self) -> Option<WidgetUpdate> {
+        let disks: Vec<_> = self
+            .disks
+            .iter_mut()
+            .filter_map(|dev| {
+                if !dev.is_active() {
+                    return None;
+                }
+                if let Some(text) = dev.render() {
+                    let name = dev.device.rsplit("/").skip(1).nth(0).unwrap();
+                    Some(format!("{}: {}", name, text))
+                } else {
+                    None
+                }
+            })
+            .collect();
+
+        let mut data = Block::new();
+        data.use_pango();
+        disks.iter().for_each(|n| {
+            data.append_full_text(n);
+        });
+        Some(WidgetUpdate {
+            refresh_interval: std::time::Duration::new(1, 0),
+            data: Some(data),
+        })
+    }
+}

src/widgets/mod.rs

@@ -1,4 +1,5 @@
 pub mod cpu;
 pub mod datetime;
+pub mod disk;
 pub mod network;
 pub mod power;

src/widgets/network.rs

@@ -86,8 +86,8 @@ impl NetworkSpeedWidget {
     pub fn new(interface: &str, num_samples: usize) -> Self {
         let last_stat = TransferStat::read_stat(interface).unwrap();
         let interface = interface.to_string();
-        let rx_history: VecDeque<f32> = (0..num_samples).map(|_| 0.).collect();
-        let tx_history: VecDeque<f32> = (0..num_samples).map(|_| 0.).collect();
+        let rx_history: VecDeque<f32> = (0..num_samples).map(|_| 1.).collect();
+        let tx_history: VecDeque<f32> = (0..num_samples).map(|_| 1.).collect();
         Self {
             last_stat,
             interface,
@@ -97,9 +97,7 @@ impl NetworkSpeedWidget {
     }
 
     fn is_active(&self) -> bool {
-        self.last_stat.rx > 0
-            || self.last_stat.tx > 0
-            || (self.rx_history.iter().sum::<f32>() + self.tx_history.iter().sum::<f32>()) > 0.
+        (self.rx_history.iter().sum::<f32>() + self.tx_history.iter().sum::<f32>()) > 0.
     }
 
     fn format_rate(rate: f64) -> String {