Get simple IO stats using Rust (throughput, ...)
Tags: linux, rust, tutorial, beginners
Published: 2020-11-11 17:06:29

The other day I wanted to grab some information about my disk performance in real-time. I saw the really useful command `iostat` and got curious. How does this tool can know how many megabytes per seconds my disks is using, etc.

So I dug into the [source code](https://github.com/sysstat/sysstat/blob/master/iostat.c) of this command and found my answer.

Basically it just read `/proc/diskstats` and that's all (if we don't care about formatting, sorting, etc).

As I'm learning Rust (and doing a project involving stats) I thought it would be great to write a simple program to only display how many megabytes per seconds my disks are using. And so my journey began.

## Setting up the project
We can create our project with **cargo**. Cargo is a package manager for Rust and you can use it to simplify your life.
```bash
cargo new nameofbinary --bin
cd nameofbinary
```
_(if you've never used Rust before, take a look at @aligoren's fantastic guide [here](https://dev.to/aligoren/getting-started-with-rust-1c19))_
## Let's code our program
What will we need ? First take a look at the `/proc/diskstats`;
```
 259       0 nvme0n1 93188 31843 6970296 27235 1834264 427303 105983322 432380 0 411608 472996 0 0 0 0 33212 13380
```
**Wow**. What a mess. However we can easily understand it by a little bit of googling ([doc from kernel.org](https://www.kernel.org/doc/Documentation/ABI/testing/procfs-diskstats)).
By reading it we understand that each value is organized and separated by whitespace. Great !

Now we can start coding (we'll do everything in the main.rs).
We'll first create a `struct` which will handle the number of megabytes one disk has read or written.
```rust
// This structure will hold our data for the disks
struct IoStats {
    pub mb_read: f64, // f64 cause it won't be an integer, we need precision
    pub mb_wrtn: f64,
}
```
Here is the pseudo code of our program first;
```
hold_prev = previous stats reference
file = file descriptor to /proc/diskstats
loop:
    hold_curr = current stats reference
    io_stats = content of the file
    for line in io_stats:
        split_field = line.split (get each value in an array)
        compute_mb_s() = calculate the diff between hold_prev and hold_curr
    write_result
    hold_prev = hold_curr
    wait 1s
```
Let's dive into the real code now;
```rust
// Hashmap of previous drives stats to compute difference from
let mut prev: HashMap<String, IoStats> = HashMap::new();
// Open the file we'll use to get the stats
let mut fd = File::open(&"/proc/diskstats").unwrap();
```
* `prev` is a hashmap which will help us find the prev disk more easily by checking only by the disk name.
* `fd` is our file descriptor (our way to read the diskstats file).
```rust
// Main program loop
loop {
    // Create the curr Hashmap, allow us to compare with the prev one
    let mut curr: HashMap<String, IoStats> = HashMap::new();

    ...

    prev = curr;
    thread::sleep(Duration::from_secs(1));
}
```
* `curr` is the same hashmap type as prev and will hold current value of disks stats.
* `prev = curr` is where we give an utility to `prev` (now prev contains all the value as curr)
* `thread::sleep` pause the loop for 1s (this avoid spamming the cpu/terminal)
```rust
let mut curr: HashMap<String, IoStats> = HashMap::new();
// Create the output string
let mut output = String::new();
// Add the header string to the output
output.push_str("\nDevice          mb_reads/s      mb_wrtn/s\n\n");
// Create a new empty string
let mut io_data = String::new();
// Read the content of the file (diskstats) to the io_data string
fd.read_to_string(&mut io_data).unwrap();
// Iterate over each line (each disk)
for line in io_data.lines() {

    ...

}
// Move the cursor to the start of the file
fd.seek(SeekFrom::Start(0)).unwrap();
// Print the result
writeln!(io::stdout().lock(), "{}", output);
prev = curr;
```
* `output` is the string we'll print (single string to print everything at once)
* `output.push_str` add the header to the output string
* `fd.read_to_string` read the content of `fd` into `io_data`
* `for` iterate over each line from `io_data`
* `fd.seek` reset the cursor position to the start of the file (allow use to read it again)
* `writeln!` write the output string in `stdout` (terminal)
```rust
for line in io_data.lines() {
    // Split field (separated by whitespace) and collect them without specific type
    let fields = line.split_whitespace().collect::<Vec<_>>();
    let ds = IoStats {
        mb_read: fields[5].parse::<f64>().unwrap() / 2048.0,
        mb_wrtn: fields[9].parse::<f64>().unwrap() / 2048.0,
    };
    // If prev already contains the info we compute the diff to get mb/s
    // Else we add to the print line the "fake" data.
    if prev.contains_key(fields[2]) {
        // Get the object from the hashmap
        let pds = prev.get(fields[2]).unwrap();
        // Construct speed line and append it to curr hashmap
        let mb_read_s = ds.mb_read - pds.mb_read;
        let mb_wrtn_s = ds.mb_wrtn - pds.mb_wrtn;
        // Add the line, formatted with color and spacing
        output.push_str(&format!("\x1b[0;32m{:16}\x1b[0m\x1b[0;34m{:10.2}{:15.2}\x1b[0m\n", fields[2], mb_read_s, mb_wrtn_s));
        // Insert the current disk data to the curr HashMap
        // the curr will later be saved as prev
        curr.insert(fields[2].to_owned(), ds);
    } else {
        // Add the line with fake data and formatting
        output.push_str(&format!("\x1b[0;32m{:16}\x1b[0m\x1b[0;34m{:10.2}{:15.2}\x1b[0m\n", fields[2], 0.00, 0.00));
    }
}
```
* `line.split_whitespace().collect::<Vec<_>>()` simply return us each values, without whitespace, as &str
* `ds` create our IoStats struct for this disk
* `ds -> mb_read/mb_wrtn` are first converted to `f64` and then divided by `2048.0` to convert from sector (one sector is almost always 512kb) to mb.
* `prev.contains_key` check if the previous hashmap contains the disk (in case of first run or usb stick)
* `pds` is the IoStats of the previous run of the loop for that particular disk
* `mb_read_s` and `mb_wrtn_s` are the mb/s value (/s as we loop every 1s)
* `output.push_str` add the current disk's info to the output string
* `curr.insert` add the current disk's info to the curr hashmap

## Final words
As we saw it's pretty easy to get stats from our disks in Rust. Everything is properly placed inside `/proc/diskstats` for us.

If you want to take a look at the full code [github](https://github.com/Martichou/iostats_demo).

Thanks for ready it, hope you learned something new ! ❤