vmstat: implement --disk --disk-sum --partition

src/uu/vmstat/src/parser.rs

@@ -5,6 +5,8 @@
 
 #[cfg(target_os = "linux")]
 use std::collections::HashMap;
+#[cfg(target_os = "linux")]
+use std::fmt::{Debug, Display, Formatter};
 
 #[cfg(target_os = "linux")]
 pub fn parse_proc_file(path: &str) -> HashMap<String, String> {
@@ -31,6 +33,7 @@ pub struct ProcData {
     pub stat: HashMap<String, String>,
     pub meminfo: HashMap<String, String>,
     pub vmstat: HashMap<String, String>,
+    pub diskstat: Vec<String>,
 }
 #[cfg(target_os = "linux")]
 impl Default for ProcData {
@@ -45,11 +48,17 @@ impl ProcData {
         let stat = parse_proc_file("/proc/stat");
         let meminfo = parse_proc_file("/proc/meminfo");
         let vmstat = parse_proc_file("/proc/vmstat");
+        let diskstat = std::fs::read_to_string("/proc/diskstats")
+            .unwrap()
+            .lines()
+            .map(|line| line.to_string())
+            .collect();
         Self {
             uptime,
             stat,
             meminfo,
             vmstat,
+            diskstat,
         }
     }
 
@@ -228,3 +237,104 @@ impl Meminfo {
         }
     }
 }
+
+#[cfg(target_os = "linux")]
+#[derive(Debug)]
+pub struct DiskStatParseError;
+
+#[cfg(target_os = "linux")]
+impl Display for DiskStatParseError {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        std::fmt::Debug::fmt("Failed to parse diskstat line", f)
+    }
+}
+
+#[cfg(target_os = "linux")]
+impl std::error::Error for DiskStatParseError {}
+
+#[cfg(target_os = "linux")]
+pub struct DiskStat {
+    // Name from https://www.kernel.org/doc/html/latest/admin-guide/iostats.html
+    pub major: u64,
+    pub minor: u64,
+    pub device: String,
+    pub reads_completed: u64,
+    pub reads_merged: u64,
+    pub sectors_read: u64,
+    pub milliseconds_spent_reading: u64,
+    pub writes_completed: u64,
+    pub writes_merged: u64,
+    pub sectors_written: u64,
+    pub milliseconds_spent_writing: u64,
+    pub ios_currently_in_progress: u64,
+    pub milliseconds_spent_doing_ios: u64,
+    pub weighted_milliseconds_spent_doing_ios: u64,
+    pub discards_completed: u64,
+    pub discards_merged: u64,
+    pub sectors_discarded: u64,
+    pub milliseconds_spent_discarding: u64,
+    pub flush_requests_completed: u64,
+    pub milliseconds_spent_flushing: u64,
+}
+
+#[cfg(target_os = "linux")]
+impl std::str::FromStr for DiskStat {
+    type Err = DiskStatParseError;
+
+    fn from_str(line: &str) -> Result<Self, Self::Err> {
+        let parts: Vec<&str> = line.split_whitespace().collect();
+        if parts.len() < 14 {
+            Err(DiskStatParseError)?;
+        }
+
+        let parse_value = |s: &str| s.parse::<u64>().map_err(|_| DiskStatParseError);
+        let parse_optional_value = |s: Option<&&str>| match s {
+            None => Ok(0),
+            Some(value) => value.parse::<u64>().map_err(|_| DiskStatParseError),
+        };
+
+        Ok(Self {
+            major: parse_value(parts[0])?,
+            minor: parse_value(parts[1])?,
+            device: parts[2].to_string(),
+            reads_completed: parse_value(parts[3])?,
+            reads_merged: parse_value(parts[4])?,
+            sectors_read: parse_value(parts[5])?,
+            milliseconds_spent_reading: parse_value(parts[6])?,
+            writes_completed: parse_value(parts[7])?,
+            writes_merged: parse_value(parts[8])?,
+            sectors_written: parse_value(parts[9])?,
+            milliseconds_spent_writing: parse_value(parts[10])?,
+            ios_currently_in_progress: parse_value(parts[11])?,
+            milliseconds_spent_doing_ios: parse_value(parts[12])?,
+            weighted_milliseconds_spent_doing_ios: parse_optional_value(parts.get(13))?,
+            discards_completed: parse_optional_value(parts.get(14))?,
+            discards_merged: parse_optional_value(parts.get(15))?,
+            sectors_discarded: parse_optional_value(parts.get(16))?,
+            milliseconds_spent_discarding: parse_optional_value(parts.get(17))?,
+            flush_requests_completed: parse_optional_value(parts.get(18))?,
+            milliseconds_spent_flushing: parse_optional_value(parts.get(19))?,
+        })
+    }
+}
+
+#[cfg(target_os = "linux")]
+impl DiskStat {
+    pub fn is_disk(&self) -> bool {
+        std::path::Path::new(&format!("/sys/block/{}", self.device)).exists()
+    }
+
+    pub fn current() -> Result<Vec<Self>, DiskStatParseError> {
+        let diskstats =
+            std::fs::read_to_string("/proc/diskstats").map_err(|_| DiskStatParseError)?;
+        let lines = diskstats.lines();
+        Self::from_proc_vec(&lines.map(|line| line.to_string()).collect::<Vec<_>>())
+    }
+
+    pub fn from_proc_vec(proc_vec: &[String]) -> Result<Vec<Self>, DiskStatParseError> {
+        proc_vec
+            .iter()
+            .map(|line| line.parse::<DiskStat>())
+            .collect()
+    }
+}

