summaryrefslogtreecommitdiffstats
path: root/src/app/data_harvester/disks.rs
blob: c4a66eff694200c3735c36257587f73988ad85f9 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
//! Data collection about disks (e.g. I/O, usage, space).

use cfg_if::cfg_if;
use hashbrown::HashMap;

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;
        pub(crate) use self::unix::*;
    } else {
        mod other;
        pub(crate) use self::other::*;
    }
}

#[derive(Clone, Debug, Default)]
pub struct DiskHarvest {
    pub name: String,
    pub mount_point: String,

    /// Windows also contains an additional volume name field.
    #[cfg(target_os = "windows")]
    pub volume_name: Option<String>,

    // TODO: Maybe unify all these?
    pub free_space: Option<u64>,
    pub used_space: Option<u64>,
    pub total_space: Option<u64>,
}

#[derive(Clone, Debug)]
pub struct IoData {
    pub read_bytes: u64,
    pub write_bytes: u64,
}

pub type IoHarvest = HashMap<String, Option<IoData>>;

cfg_if! {
    if #[cfg(any(target_os = "linux", target_os = "macos", target_os = "windows"))] {
        mod io_counters;
        pub use io_counters::IoCounters;

        /// Returns the I/O usage of certain mount points.
        pub fn get_io_usage() -> anyhow::Result<IoHarvest> {
            let mut io_hash: HashMap<String, Option<IoData>> = HashMap::new();

            // TODO: Maybe rewrite this to not do a result of vec of result...
            for io in io_stats()?.into_iter() {
                let mount_point = io.device_name().to_string_lossy();

                io_hash.insert(
                    mount_point.to_string(),
                    Some(IoData {
                        read_bytes: io.read_bytes(),
                        write_bytes: io.write_bytes(),
                    }),
                );
            }

            Ok(io_hash)
        }
    } else if #[cfg(not(target_os = "freebsd"))] {
        pub fn get_io_usage() -> anyhow::Result<IoHarvest> {
            anyhow::bail!("Unsupported OS");
        }
    }
}

/// Whether to keep the current disk entry given the filters, disk name, and disk mount.
/// Precedence ordering in the case where name and mount filters disagree, "allow"
/// takes precedence over "deny".
///
/// For implementation, we do this as follows:
///
/// 1. Is the entry allowed through any filter? That is, does it match an entry in a
///    filter where `is_list_ignored` is `false`? If so, we always keep this entry.
/// 2. Is the entry denied through any filter? That is, does it match an entry in a
///    filter where `is_list_ignored` is `true`? If so, we always deny this entry.
/// 3. Anything else is allowed.
pub fn keep_disk_entry(
    disk_name: &str, mount_point: &str, disk_filter: &Option<Filter>, mount_filter: &Option<Filter>,
) -> bool {
    match (disk_filter, mount_filter) {
        (Some(d), Some(m)) => match (d.is_list_ignored, m.is_list_ignored) {
            (true, true) => !(d.has_match(disk_name) || m.has_match(mount_point)),
            (true, false) => {
                if m.has_match(mount_point) {
                    true
                } else {
                    d.keep_entry(disk_name)
                }
            }
            (false, true) => {
                if d.has_match(disk_name) {
                    true
                } else {
                    m.keep_entry(mount_point)
                }
            }
            (false, false) => d.has_match(disk_name) || m.has_match(mount_point),
        },
        (Some(d), None) => d.keep_entry(disk_name),
        (None, Some(m)) => m.keep_entry(mount_point),
        (None, None) => true,
    }
}

#[cfg(test)]
mod test {
    use regex::Regex;

    use super::keep_disk_entry;
    use crate::app::filter::Filter;

    fn run_filter(disk_filter: &Option<Filter>, mount_filter: &Option<Filter>) -> Vec<usize> {
        let targets = [
            ("/dev/nvme0n1p1", "/boot"),
            ("/dev/nvme0n1p2", "/"),
            ("/dev/nvme0n1p3", "/home"),
            ("/dev/sda1", "/mnt/test"),
            ("/dev/sda2", "/mnt/boot"),
        ];

        targets
            .into_iter()
            .enumerate()
            .filter_map(|(itx, (name, mount))| {
                if keep_disk_entry(name, mount, disk_filter, mount_filter) {
                    Some(itx)
                } else {
                    None
                }
            })
            .collect()
    }

    #[test]
    fn test_keeping_disk_entry() {
        let disk_ignore = Some(Filter {
            is_list_ignored: true,
            list: vec![Regex::new("nvme").unwrap()],
        });

        let disk_keep = Some(Filter {
            is_list_ignored: false,
            list: vec![Regex::new("nvme").unwrap()],
        });

        let mount_ignore = Some(Filter {
            is_list_ignored: true,
            list: vec![Regex::new("boot").unwrap()],
        });

        let mount_keep = Some(Filter {
            is_list_ignored: false,
            list: vec![Regex::new("boot").unwrap()],
        });

        assert_eq!(run_filter(&None, &None), vec![0, 1, 2, 3, 4]);

        assert_eq!(run_filter(&disk_ignore, &None), vec![3, 4]);
        assert_eq!(run_filter(&disk_keep, &None), vec![0, 1, 2]);

        assert_eq!(run_filter(&None, &mount_ignore), vec![1, 2, 3]);
        assert_eq!(run_filter(&None, &mount_keep), vec![0, 4]);

        assert_eq!(run_filter(&disk_ignore, &mount_ignore), vec![3]);
        assert_eq!(run_filter(&disk_keep, &mount_ignore), vec![0, 1, 2, 3]);

        assert_eq!(run_filter(&disk_ignore, &mount_keep), vec![0, 3, 4]);
        assert_eq!(run_filter(&disk_keep, &mount_keep), vec![0, 1, 2, 4]);
    }
}