dupehell_core/
disk.rs

1// DupeHell -- MIT License
2//
3// Synthetic multi-domain dataset generator for record linkage benchmarking.
4// No liability for misuse.
5
6//! Free-space checks for the volume holding the output directory: an
7//! up-front comparison against the run's estimated output size (see
8//! `pipeline::estimate_output_bytes`) and a [`DiskGuard`] polled during
9//! generation, so a run stops cleanly instead of filling the disk.
10
11use std::path::{Path, PathBuf};
12
13use sysinfo::{DiskRefreshKind, Disks};
14
15/// Free space below which a running generation stops. Comfortably more
16/// than one batch's worth of output (a 500K-row IPC batch of the widest
17/// schemas is ~200 MB), so the stop happens before a write can fail.
18pub const MIN_FREE_BYTES: u64 = 2 * 1024 * 1024 * 1024;
19
20/// `path`, or its nearest existing ancestor (the output directory may not
21/// exist yet), made absolute — without Windows' `\\?\` verbatim prefix,
22/// which `Disk::mount_point` (`C:\`) never carries.
23fn resolve(path: &Path) -> Option<PathBuf> {
24    let mut p = path.to_path_buf();
25    loop {
26        if let Ok(c) = std::fs::canonicalize(&p) {
27            let s = c.to_string_lossy();
28            return Some(PathBuf::from(s.strip_prefix(r"\\?\").unwrap_or(&s)));
29        }
30        if !p.pop() {
31            return None;
32        }
33    }
34}
35
36/// Index of the disk whose mount point is the longest prefix of `path`.
37fn disk_index(disks: &Disks, path: &Path) -> Option<usize> {
38    disks
39        .iter()
40        .enumerate()
41        .filter(|(_, d)| path.starts_with(d.mount_point()))
42        .max_by_key(|(_, d)| d.mount_point().as_os_str().len())
43        .map(|(i, _)| i)
44}
45
46/// Free space (bytes) on the volume holding `path`, or `None` if it can't
47/// be determined (unusual filesystem, sandbox) — callers then skip the
48/// check rather than block a run on a guess.
49pub fn available_space(path: &Path) -> Option<u64> {
50    let path = resolve(path)?;
51    let disks = Disks::new_with_refreshed_list().ok()?;
52    disk_index(&disks, &path).map(|i| disks[i].available_space())
53}
54
55pub fn gb(bytes: u64) -> f64 {
56    bytes as f64 / 1e9
57}
58
59/// Polled during generation: fails once free space on the output volume
60/// drops below [`MIN_FREE_BYTES`]. Only that one disk's storage figures are
61/// refreshed per check, so polling it once per batch is cheap.
62pub struct DiskGuard {
63    disks: Disks,
64    index: usize,
65    path: PathBuf,
66}
67
68impl DiskGuard {
69    /// `None` when the output volume can't be identified — no guard then.
70    pub fn new(output_dir: &Path) -> Option<Self> {
71        let path = resolve(output_dir)?;
72        let disks = Disks::new_with_refreshed_list().ok()?;
73        let index = disk_index(&disks, &path)?;
74        Some(Self { disks, index, path })
75    }
76
77    pub fn check(&mut self) -> Result<(), String> {
78        let disk = &mut self.disks[self.index];
79        disk.refresh_specifics(DiskRefreshKind::nothing().with_storage());
80        let free = disk.available_space();
81        if free < MIN_FREE_BYTES {
82            return Err(format!(
83                "stopped: only {:.1} GB left on the disk holding {} (minimum {:.1} GB) — \
84                 free some space or choose another --output-dir. Files written so far \
85                 are incomplete.",
86                gb(free),
87                self.path.display(),
88                gb(MIN_FREE_BYTES)
89            ));
90        }
91        Ok(())
92    }
93}
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98
99    #[test]
100    fn test_available_space_on_temp_dir() {
101        // The temp dir's volume is always a real, listed disk.
102        let free = available_space(&std::env::temp_dir());
103        assert!(free.is_some_and(|f| f > 0));
104    }
105
106    #[test]
107    fn test_resolve_falls_back_to_existing_ancestor() {
108        let missing = std::env::temp_dir().join("dupehell_missing_dir_xyz/sub");
109        let resolved = resolve(&missing).unwrap();
110        assert!(!resolved.to_string_lossy().starts_with(r"\\?\"));
111        assert!(available_space(&missing).is_some());
112    }
113
114    #[test]
115    fn test_guard_passes_with_space_left() {
116        let mut guard = DiskGuard::new(&std::env::temp_dir()).unwrap();
117        // CI runners and dev machines keep well over 2 GB free on temp.
118        assert!(guard.check().is_ok());
119    }
120}
dupehell_core/disk.rs

dupehell_core/
disk.rs
