workflow_perf_monitor/io/
mod.rs

1//! Get io usage for current process.
2use thiserror::Error;
3
4#[derive(Error, Debug)]
5#[error("IOStatsError({code}):{msg}")]
6pub struct IOStatsError {
7    pub code: i32,
8    pub msg: String,
9}
10
11impl From<std::io::Error> for IOStatsError {
12    fn from(e: std::io::Error) -> Self {
13        Self {
14            code: e.kind() as i32,
15            msg: e.to_string(),
16        }
17    }
18}
19
20impl From<std::num::ParseIntError> for IOStatsError {
21    fn from(e: std::num::ParseIntError) -> Self {
22        Self {
23            code: 0,
24            msg: e.to_string(),
25        }
26    }
27}
28/// A struct represents io status.
29#[derive(Debug, Clone, Default)]
30pub struct IOStats {
31    /// (linux & windows)  the number of read operations performed (cumulative)
32    pub read_count: u64,
33
34    /// (linux & windows) the number of write operations performed (cumulative)
35    pub write_count: u64,
36
37    /// the number of bytes read (cumulative).
38    pub read_bytes: u64,
39
40    /// the number of bytes written (cumulative)
41    pub write_bytes: u64,
42}
43/// Get the io stats of current process. Most platforms are supported.
44#[cfg(any(
45    target_os = "linux",
46    target_os = "android",
47    target_os = "macos",
48    target_os = "windows"
49))]
50pub fn get_process_io_stats() -> Result<IOStats, IOStatsError> {
51    get_process_io_stats_impl()
52}
53
54#[cfg(any(target_os = "linux", target_os = "android"))]
55fn get_process_io_stats_impl() -> Result<IOStats, IOStatsError> {
56    use std::{
57        io::{BufRead, BufReader},
58        str::FromStr,
59    };
60    let mut io_stats = IOStats::default();
61    let reader = BufReader::new(std::fs::File::open("/proc/self/io")?);
62
63    for line in reader.lines() {
64        let line = line?;
65        let mut s = line.split_whitespace();
66        if let (Some(field), Some(value)) = (s.next(), s.next()) {
67            match field {
68                "syscr:" => io_stats.read_count = u64::from_str(value)?,
69                "syscw:" => io_stats.write_count = u64::from_str(value)?,
70                "read_bytes:" => io_stats.read_bytes = u64::from_str(value)?,
71                "write_bytes:" => io_stats.write_bytes = u64::from_str(value)?,
72                _ => continue,
73            }
74        }
75    }
76
77    Ok(io_stats)
78}
79
80#[cfg(target_os = "windows")]
81fn get_process_io_stats_impl() -> Result<IOStats, IOStatsError> {
82    use std::mem::MaybeUninit;
83    use windows_sys::Win32::System::Threading::GetCurrentProcess;
84    use windows_sys::Win32::System::Threading::GetProcessIoCounters;
85    use windows_sys::Win32::System::Threading::IO_COUNTERS;
86    let mut io_counters = MaybeUninit::<IO_COUNTERS>::uninit();
87    let ret = unsafe {
88        // If the function succeeds, the return value is nonzero.
89        // If the function fails, the return value is zero.
90        // https://docs.microsoft.com/en-us/windows/win32/api/winbase/nf-winbase-getprocessiocounters
91        GetProcessIoCounters(GetCurrentProcess(), io_counters.as_mut_ptr())
92    };
93    if ret == 0 {
94        return Err(std::io::Error::last_os_error().into());
95    }
96    let io_counters = unsafe { io_counters.assume_init() };
97    Ok(IOStats {
98        read_count: io_counters.ReadOperationCount,
99        write_count: io_counters.WriteOperationCount,
100        read_bytes: io_counters.ReadTransferCount,
101        write_bytes: io_counters.WriteTransferCount,
102    })
103}
104
105#[cfg(target_os = "macos")]
106fn get_process_io_stats_impl() -> Result<IOStats, IOStatsError> {
107    use libc::{rusage_info_v2, RUSAGE_INFO_V2};
108    use std::{mem::MaybeUninit, os::raw::c_int};
109
110    let mut rusage_info_v2 = MaybeUninit::<rusage_info_v2>::uninit();
111    let ret_code = unsafe {
112        libc::proc_pid_rusage(
113            std::process::id() as c_int,
114            RUSAGE_INFO_V2,
115            rusage_info_v2.as_mut_ptr() as *mut _,
116        )
117    };
118    if ret_code != 0 {
119        return Err(std::io::Error::last_os_error().into());
120    }
121    let rusage_info_v2 = unsafe { rusage_info_v2.assume_init() };
122    Ok(IOStats {
123        read_bytes: rusage_info_v2.ri_diskio_bytesread,
124        write_bytes: rusage_info_v2.ri_diskio_byteswritten,
125        ..Default::default()
126    })
127}