add zfs io counters for linux and freebsd (#1248)

* add zfs io counters for linux and freebsd

* ci

* freebsd clippy

* code review: remove dead code for zfs feature gate of freebsd iocounters and squash if statement in zfs_io_counters

5 files changed, 205 insertions, 8 deletions

diff --git a/src/app/data_farmer.rs b/src/app/data_farmer.rs
index 37e58229..d2edb53a 100644
--- a/src/app/data_farmer.rs
+++ b/src/app/data_farmer.rs
@@ -347,6 +347,13 @@ impl DataCollection {
                             None => device.name.split('/').last(),
                         }
                     } else {
+                        #[cfg(feature = "zfs")]
+                        if ! device.name.starts_with('/'){
+                            Some(device.name.as_str()) // use the whole zfs dataset name
+                        } else {
+                            device.name.split('/').last()
+                        }
+                        #[cfg(not(feature = "zfs"))]
                         device.name.split('/').last()
                     }
                 }
diff --git a/src/app/data_harvester/disks.rs b/src/app/data_harvester/disks.rs
index 3359fcc7..5075f335 100644
--- a/src/app/data_harvester/disks.rs
+++ b/src/app/data_harvester/disks.rs
@@ -8,12 +8,20 @@ use crate::app::filter::Filter;
 cfg_if! {
     if #[cfg(target_os = "freebsd")] {
         mod freebsd;
+        #[cfg(feature = "zfs")]
+        mod io_counters;
+        #[cfg(feature = "zfs")]
+        mod zfs_io_counters;
+        #[cfg(feature = "zfs")]
+        pub use io_counters::IoCounters;
         pub(crate) use self::freebsd::*;
     } else if #[cfg(target_os = "windows")] {
         mod windows;
         pub(crate) use self::windows::*;
     } else if #[cfg(target_os = "linux")] {
         mod unix;
+        #[cfg(feature = "zfs")]
+        mod zfs_io_counters;
         pub(crate) use self::unix::*;
     } else if #[cfg(target_os = "macos")] {
         mod unix;
diff --git a/src/app/data_harvester/disks/freebsd.rs b/src/app/data_harvester/disks/freebsd.rs
index aa127ef6..a10fa642 100644
--- a/src/app/data_harvester/disks/freebsd.rs
+++ b/src/app/data_harvester/disks/freebsd.rs
@@ -6,7 +6,11 @@ use serde::Deserialize;
 
 use super::{keep_disk_entry, DiskHarvest, IoHarvest};
 
-use crate::{app::data_harvester::DataCollector, data_harvester::deserialize_xo, utils::error};
+use crate::{
+    app::data_harvester::DataCollector, data_harvester::deserialize_xo,
+    data_harvester::disks::IoData, utils::error,
+};
+use hashbrown::HashMap;
 
 #[derive(Deserialize, Debug, Default)]
 #[serde(rename_all = "kebab-case")]
@@ -25,14 +29,32 @@ struct FileSystem {
 }
 
 pub fn get_io_usage() -> error::Result<IoHarvest> {
-    let io_harvest = get_disk_info().map(|storage_system_information| {
-        storage_system_information
-            .filesystem
-            .into_iter()
-            .map(|disk| (disk.name, None))
-            .collect()
-    })?;
+    #[allow(unused_mut)]
+    let mut io_harvest: HashMap<String, Option<IoData>> =
+        get_disk_info().map(|storage_system_information| {
+            storage_system_information
+                .filesystem
+                .into_iter()
+                .map(|disk| (disk.name, None))
+                .collect()
+        })?;
 
+    #[cfg(feature = "zfs")]
+    {
+        use crate::app::data_harvester::disks::zfs_io_counters;
+        if let Ok(zfs_io) = zfs_io_counters::zfs_io_stats() {
+            for io in zfs_io.into_iter().flatten() {
+                let mount_point = io.device_name().to_string_lossy();
+                io_harvest.insert(
+                    mount_point.to_string(),
+                    Some(IoData {
+                        read_bytes: io.read_bytes(),
+                        write_bytes: io.write_bytes(),
+                    }),
+                );
+            }
+        }
+    }
     Ok(io_harvest)
 }
 
diff --git a/src/app/data_harvester/disks/unix/linux/counters.rs b/src/app/data_harvester/disks/unix/linux/counters.rs
index f01d5b7d..d65863ba 100644
--- a/src/app/data_harvester/disks/unix/linux/counters.rs
+++ b/src/app/data_harvester/disks/unix/linux/counters.rs
@@ -81,5 +81,13 @@ pub fn io_stats() -> anyhow::Result<Vec<anyhow::Result<IoCounters>>> {
         }
     }
 
+    #[cfg(feature = "zfs")]
+    {
+        use crate::app::data_harvester::disks::zfs_io_counters;
+        if let Ok(mut zfs_io) = zfs_io_counters::zfs_io_stats() {
+            results.append(&mut zfs_io);
+        }
+    }
+
     Ok(results)
 }