src/uu/vmstat/src/picker.rs

@@ -4,9 +4,11 @@
 // file that was distributed with this source code.
 
 #[cfg(target_os = "linux")]
-use crate::{CpuLoad, CpuLoadRaw, Meminfo, ProcData};
+use crate::{CpuLoad, CpuLoadRaw, DiskStat, Meminfo, ProcData};
 #[cfg(target_os = "linux")]
 use clap::ArgMatches;
+#[cfg(target_os = "linux")]
+use uucore::error::{UResult, USimpleError};
 
 #[cfg(target_os = "linux")]
 pub type Picker = (
@@ -203,6 +205,61 @@ pub fn get_stats() -> Vec<(String, u64)> {
     ]
 }
 
+#[cfg(target_os = "linux")]
+pub fn get_disk_sum() -> UResult<Vec<(String, u64)>> {
+    let disk_data = DiskStat::current()
+        .map_err(|_| USimpleError::new(1, "Unable to retrieve disk statistics"))?;
+
+    let mut disks = 0;
+    let mut partitions = 0;
+    let mut total_reads = 0;
+    let mut merged_reads = 0;
+    let mut read_sectors = 0;
+    let mut milli_reading = 0;
+    let mut writes = 0;
+    let mut merged_writes = 0;
+    let mut written_sectors = 0;
+    let mut milli_writing = 0;
+    let mut inprogress_io = 0;
+    let mut milli_spent_io = 0;
+    let mut milli_weighted_io = 0;
+
+    for disk in disk_data.iter() {
+        if disk.is_disk() {
+            disks += 1;
+            total_reads += disk.reads_completed;
+            merged_reads += disk.reads_merged;
+            read_sectors += disk.sectors_read;
+            milli_reading += disk.milliseconds_spent_reading;
+            writes += disk.writes_completed;
+            merged_writes += disk.writes_merged;
+            written_sectors += disk.sectors_written;
+            milli_writing += disk.milliseconds_spent_writing;
+            inprogress_io += disk.ios_currently_in_progress;
+            milli_spent_io += disk.milliseconds_spent_doing_ios / 1000;
+            milli_weighted_io += disk.weighted_milliseconds_spent_doing_ios / 1000;
+        } else {
+            partitions += 1;
+        }
+    }
+
+    Ok(vec![
+        ("disks".to_string(), disks),
+        ("partitions".to_string(), partitions),
+        ("total reads".to_string(), total_reads),
+        ("merged reads".to_string(), merged_reads),
+        ("read sectors".to_string(), read_sectors),
+        ("milli reading".to_string(), milli_reading),
+        ("writes".to_string(), writes),
+        ("merged writes".to_string(), merged_writes),
+        ("written sectors".to_string(), written_sectors),
+        ("milli writing".to_string(), milli_writing),
+        ("in progress IO".to_string(), inprogress_io),
+        ("milli spent IO".to_string(), milli_spent_io),
+        ("milli weighted IO".to_string(), milli_weighted_io),
+    ])
+}
+
 #[cfg(target_os = "linux")]
 fn with_unit(x: u64, arg: &ArgMatches) -> u64 {
     if let Some(unit) = arg.get_one::<String>("unit") {

src/uu/vmstat/src/vmstat.rs

@@ -7,7 +7,7 @@ mod parser;
 mod picker;
 
 #[cfg(target_os = "linux")]
-use crate::picker::{get_pickers, get_stats, Picker};
+use crate::picker::{get_disk_sum, get_pickers, get_stats, Picker};
 use clap::value_parser;
 #[allow(unused_imports)]
 use clap::{arg, crate_version, ArgMatches, Command};
@@ -26,22 +26,31 @@ pub fn uumain(args: impl uucore::Args) -> UResult<()> {
     let matches = uu_app().try_get_matches_from(args)?;
     #[cfg(target_os = "linux")]
     {
-        if matches.get_flag("forks") {
-            return print_forks();
-        }
-        if matches.get_flag("stats") {
-            return print_stats();
-        }
-
+        let wide = matches.get_flag("wide");
         let one_header = matches.get_flag("one-header");
         let no_first = matches.get_flag("no-first");
         let term_height = terminal_size::terminal_size()
             .map(|size| size.1 .0)
             .unwrap_or(0);
 
+        if matches.get_flag("forks") {
+            return print_forks();
+        }
         if matches.get_flag("slabs") {
             return print_slabs(one_header, term_height);
         }
+        if matches.get_flag("stats") {
+            return print_stats();
+        }
+        if matches.get_flag("disk") {
+            return print_disk(wide, one_header, term_height);
+        }
+        if matches.get_flag("disk-sum") {
+            return print_disk_sum();
+        }
+        if let Some(device) = matches.get_one::<String>("partition") {
+            return print_partition(device);
+        }
 
         // validate unit
         if let Some(unit) = matches.get_one::<String>("unit") {
@@ -141,6 +150,110 @@ fn print_slab_header() {
     );
 }
 
+#[cfg(target_os = "linux")]
+fn print_disk_header(wide: bool) {
+    if wide {
+        println!("disk- -------------------reads------------------- -------------------writes------------------ ------IO-------");
+        println!(
+            "{:>15} {:>9} {:>11} {:>11} {:>9} {:>9} {:>11} {:>11} {:>7} {:>7}",
+            "total", "merged", "sectors", "ms", "total", "merged", "sectors", "ms", "cur", "sec"
+        );
+    } else {
+        println!("disk- ------------reads------------ ------------writes----------- -----IO------");
+        println!(
+            "{:>12} {:>6} {:>7} {:>7} {:>6} {:>6} {:>7} {:>7} {:>6} {:>6}",
+            "total", "merged", "sectors", "ms", "total", "merged", "sectors", "ms", "cur", "sec"
+        );
+    }
+}
+
+#[cfg(target_os = "linux")]
+fn print_disk(wide: bool, one_header: bool, term_height: u16) -> UResult<()> {
+    let disk_data = DiskStat::current()
+        .map_err(|_| USimpleError::new(1, "Unable to retrieve disk statistics"))?;
+
+    let mut line_count = 0;
+
+    print_disk_header(wide);
+
+    for disk in disk_data {
+        if !disk.is_disk() {
+            continue;
+        }
+
+        if needs_header(one_header, term_height, line_count) {
+            print_disk_header(wide);
+        }
+        line_count += 1;
+
+        if wide {
+            println!(
+                "{:<5} {:>9} {:>9} {:>11} {:>11} {:>9} {:>9} {:>11} {:>11} {:>7} {:>7}",
+                disk.device,
+                disk.reads_completed,
+                disk.reads_merged,
+                disk.sectors_read,
+                disk.milliseconds_spent_reading,
+                disk.writes_completed,
+                disk.writes_merged,
+                disk.sectors_written,
+                disk.milliseconds_spent_writing,
+                disk.ios_currently_in_progress / 1000,
+                disk.milliseconds_spent_doing_ios / 1000
+            );
+        } else {
+            println!(
+                "{:<5} {:>6} {:>6} {:>7} {:>7} {:>6} {:>6} {:>7} {:>7} {:>6} {:>6}",
+                disk.device,
+                disk.reads_completed,
+                disk.reads_merged,
+                disk.sectors_read,
+                disk.milliseconds_spent_reading,
+                disk.writes_completed,
+                disk.writes_merged,
+                disk.sectors_written,
+                disk.milliseconds_spent_writing,
+                disk.ios_currently_in_progress / 1000,
+                disk.milliseconds_spent_doing_ios / 1000
+            );
+        }
+    }
+
+    Ok(())
+}
+
+#[cfg(target_os = "linux")]
+fn print_disk_sum() -> UResult<()> {
+    let data = get_disk_sum()?;
+
+    data.iter()
+        .for_each(|(name, value)| println!("{value:>13} {name}"));
+
+    Ok(())
+}
+
+#[cfg(target_os = "linux")]
+fn print_partition(device: &str) -> UResult<()> {
+    let disk_data = DiskStat::current()
+        .map_err(|_| USimpleError::new(1, "Unable to retrieve disk statistics"))?;
+
+    let disk = disk_data
+        .iter()
+        .find(|disk| disk.device == device)
+        .ok_or_else(|| USimpleError::new(1, format!("Disk/Partition {device} not found")))?;
+
+    println!(
+        "{device:<9} {:>11} {:>17} {:>11} {:>17}",
+        "reads", "read sectors", "writes", "requested writes"
+    );
+    println!(
+        "{:>21} {:>17} {:>11} {:>17}",
+        disk.reads_completed, disk.sectors_read, disk.writes_completed, disk.sectors_written
+    );
+
+    Ok(())
+}
+
 #[cfg(target_os = "linux")]
 fn print_header(pickers: &[Picker]) {
     let mut section: Vec<&str> = vec![];
@@ -191,15 +304,18 @@ pub fn uu_app() -> Command {
                 .value_parser(value_parser!(u64)),
             arg!(-a --active "Display active and inactive memory"),
             arg!(-f --forks "switch displays the number of forks since boot")
-                .conflicts_with_all(["slabs", "stats", /*"disk", "disk-sum", "partition"*/]),
+                .conflicts_with_all(["slabs", "stats", "disk", "disk-sum", "partition"]),
             arg!(-m --slabs "Display slabinfo")
-                .conflicts_with_all(["forks", "stats", /*"disk", "disk-sum", "partition"*/]),
+                .conflicts_with_all(["forks", "stats", "disk", "disk-sum", "partition"]),
             arg!(-n --"one-header" "Display the header only once rather than periodically"),
             arg!(-s --stats "Displays a table of various event counters and memory statistics")
-                .conflicts_with_all(["forks", "slabs", /*"disk", "disk-sum", "partition"*/]),
-            // arg!(-d --disk "Report disk statistics"),
-            // arg!(-D --"disk-sum" "Report some summary statistics about disk activity"),
-            // arg!(-p --partition <device> "Detailed statistics about partition"),
+                .conflicts_with_all(["forks", "slabs", "disk", "disk-sum", "partition"]),
+            arg!(-d --disk "Report disk statistics")
+                .conflicts_with_all(["forks", "slabs", "stats", "disk-sum", "partition"]),
+            arg!(-D --"disk-sum" "Report some summary statistics about disk activity")
+                .conflicts_with_all(["forks", "slabs", "stats", "disk", "partition"]),
+            arg!(-p --partition <device> "Detailed statistics about partition")
+                .conflicts_with_all(["forks", "slabs", "stats", "disk", "disk-sum"]),
             arg!(-S --unit <character> "Switches outputs between 1000 (k), 1024 (K), 1000000 (m), or 1048576 (M) bytes"),
             arg!(-t --timestamp "Append timestamp to each line"),
             arg!(-w --wide "Wide output mode"),

tests/by-util/test_vmstat.rs

@@ -92,3 +92,15 @@ fn test_timestamp() {
 fn test_stats() {
     new_ucmd!().arg("-s").succeeds();
 }
+
+#[test]
+#[cfg(target_os = "linux")]
+fn test_disk() {
+    new_ucmd!().arg("-d").succeeds();
+}
+
+#[test]
+#[cfg(target_os = "linux")]
+fn test_disk_sum() {
+    new_ucmd!().arg("-D").succeeds();
+}