mprober_lib/volume/
volume.rs

1use std::{
2    collections::HashSet,
3    ffi::CString,
4    hash::{Hash, Hasher},
5    io::{self, ErrorKind},
6    mem::zeroed,
7    thread::sleep,
8    time::Duration,
9};
10
11use crate::{
12    scanner_rust::{ScannerAscii, ScannerError},
13    volume::{get_mounts, VolumeSpeed, VolumeStat},
14};
15
16#[derive(Debug, Clone, Eq)]
17pub struct Volume {
18    pub device: String,
19    pub stat:   VolumeStat,
20    pub size:   u64,
21    pub used:   u64,
22    pub points: Vec<String>,
23}
24
25impl Hash for Volume {
26    #[inline]
27    fn hash<H: Hasher>(&self, state: &mut H) {
28        self.device.hash(state)
29    }
30}
31
32impl PartialEq for Volume {
33    #[inline]
34    fn eq(&self, other: &Volume) -> bool {
35        self.device.eq(&other.device)
36    }
37}
38
39/// Get volume information by reading the `/proc/diskstats` file and using the `statvfs` function in libc.
40///
41/// ```rust
42/// use mprober_lib::volume;
43///
44/// let volumes = volume::get_volumes().unwrap();
45///
46/// println!("{volumes:#?}");
47/// ```
48pub fn get_volumes() -> Result<Vec<Volume>, ScannerError> {
49    let mut mounts = get_mounts()?;
50
51    let mut sc = ScannerAscii::scan_path("/proc/diskstats")?;
52
53    let mut volumes = Vec::with_capacity(1);
54
55    loop {
56        if sc.drop_next()?.is_none() {
57            break;
58        }
59
60        sc.drop_next()?.ok_or(ErrorKind::UnexpectedEof)?;
61
62        let device =
63            unsafe { String::from_utf8_unchecked(sc.next_raw()?.ok_or(ErrorKind::UnexpectedEof)?) };
64
65        if let Some(points) = mounts.remove(&device) {
66            for _ in 0..2 {
67                sc.drop_next()?.ok_or(ErrorKind::UnexpectedEof)?;
68            }
69
70            let read_bytes = sc.next_u64()?.ok_or(ErrorKind::UnexpectedEof)?;
71
72            for _ in 0..3 {
73                sc.drop_next()?.ok_or(ErrorKind::UnexpectedEof)?;
74            }
75
76            let write_bytes = sc.next_u64()?.ok_or(ErrorKind::UnexpectedEof)?;
77
78            for _ in 0..2 {
79                sc.drop_next()?.ok_or(ErrorKind::UnexpectedEof)?;
80            }
81
82            let time_spent = sc.next_u64()?.ok_or(ErrorKind::UnexpectedEof)?;
83
84            if time_spent > 0 {
85                let (size, used) = {
86                    let path = CString::new(points[0].as_bytes()).unwrap();
87
88                    let mut stats: libc::statvfs = unsafe { zeroed() };
89
90                    let rtn = unsafe { libc::statvfs(path.as_ptr(), &mut stats as *mut _) };
91
92                    if rtn != 0 {
93                        return Err(io::Error::last_os_error().into());
94                    }
95
96                    #[allow(clippy::unnecessary_cast)]
97                    (
98                        stats.f_bsize as u64 * stats.f_blocks as u64,
99                        stats.f_bsize as u64 * (stats.f_blocks - stats.f_bavail) as u64,
100                    )
101                };
102
103                let stat = VolumeStat {
104                    read_bytes,
105                    write_bytes,
106                };
107
108                let volume = Volume {
109                    device,
110                    stat,
111                    size,
112                    used,
113                    points,
114                };
115
116                volumes.push(volume);
117            }
118        }
119
120        sc.drop_next_line()?.ok_or(ErrorKind::UnexpectedEof)?;
121    }
122
123    Ok(volumes)
124}
125
126/// Get volume information by reading the `/proc/diskstats` file and using the `statvfs` function in libc. And measure the speed within a specific time interval.
127///
128/// ```rust
129/// use std::time::Duration;
130///
131/// use mprober_lib::volume;
132///
133/// let volumes_with_speed =
134///     volume::get_volumes_with_speed(Duration::from_millis(100)).unwrap();
135///
136/// for (volume, volume_with_speed) in volumes_with_speed {
137///     println!("{}: ", volume.device);
138///     println!("    Read: {:.1} B/s", volume_with_speed.read);
139///     println!("    Write: {:.1} B/s", volume_with_speed.write);
140/// }
141/// ```
142pub fn get_volumes_with_speed(
143    interval: Duration,
144) -> Result<Vec<(Volume, VolumeSpeed)>, ScannerError> {
145    let pre_volumes = get_volumes()?;
146
147    let pre_volumes_length = pre_volumes.len();
148
149    let mut pre_volumes_hashset = HashSet::with_capacity(pre_volumes_length);
150
151    for pre_volume in pre_volumes {
152        pre_volumes_hashset.insert(pre_volume);
153    }
154
155    sleep(interval);
156
157    let volumes = get_volumes()?;
158
159    let mut volumes_with_speed = Vec::with_capacity(volumes.len().min(pre_volumes_length));
160
161    for volume in volumes {
162        if let Some(pre_volume) = pre_volumes_hashset.get(&volume) {
163            let volume_speed = pre_volume.stat.compute_speed(&volume.stat, interval);
164
165            volumes_with_speed.push((volume, volume_speed));
166        }
167    }
168
169    Ok(volumes_with_speed)
170}