use std::fs::File;
use std::io;
use std::collections::HashMap;
use std::time::Duration;
use clap::Parser;
use diskstats::IoStatsReport;
pub mod diskstats;
#[derive(Parser)]
#[command(author, version, about)]
struct Args {
#[arg(short = 'x', default_value = "1")]
interval: u64,
#[arg(short = 'd', long = "devices", value_delimiter = ' ', num_args = 1..)]
devices: Option<Vec<String>>,
}
fn main() -> io::Result<()> {
let args = Args::parse();
let mut file = File::open("/proc/diskstats")?;
let mut buffer = String::new();
let mut prev_report = None;
let mut tree_view: String;
let mut tree_devices: HashMap<String, ()>;
match get_block_tree(&args.devices) {
Err(e) => {
eprintln!("{}", e);
return Err(e);
},
Ok((devices, view)) => {
tree_view = view;
tree_devices = devices;
}
}
loop {
let err = io::Read::read_to_string(&mut file, &mut buffer);
if err.is_err() {
eprintln!("Failed to read diskstats: {}", err.as_ref().unwrap());
return Err(err.unwrap_err());
}
let err = io::Seek::seek(&mut file, io::SeekFrom::Start(0));
if err.is_err() {
eprintln!("Failed to reset diskstats: {}", err.as_ref().unwrap());
return Err(err.unwrap_err());
}
let report = diskstats::DiskStatsReport::from_output(&buffer);
if let Some(prev_report) = prev_report {
let delta = report.delta(&prev_report, args.interval as f64);
if args.devices.is_none() && !are_all_devices_in_tree(&delta, &tree_devices) {
match get_block_tree(&args.devices) {
Ok((devices, view)) => {
tree_view = view;
tree_devices = devices;
},
Err(e) => {
eprintln!("Failed to get block tree: {}", e);
return Err(e);
}
}
}
println!("{}", build_iotree(&tree_view, &tree_devices, &delta));
}
prev_report = Some(report);
std::thread::sleep(Duration::from_secs(args.interval));
}
}
fn get_block_tree(filter: &Option<Vec<String>>) -> Result<(HashMap<String, ()>, String), io::Error> {
let tree_view = block_tree_view(filter)?;
let tree_devices = get_tree_devices(&tree_view);
Ok((tree_devices, tree_view))
}
fn block_tree_view(filter: &Option<Vec<String>>) -> Result<String, io::Error> {
let mut cmd = "lsblk -s -o NAME,MAJ:MIN".to_string();
if let Some(filter) = filter {
for device in filter {
cmd.push_str(format!(" {}", device).as_str());
}
}
let output = std::process::Command::new("bash")
.arg("-c")
.arg(cmd)
.output()?;
if !output.status.success() {
Err(io::Error::new(io::ErrorKind::Other, String::from_utf8_lossy(&output.stderr)))
} else {
Ok(String::from_utf8_lossy(&output.stdout).to_string())
}
}
fn get_tree_devices(tree_view: &str) -> HashMap<String, ()> {
let mut devices = HashMap::new();
for line in tree_view.lines().skip(1) {
let parts: Vec<&str> = line.split_whitespace().collect();
if parts.len() >= 2 {
let name = parts.last().unwrap().to_string();
devices.insert(name, ());
}
}
devices
}
fn are_all_devices_in_tree(delta : &IoStatsReport, tree_devices : &HashMap<String, ()>) -> bool {
for (device, _) in delta.devices.iter() {
if !tree_devices.contains_key(device) {
return false;
}
}
true
}
fn build_iotree(tree_view: &str, tree_devices: &HashMap<String, ()>, delta: &IoStatsReport) -> String {
let mut iotree = String::new();
let max_len = tree_view.lines().map(|l| l.len()).max().unwrap();
let mut lines = tree_view.lines();
let header = lines.next().unwrap();
let mut row = String::new();
row.push_str(&format!("{:<1$}", header, max_len + 1));
row.push_str(&format!(" {:>8} {:>10} {:>10} {:>10} {:>10} {:>10} {:>10} {:>10} {:>10} {:>10} {:>10} {:>10} {:>8}\n",
"R_IO/s", "R_BW/s", "R_LAT", "R_SZ", "W_IO/s", "W_BW/s", "W_LAT", "W_SZ", "D_IO/s", "D_BW/s", "D_LAT", "D_SZ", "QD"));
iotree.push_str(&row);
for line in lines {
let parts: Vec<&str> = line.split_whitespace().collect();
if parts.len() >= 2 {
let device_id = parts.last().unwrap().to_string();
if !tree_devices.contains_key(&device_id) {
continue;
}
if let Some(stats) = delta.devices.get(&device_id) {
let mut row = String::new();
row.push_str(&format!("{:<1$}", line, max_len + 1));
row.push_str(&format!(" {:>8} {:>10} {:>10} {:>10} {:>10} {:>10} {:>10} {:>10} {:>10} {:>10} {:>10} {:>10} {:>8}\n",
ios_to_str(stats.read_ios),
bytes_to_str(stats.read_bw),
latency_to_str(stats.read_latency),
bytes_to_str(stats.read_io_size),
ios_to_str(stats.write_ios),
bytes_to_str(stats.write_bw),
latency_to_str(stats.write_latency),
bytes_to_str(stats.write_io_size),
ios_to_str(stats.discard_ios),
bytes_to_str(stats.discard_bw),
latency_to_str(stats.discard_latency),
bytes_to_str(stats.discard_io_size),
stats.queue_depth));
iotree.push_str(&row);
}
}
}
iotree
}
fn bytes_to_str(bytes: f64) -> String {
if bytes < 1024.0 {
return format!("{:0.2}B", bytes);
}
let mut size = bytes;
let mut i = 0;
while size >= 1024.0 && i < 4 {
size /= 1024.0;
i += 1;
}
let unit = match i {
1 => "KiB",
2 => "MiB",
3 => "GiB",
_ => "TiB",
};
format!("{:0.2}{}", size, unit)
}
fn ios_to_str(ios: f64) -> String {
if ios < 1000.0 {
return format!("{:0.2}", ios);
}
let mut size = ios;
let mut i = 0;
while size >= 1000.0 && i < 3 {
size /= 1000.0;
i += 1;
}
let unit = match i {
1 => "K",
2 => "M",
_ => "G"
};
format!("{:0.2}{}", size, unit)
}
fn latency_to_str(latency: f64) -> String {
if latency < 1000.0 {
return format!("{:0.2}us", latency);
}
let mut size = latency;
let mut i = 0;
while size >= 1000.0 && i < 3 {
size /= 1000.0;
i += 1;
}
let unit = match i {
1 => "ms",
_ => "s",
};
format!("{:0.2}{}", size, unit)
}