Retrieve io stats from cgroup v2

src/cgroups/stats.rs

@@ -1,5 +1,5 @@
-use anyhow::{Context, Result};
-use std::{collections::HashMap, fmt::Display, fs, path::Path};
+use anyhow::{bail, Context, Result};
+use std::{collections::HashMap, fmt::Display, fs, hash::Hash, path::Path};
 
 use super::common;
 
@@ -294,6 +294,12 @@ pub fn supported_page_sizes() -> Result<Vec<String>> {
     Ok(sizes)
 }
 
+pub fn parse_value(value: &str) -> Result<u64> {
+    value
+        .parse()
+        .with_context(|| format!("failed to parse {}", value))
+}
+
 pub fn parse_single_value(file_path: &Path) -> Result<u64> {
     let value = common::read_cgroup_file(file_path)?;
     let value = value.trim();
@@ -334,6 +340,37 @@ pub fn parse_flat_keyed_data(file_path: &Path) -> Result<HashMap<String, u64>> {
     Ok(stats)
 }
 
+pub fn parse_nested_keyed_data(file_path: &Path) -> Result<HashMap<String, Vec<String>>> {
+    let mut stats: HashMap<String, Vec<String>> = HashMap::new();
+    let keyed_data = common::read_cgroup_file(file_path)?;
+    for entry in keyed_data.lines() {
+        let entry_fields: Vec<&str> = entry.split_ascii_whitespace().collect();
+        if entry_fields.len() < 2 {
+            continue;
+        }
+
+        stats.insert(
+            entry_fields[0].to_owned(),
+            entry_fields[1..]
+                .iter()
+                .copied()
+                .map(|p| p.to_owned())
+                .collect(),
+        );
+    }
+
+    Ok(stats)
+}
+
+pub fn parse_device_number(device: &str) -> Result<(u64, u64)> {
+    let numbers: Vec<&str> = device.split_terminator(':').collect();
+    if numbers.len() != 2 {
+        bail!("failed to parse device number {}", device);
+    }
+
+    Ok((numbers[0].parse()?, numbers[1].parse()?))
+}
+
 pub fn pid_stats(cgroup_path: &Path) -> Result<PidStats> {
     let mut stats = PidStats::default();
 

src/cgroups/v1/blkio.rs

@@ -2,7 +2,7 @@ use std::path::Path;
 
 use crate::cgroups::{
     common,
-    stats::{BlkioDeviceStat, BlkioStats, StatsProvider},
+    stats::{self, BlkioDeviceStat, BlkioStats, StatsProvider},
     v1::Controller,
 };
 use anyhow::{bail, Context, Result};
@@ -173,7 +173,7 @@ impl Blkio {
                 continue;
             }
 
-            let (major, minor) = Self::parse_device_number(entry_fields[0])?;
+            let (major, minor) = stats::parse_device_number(entry_fields[0])?;
             let op_type = if entry_fields.len() == 3 {
                 Some(entry_fields[1].to_owned())
             } else {
@@ -209,15 +209,6 @@ impl Blkio {
 
         Ok(stats)
     }
-
-    fn parse_device_number(entry: &str) -> Result<(u64, u64)> {
-        let numbers: Vec<&str> = entry.split_terminator(':').collect();
-        if numbers.len() != 2 {
-            bail!("failed to parse device number {}", entry);
-        }
-
-        Ok((numbers[0].parse()?, numbers[1].parse()?))
-    }
 }
 
 #[cfg(test)]

src/cgroups/v2/io.rs

@@ -2,13 +2,17 @@ use std::path::{Path, PathBuf};
 
 use anyhow::{bail, Result};
 
-use crate::cgroups::common;
+use crate::cgroups::{
+    common,
+    stats::{self, BlkioDeviceStat, BlkioStats, StatsProvider},
+};
 
 use super::controller::Controller;
 use oci_spec::{LinuxBlockIo, LinuxResources};
 
 const CGROUP_BFQ_IO_WEIGHT: &str = "io.bfq.weight";
 const CGROUP_IO_WEIGHT: &str = "io.weight";
+const CGROUP_IO_STAT: &str = "io.stat";
 
 pub struct Io {}
 
@@ -22,6 +26,58 @@ impl Controller for Io {
     }
 }
 
+impl StatsProvider for Io {
+    type Stats = BlkioStats;
+
+    fn stats(cgroup_path: &Path) -> Result<Self::Stats> {
+        let keyed_data = stats::parse_nested_keyed_data(&cgroup_path.join(CGROUP_IO_STAT))?;
+        let mut service_bytes = Vec::with_capacity(keyed_data.len());
+        let mut serviced = Vec::with_capacity(keyed_data.len());
+        for entry in keyed_data {
+            let (major, minor) = stats::parse_device_number(&entry.0)?;
+            for value in &entry.1 {
+                if value.starts_with("rbytes") {
+                    service_bytes.push(BlkioDeviceStat {
+                        major,
+                        minor,
+                        op_type: Some("read".to_owned()),
+                        value: stats::parse_value(&value[7..])?,
+                    });
+                } else if value.starts_with("wbytes") {
+                    service_bytes.push(BlkioDeviceStat {
+                        major,
+                        minor,
+                        op_type: Some("write".to_owned()),
+                        value: stats::parse_value(&value[7..])?,
+                    });
+                } else if value.starts_with("rios") {
+                    serviced.push(BlkioDeviceStat {
+                        major,
+                        minor,
+                        op_type: Some("read".to_owned()),
+                        value: stats::parse_value(&value[5..])?,
+                    });
+                } else if value.starts_with("wios") {
+                    serviced.push(BlkioDeviceStat {
+                        major,
+                        minor,
+                        op_type: Some("write".to_owned()),
+                        value: stats::parse_value(&value[5..])?,
+                    });
+                }
+            }
+        }
+
+        let stats = BlkioStats {
+            service_bytes,
+            serviced,
+            ..Default::default()
+        };
+
+        Ok(stats)
+    }
+}
+
 impl Io {
     fn io_max_path(path: &Path) -> PathBuf {
         path.join("io.max")

src/cgroups/v2/manager.rs

@@ -165,6 +165,7 @@ impl CgroupManager for Manager {
                 ControllerType::HugeTlb => stats.hugetlb = HugeTlb::stats(&self.full_path)?,
                 ControllerType::Pids => stats.pids = Pids::stats(&self.full_path)?,
                 ControllerType::Memory => stats.memory = Memory::stats(&self.full_path)?,
+                ControllerType::Io => stats.blkio = Io::stats(&self.full_path)?,
                 _ => continue,
             }
         }