diff --git a/src/app/data_harvester/disks/zfs_io_counters.rs b/src/app/data_harvester/disks/zfs_io_counters.rs
new file mode 100644
index 00000000..a53e8af3
--- /dev/null
+++ b/src/app/data_harvester/disks/zfs_io_counters.rs
@@ -0,0 +1,152 @@
+use crate::app::data_harvester::disks::IoCounters;
+
+/// Returns zpool I/O stats. Pulls data from `sysctl kstat.zfs.{POOL}.dataset.{objset-*}`
+#[cfg(target_os = "freebsd")]
+pub fn zfs_io_stats() -> anyhow::Result<Vec<anyhow::Result<IoCounters>>> {
+    use sysctl::Sysctl;
+    let zfs_ctls: Vec<_> = sysctl::Ctl::new("kstat.zfs.")?
+        .into_iter()
+        .filter_map(|e| {
+            e.ok().and_then(|ctl| {
+                let name = ctl.name();
+                if let Ok(name) = name {
+                    if name.contains("objset-")
+                        && (name.contains("dataset_name")
+                            || name.contains("nwritten")
+                            || name.contains("nread"))
+                    {
+                        Some(ctl)
+                    } else {
+                        None
+                    }
+                } else {
+                    None
+                }
+            })
+        })
+        .collect();
+
+    use itertools::Itertools;
+    let results: Vec<anyhow::Result<IoCounters>> = zfs_ctls
+        .iter()
+        .chunks(3)
+        .into_iter()
+        .filter_map(|chunk| {
+            let mut nread = 0;
+            let mut nwrite = 0;
+            let mut ds_name = String::new();
+            for ctl in chunk {
+                if let Ok(name) = ctl.name() {
+                    if name.contains("dataset_name") {
+                        ds_name = ctl.value_string().ok()?;
+                    } else if name.contains("nread") {
+                        if let Ok(sysctl::CtlValue::U64(val)) = ctl.value() {
+                            nread = val;
+                        }
+                    } else if name.contains("nwritten") {
+                        if let Ok(sysctl::CtlValue::U64(val)) = ctl.value() {
+                            nwrite = val;
+                        }
+                    }
+                }
+            }
+            Some(Ok(IoCounters::new(ds_name, nread, nwrite)))
+        })
+        .collect();
+    Ok(results)
+}
+
+/// Returns zpool I/O stats. Pulls data from `/proc/spl/kstat/zfs/*/objset-*`.
+#[cfg(target_os = "linux")]
+pub fn zfs_io_stats() -> anyhow::Result<Vec<anyhow::Result<IoCounters>>> {
+    if let Ok(zpools) = std::fs::read_dir("/proc/spl/kstat/zfs") {
+        let zpools_vec: Vec<std::path::PathBuf> = zpools
+            .filter_map(|e| {
+                e.ok().and_then(|d| {
+                    let p = d.path();
+                    if p.is_dir() {
+                        Some(p)
+                    } else {
+                        None
+                    }
+                })
+            })
+            .collect();
+        let results = zpools_vec
+            .iter()
+            .filter_map(|zpool| {
+                // go through each pool
+                if let Ok(datasets) = std::fs::read_dir(zpool) {
+                    let datasets_vec: Vec<std::path::PathBuf> =
+                        datasets // go through dataset
+                            .filter_map(|e| {
+                                e.ok().and_then(|d| {
+                                    let p = d.path();
+                                    if p.is_file() && p.to_str()?.contains("objset-") {
+                                        Some(p)
+                                    } else {
+                                        None
+                                    }
+                                })
+                            })
+                            .collect();
+                    let io_counters: Vec<anyhow::Result<IoCounters>> = datasets_vec
+                        .iter()
+                        .filter_map(|ds| {
+                            // get io-counter from each dataset
+                            if let Ok(contents) = std::fs::read_to_string(ds) {
+                                let mut read = 0;
+                                let mut write = 0;
+                                let mut name = "";
+                                contents.lines().for_each(|line| {
+                                    if let Some((label, value)) = line.split_once(' ') {
+                                        match label {
+                                            "dataset_name" => {
+                                                if let Some((_type, val)) =
+                                                    value.trim_start().rsplit_once(' ')
+                                                {
+                                                    name = val;
+                                                }
+                                            }
+                                            "nwritten" => {
+                                                if let Some((_type, val)) =
+                                                    value.trim_start().rsplit_once(' ')
+                                                {
+                                                    if let Ok(number) = val.parse::<u64>() {
+                                                        write = number;
+                                                    }
+                                                }
+                                            }
+                                            "nread" => {
+                                                if let Some((_type, val)) =
+                                                    value.trim_start().rsplit_once(' ')
+                                                {
+                                                    if let Ok(number) = val.parse::<u64>() {
+                                                        read = number;
+                                                    }
+                                                }
+                                            }
+                                            _ => {}
+                                        }
+                                    }
+                                });
+                                let counter = IoCounters::new(name.to_owned(), read, write);
+                                //log::debug!("adding io counter for zfs {:?}", counter);
+                                Some(Ok(counter))
+                            } else {
+                                None
+                            }
+                        })
+                        .collect();
+                    Some(io_counters)
+                } else {
+                    None
+                }
+            })
+            .flatten()
+            .collect(); // combine io-counters
+        Ok(results)
+    } else {
+        Err(anyhow::anyhow!("Unable to open zfs proc directory"))
+    }
+}