Add support for disk graphs

Cargo.lock

@@ -1182,8 +1182,8 @@ dependencies = [
 ]
 
 [[package]]
-name = "wk-cpu-graph"
-version = "0.1.0"
+name = "wk-graph"
+version = "0.2.0"
 dependencies = [
  "chrono",
  "clap",

Cargo.toml

@@ -1,8 +1,8 @@
 [package]
-name = "wk-cpu-graph"
+name = "wk-graph"
 authors = ["2Shirt <2xShirt@gmail.com>"]
 license = "GPL"
-version = "0.1.0"
+version = "0.2.0"
 edition = "2024"
 
 [dependencies]

src/main.rs

@@ -1,17 +1,31 @@
-use clap::Parser;
+use clap::{Parser, Subcommand};
+use core::f64;
 use plotters::prelude::*;
 use serde::Deserialize;
-use std::{error::Error, fs, path::PathBuf};
+use std::{fs, path::PathBuf};
 use toml;
 
-#[derive(Parser, Debug)]
+#[derive(Parser)]
 #[command(version, about, long_about = None)]
+#[command(propagate_version = true)]
 struct Args {
+    #[command(subcommand)]
+    command: Commands,
+
     /// Sets a custom config file
     #[arg(short, long, value_name = "FILE")]
     config: Option<PathBuf>,
 }
 
+#[derive(Subcommand)]
+enum Commands {
+    /// Export CPU graph
+    Cpu,
+
+    /// Export drive graph
+    Disk,
+}
+
 struct LineColors {
     index: usize,
     colors: Vec<RGBColor>,
@@ -35,14 +49,14 @@ impl LineColors {
 }
 
 #[derive(Deserialize)]
-struct Data {
+struct CpuData {
     cpu: String,
     temps: Vec<Temps>,
     tests: Vec<Tests>,
     averages: Vec<Averages>,
 }
 
-impl Data {
+impl CpuData {
     fn len(&self) -> i32 {
         self.temps
             .iter()
@@ -70,6 +84,26 @@ impl Data {
     }
 }
 
+#[derive(Deserialize)]
+struct DiskData {
+    label: String,
+    read_rates: Vec<f64>,
+}
+
+impl DiskData {
+    fn len(&self) -> i32 {
+        self.read_rates.len() as i32
+    }
+    fn max_rate(&self) -> f64 {
+        self.read_rates
+            .iter()
+            .fold(f64::MIN_POSITIVE, |a, &b| a.max(b))
+    }
+    fn min_rate(&self) -> f64 {
+        self.read_rates.iter().fold(f64::INFINITY, |a, &b| a.min(b))
+    }
+}
+
 #[derive(Deserialize)]
 struct Temps {
     name: String,
@@ -98,13 +132,25 @@ struct Averages {
     temp: f64,
 }
 
-fn main() -> Result<(), Box<dyn Error>> {
+type Error = Box<dyn std::error::Error>;
+
+type Result<T> = core::result::Result<T, Error>;
+
+fn main() -> Result<()> {
     // Parse data
     let args = Args::parse();
     let config_path = args.config.expect("Config not specified.");
+    let _ = match &args.command {
+        Commands::Cpu => export_cpu_graph(config_path),
+        Commands::Disk => export_drive_graph(config_path),
+    };
+    Ok(())
+}
+
+fn export_cpu_graph(config_path: PathBuf) -> Result<()> {
     let out_path = config_path.with_extension("png");
     let toml_data = fs::read_to_string(&config_path).expect("Failed to read data file.");
-    let data: Data = toml::from_str(&toml_data).expect("Failed to parse TOML data");
+    let data: CpuData = toml::from_str(&toml_data).expect("Failed to parse TOML data");
     let max_temp = f64::ceil((data.max_temp() + 1.0) / 5.0) * 5.0;
     let min_temp = f64::floor((data.min_temp() - 1.0) / 5.0) * 5.0;
     let num_temps = data.len();
@@ -243,6 +289,70 @@ fn main() -> Result<(), Box<dyn Error>> {
     Ok(())
 }
 
+fn export_drive_graph(config_path: PathBuf) -> Result<()> {
+    let out_path = config_path.with_extension("png");
+    let toml_data = fs::read_to_string(&config_path).expect("Failed to read data file.");
+    let data: DiskData = toml::from_str(&toml_data).expect("Failed to parse TOML data");
+    let max_rate = f64::ceil((data.max_rate() + 1.0) / 50.0) * 50.0;
+    let min_rate = f64::floor((data.min_rate() - 1.0) / 50.0) * 50.0;
+    let num_rates = data.len();
+
+    // Chart data
+    // TODO: Switch to SVG
+    // let root = SVGBackend::new(out_path, (1600, 900)).into_drawing_area();
+    let root = BitMapBackend::new(&out_path, (1600, 900)).into_drawing_area();
+
+    root.fill(&WHITE)?;
+
+    let mut chart = ChartBuilder::on(&root)
+        .margin(5)
+        .caption("I/O Benchmark", ("sans-serif", 30))
+        .set_label_area_size(LabelAreaPosition::Left, 60)
+        .set_label_area_size(LabelAreaPosition::Right, 60)
+        .set_label_area_size(LabelAreaPosition::Bottom, 20)
+        .build_cartesian_2d(0..100, min_rate..max_rate)?
+        .set_secondary_coord(0..100, min_rate..max_rate);
+
+    chart
+        .configure_mesh()
+        .disable_x_mesh()
+        .disable_y_mesh()
+        .x_labels(1)
+        .max_light_lines(4)
+        .y_desc("Rates (MiB/s)")
+        .draw()?;
+
+    // Read rates
+    chart
+        .draw_series(LineSeries::new(
+            data.read_rates.iter().enumerate().map(|(i, rate)| {
+                let index = (i + 1) as i32;
+                (index, *rate)
+            }),
+            &BLUE,
+        ))
+        .unwrap()
+        .label(&data.label)
+        .legend(move |(x, y)| Rectangle::new([(x - 15, y + 1), (x, y)], BLUE));
+
+    // Export to file
+    chart
+        .configure_series_labels()
+        .position(SeriesLabelPosition::LowerRight)
+        .margin(20)
+        .legend_area_size(5)
+        .border_style(BLACK)
+        .background_style(RGBColor(192, 192, 192))
+        .label_font(("Calibri", 20))
+        .draw()?;
+
+    // To avoid the IO failure being ignored silently, we manually call the present function
+    root.present().expect("Unable to write result to file, please make sure 'plotters-doc-data' dir exists under current dir");
+
+    // Done
+    Ok(())
+}
+
 #[test]
 fn entry_point() {
     main().unwrap()