Skip to main content

dua/
common.rs

1use crate::{crossdev, walk};
2use byte_unit::{ByteUnit, n_gb_bytes, n_gib_bytes, n_mb_bytes, n_mib_bytes};
3use serde::Deserialize;
4use std::collections::BTreeSet;
5use std::path::PathBuf;
6use std::sync::Arc;
7use std::sync::atomic::{AtomicBool, Ordering};
8use std::time::Duration;
9use std::{fmt, path::Path};
10
11/// Specifies a way to format bytes
12#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize)]
13pub enum ByteFormat {
14    /// metric format, based on 1000.
15    #[serde(rename = "metric")]
16    Metric,
17    /// binary format, based on 1024
18    #[serde(rename = "binary")]
19    Binary,
20    /// raw bytes, without additional formatting
21    #[serde(rename = "bytes")]
22    Bytes,
23    /// only gigabytes without smart-unit
24    #[serde(rename = "gb")]
25    GB,
26    /// only gibibytes without smart-unit
27    #[serde(rename = "gib")]
28    GiB,
29    /// only megabytes without smart-unit
30    #[serde(rename = "mb")]
31    MB,
32    /// only mebibytes without smart-unit
33    #[serde(rename = "mib")]
34    MiB,
35}
36
37impl ByteFormat {
38    /// Return the content width (without unit suffix) needed to display values in this format.
39    #[must_use]
40    pub fn width(self) -> usize {
41        use ByteFormat::{Binary, Bytes, MB, MiB};
42        match self {
43            Binary => 11,
44            Bytes | MB | MiB => 12,
45            _ => 10,
46        }
47    }
48    /// Return the full width (value plus unit and separator) used by this format.
49    #[must_use]
50    pub fn total_width(self) -> usize {
51        use ByteFormat::{Binary, Bytes, GB, GiB, MB, Metric, MiB};
52        const THE_SPACE_BETWEEN_UNIT_AND_NUMBER: usize = 1;
53
54        self.width()
55            + match self {
56                Binary | MiB | GiB => 3,
57                Metric | MB | GB => 2,
58                Bytes => 1,
59            }
60            + THE_SPACE_BETWEEN_UNIT_AND_NUMBER
61    }
62    /// Create a display adapter for `bytes` using this format.
63    #[must_use]
64    pub fn display(self, bytes: u128) -> impl fmt::Display {
65        ByteFormatDisplay {
66            format: self,
67            bytes,
68        }
69    }
70}
71
72/// A lightweight display adapter created by [`ByteFormat::display`].
73struct ByteFormatDisplay {
74    format: ByteFormat,
75    bytes: u128,
76}
77
78#[expect(
79    clippy::cast_precision_loss,
80    reason = "byte_unit requires floating-point conversion"
81)]
82impl fmt::Display for ByteFormatDisplay {
83    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
84        use ByteFormat::{Binary, Bytes, GB, GiB, MB, Metric, MiB};
85        use byte_unit::Byte;
86
87        let format = match self.format {
88            Bytes => return write!(f, "{} b", self.bytes),
89            Binary => (true, None),
90            Metric => (false, None),
91            GB => (false, Some((n_gb_bytes!(1), ByteUnit::GB))),
92            GiB => (false, Some((n_gib_bytes!(1), ByteUnit::GiB))),
93            MB => (false, Some((n_mb_bytes!(1), ByteUnit::MB))),
94            MiB => (false, Some((n_mib_bytes!(1), ByteUnit::MiB))),
95        };
96
97        let b = match format {
98            (_, Some((divisor, unit))) => Byte::from_unit(self.bytes as f64 / divisor as f64, unit)
99                .expect("byte count > 0")
100                .get_adjusted_unit(unit),
101            (binary, None) => Byte::from_bytes(self.bytes).get_appropriate_unit(binary),
102        }
103        .format(2);
104        let mut splits = b.split(' ');
105        match (splits.next(), splits.next()) {
106            (Some(bytes), Some(unit)) => write!(
107                f,
108                "{} {:>unit_width$}",
109                bytes,
110                unit,
111                unit_width = match self.format {
112                    Binary => 3,
113                    _ => 2,
114                }
115            ),
116            _ => f.write_str(&b),
117        }
118    }
119}
120
121/// Throttle access to an optional `io::Write` to the specified `Duration`
122#[derive(Debug)]
123pub(crate) struct Throttle {
124    trigger: Arc<AtomicBool>,
125}
126
127impl Throttle {
128    /// Create a new throttle that allows updates at most once per `duration`.
129    ///
130    /// If `initial_sleep` is set, the first update is delayed by that amount.
131    pub(crate) fn new(duration: Duration, initial_sleep: Option<Duration>) -> Self {
132        let instance = Self {
133            trigger: Arc::default(),
134        };
135
136        let trigger = Arc::downgrade(&instance.trigger);
137        std::thread::spawn(move || {
138            if let Some(duration) = initial_sleep {
139                std::thread::sleep(duration);
140            }
141            while let Some(t) = trigger.upgrade() {
142                t.store(true, Ordering::Relaxed);
143                std::thread::sleep(duration);
144            }
145        });
146
147        instance
148    }
149
150    /// Execute `f` only if the throttle currently allows an update.
151    pub(crate) fn throttled<F>(&self, f: F)
152    where
153        F: FnOnce(),
154    {
155        if self.can_update() {
156            f();
157        }
158    }
159
160    /// Return `true` if we are not currently throttled.
161    pub(crate) fn can_update(&self) -> bool {
162        self.trigger.swap(false, Ordering::Relaxed)
163    }
164}
165
166/// Configures a filesystem walk, including output and formatting options.
167#[derive(Clone)]
168pub struct WalkOptions {
169    /// The amount of filesystem worker threads to use.
170    pub threads: usize,
171    /// If `true`, count every hard-link occurrence independently.
172    pub count_hard_links: bool,
173    /// If `true`, use apparent size (`metadata.len()`), not allocated blocks on disk.
174    pub apparent_size: bool,
175    /// If `false`, traversal is constrained to the root filesystem/device.
176    pub cross_filesystems: bool,
177    /// Canonicalized directories to skip from traversal.
178    pub ignore_dirs: BTreeSet<PathBuf>,
179}
180
181impl WalkOptions {
182    /// Create an iterator over `root` honoring this walk configuration.
183    ///
184    /// `root_device_id` is used to filter entries when `cross_filesystems == false`.
185    /// If `skip_root` is `true`, the root directory itself is omitted from yielded entries.
186    pub(crate) fn iter_from_path(
187        &self,
188        root: &Path,
189        root_device_id: u64,
190        skip_root: bool,
191        order: walk::Order,
192    ) -> impl Iterator<Item = std::io::Result<walk::Entry>> {
193        let ignore_dirs = self.ignore_dirs.clone();
194        let cwd = std::env::current_dir().unwrap_or_else(|_| root.to_owned());
195        let cross_filesystems = self.cross_filesystems;
196        walk::walk(root, self.threads, order, move |entry| {
197            (cross_filesystems
198                || entry.metadata.as_ref().map_or(true, |metadata| {
199                    crossdev::is_same_device(root_device_id, metadata)
200                }))
201                && (entry.depth == 0 || !ignore_directory(&entry.path(), &ignore_dirs, &cwd))
202        })
203        .filter(move |entry| !skip_root || entry.as_ref().map_or(true, |entry| entry.depth > 0))
204    }
205}
206
207/// Information we gather during a filesystem walk
208#[derive(Default)]
209pub struct WalkResult {
210    /// The amount of `io::errors` we encountered. Can happen when fetching meta-data, or when reading the directory contents.
211    pub num_errors: u64,
212}
213
214impl WalkResult {
215    /// Convert traversal result into a process exit code.
216    ///
217    /// Returns `0` if no I/O errors occurred, otherwise `1`.
218    #[must_use]
219    pub fn to_exit_code(&self) -> i32 {
220        i32::from(self.num_errors > 0)
221    }
222}
223
224/// Canonicalize user-provided ignore directory paths.
225///
226/// Non-canonicalizable paths are ignored.
227pub fn canonicalize_ignore_dirs(ignore_dirs: &[PathBuf]) -> BTreeSet<PathBuf> {
228    let dirs = ignore_dirs
229        .iter()
230        .map(gix::path::realpath)
231        .filter_map(Result::ok)
232        .collect();
233    log::info!("Ignoring canonicalized {dirs:?}");
234    dirs
235}
236
237fn ignore_directory(path: &Path, ignore_dirs: &BTreeSet<PathBuf>, cwd: &Path) -> bool {
238    if ignore_dirs.is_empty() {
239        return false;
240    }
241    let path = gix::path::realpath_opts(path, cwd, 32);
242    path.is_ok_and(|path| {
243        let ignored = ignore_dirs.contains(&path);
244        if ignored {
245            log::debug!("Ignored {}", path.display());
246        }
247        ignored
248    })
249}
250
251#[cfg(test)]
252mod tests {
253    use super::*;
254
255    #[test]
256    fn test_ignore_directories() {
257        let cwd = std::env::current_dir().unwrap();
258        #[cfg(unix)]
259        let mut parameters = vec![
260            ("/usr", vec!["/usr"], true),
261            ("/usr/local", vec!["/usr"], false),
262            ("/smth", vec!["/usr"], false),
263            ("/usr/local/..", vec!["/usr/local/.."], true),
264            ("/usr", vec!["/usr/local/.."], true),
265            ("/usr/local/share/../..", vec!["/usr"], true),
266        ];
267
268        #[cfg(windows)]
269        let mut parameters = vec![
270            ("C:\\Windows", vec!["C:\\Windows"], true),
271            ("C:\\Windows\\System", vec!["C:\\Windows"], false),
272            ("C:\\Smth", vec!["C:\\Windows"], false),
273            (
274                "C:\\Windows\\System\\..",
275                vec!["C:\\Windows\\System\\.."],
276                true,
277            ),
278            ("C:\\Windows", vec!["C:\\Windows\\System\\.."], true),
279            (
280                "C:\\Windows\\System\\Speech\\..\\..",
281                vec!["C:\\Windows"],
282                true,
283            ),
284        ];
285
286        parameters.extend([
287            ("src", vec!["src"], true),
288            ("src/interactive", vec!["src"], false),
289            ("src/interactive/..", vec!["src"], true),
290        ]);
291
292        for (path, ignore_dirs, expected_result) in parameters {
293            let ignore_dirs = canonicalize_ignore_dirs(
294                &ignore_dirs.into_iter().map(Into::into).collect::<Vec<_>>(),
295            );
296            assert_eq!(
297                ignore_directory(path.as_ref(), &ignore_dirs, &cwd),
298                expected_result,
299                "result='{expected_result}' for path='{path}' and ignore_dir='{ignore_dirs:?}' "
300            );
301        }
302    }
303
304    #[test]
305    fn explicitly_selected_ignored_root_is_traversed() {
306        let root = tempfile::tempdir().unwrap();
307        let child = root.path().join("child");
308        std::fs::create_dir(&child).unwrap();
309        let options = WalkOptions {
310            threads: 2,
311            count_hard_links: false,
312            apparent_size: false,
313            cross_filesystems: true,
314            ignore_dirs: canonicalize_ignore_dirs(&[root.path().to_owned()]),
315        };
316
317        let paths = options
318            .iter_from_path(
319                root.path(),
320                crossdev::init(root.path()).unwrap(),
321                false,
322                walk::Order::Completion,
323            )
324            .map(|entry| entry.unwrap().path())
325            .collect::<Vec<_>>();
326
327        assert!(paths.contains(&child));
328    }
329}