dua/
aggregate.rs

1use crate::{ByteFormat, InodeFilter, Throttle, WalkOptions, WalkResult, WalkRoot, crossdev};
2use anyhow::Result;
3use filesize::PathExt;
4use owo_colors::{AnsiColors as Color, OwoColorize};
5use std::path::PathBuf;
6use std::time::Duration;
7use std::{io, path::Path};
8
9const CLEAR_CURRENT_LINE: &str = "\x1b[2K\r";
10
11/// Aggregate the given `paths` and write information about them to `out` in a human-readable format.
12/// If `compute_total` is set, it will write an additional line with the total size across all given `paths`.
13/// If `sort_by_size_in_bytes` is set, we will sort all sizes (ascending) before outputting them.
14pub fn aggregate(
15    mut out: impl io::Write,
16    mut err: Option<impl io::Write>,
17    walk_options: WalkOptions,
18    compute_total: bool,
19    sort_by_size_in_bytes: bool,
20    byte_format: ByteFormat,
21    paths: Vec<PathBuf>,
22) -> Result<(WalkResult, Statistics)> {
23    let mut res = WalkResult::default();
24    let mut stats = Statistics {
25        smallest_file_in_bytes: u128::MAX,
26        ..Default::default()
27    };
28    let num_roots = paths.len();
29    let mut aggregates = paths
30        .iter()
31        .cloned()
32        .map(|path| {
33            (
34                path,  // root path
35                0u128, // accumulated byte size
36                0u64,  // I/O error count
37            )
38        })
39        .collect::<Vec<_>>();
40    let mut device_ids = vec![0; num_roots];
41    let mut completed = vec![false; num_roots];
42    let mut roots = Vec::with_capacity(num_roots);
43    for (root_idx, path) in paths.into_iter().enumerate() {
44        let device_id = if walk_options.cross_filesystems {
45            0
46        } else {
47            crossdev::init(&path).unwrap_or(0)
48        };
49        device_ids[root_idx] = device_id;
50        roots.push(WalkRoot {
51            index: root_idx,
52            path,
53            device_id,
54        });
55    }
56    let mut inodes = InodeFilter::default();
57    let progress = Throttle::new(Duration::from_millis(100), Duration::from_secs(1).into());
58    let mut progress_visible = false;
59    let mut next_output = 0;
60
61    for (root_idx, event) in
62        walk_options.iter_from_paths(roots, false, crate::walk::Order::Completion)
63    {
64        let entry = match event {
65            crate::walk::RootEvent::Entry(entry) => entry,
66            crate::walk::RootEvent::Finished => {
67                completed[root_idx] = true;
68                if !sort_by_size_in_bytes {
69                    output_completed(
70                        &mut out,
71                        &mut err,
72                        &aggregates,
73                        &completed,
74                        &mut next_output,
75                        &mut progress_visible,
76                        byte_format,
77                    )?;
78                }
79                continue;
80            }
81        };
82        let (_, num_bytes, num_errors) = &mut aggregates[root_idx];
83        stats.entries_traversed += 1;
84        progress.throttled(|| {
85            if let Some(err) = err.as_mut() {
86                write!(err, "Enumerating {} items\r", stats.entries_traversed).ok();
87                progress_visible = true;
88            }
89        });
90        match entry {
91            Ok(entry) => {
92                let file_size = u128::from(match &entry.metadata {
93                    Ok(m)
94                        if (walk_options.count_hard_links || inodes.add(m))
95                            && (walk_options.cross_filesystems
96                                || crossdev::is_same_device(device_ids[root_idx], m)) =>
97                    {
98                        if walk_options.apparent_size {
99                            m.len()
100                        } else {
101                            entry.path().size_on_disk_fast(m).unwrap_or_else(|_| {
102                                *num_errors += 1;
103                                0
104                            })
105                        }
106                    }
107                    Ok(_) => 0,
108                    Err(_) => {
109                        *num_errors += 1;
110                        0
111                    }
112                });
113                stats.largest_file_in_bytes = stats.largest_file_in_bytes.max(file_size);
114                stats.smallest_file_in_bytes = stats.smallest_file_in_bytes.min(file_size);
115                *num_bytes += file_size;
116            }
117            Err(_) => *num_errors += 1,
118        }
119    }
120
121    let total = aggregates.iter().map(|(_, bytes, _)| bytes).sum();
122    res.num_errors = aggregates.iter().map(|(_, _, errors)| errors).sum();
123
124    if stats.entries_traversed == 0 {
125        stats.smallest_file_in_bytes = 0;
126    }
127
128    if progress_visible && let Some(err) = err.as_mut() {
129        write!(err, "{CLEAR_CURRENT_LINE}").ok();
130    }
131
132    if sort_by_size_in_bytes {
133        output_sorted(&mut out, aggregates, byte_format)?;
134    } else {
135        debug_assert_eq!(next_output, num_roots);
136    }
137
138    if num_roots > 1 && compute_total {
139        output_colored_path(
140            &mut out,
141            Path::new("total"),
142            total,
143            res.num_errors,
144            None,
145            byte_format,
146        )?;
147    }
148    Ok((res, stats))
149}
150
151/// Write the contiguous run of completed roots starting at `next_output`, preserving input order.
152/// Clears a visible progress line before writing the first completed root.
153/// `progress_visible` tracks if progress information is currently shown, taking up the last line.
154fn output_completed<W: io::Write, E: io::Write>(
155    out: &mut W,
156    err: &mut Option<E>,
157    aggregates: &[(std::path::PathBuf, u128, u64)],
158    completed: &[bool],
159    next_output: &mut usize,
160    progress_visible: &mut bool,
161    byte_format: ByteFormat,
162) -> io::Result<()> {
163    let must_report_completed_path = completed.get(*next_output).copied() == Some(true);
164    // Remove the transient progress line before writing permanent results to the terminal.
165    if must_report_completed_path && *progress_visible {
166        if let Some(err) = err.as_mut() {
167            write!(err, "{CLEAR_CURRENT_LINE}").ok();
168        }
169        *progress_visible = false;
170    }
171    while completed.get(*next_output).copied() == Some(true) {
172        let (path, num_bytes, num_errors) = &aggregates[*next_output];
173        output_colored_path(
174            out,
175            path,
176            *num_bytes,
177            *num_errors,
178            path_color_of(path),
179            byte_format,
180        )?;
181        *next_output += 1;
182    }
183    Ok(())
184}
185
186fn output_sorted(
187    out: &mut impl io::Write,
188    mut aggregates: Vec<(std::path::PathBuf, u128, u64)>,
189    byte_format: ByteFormat,
190) -> std::result::Result<(), io::Error> {
191    aggregates.sort_by_key(|&(_, num_bytes, _)| num_bytes);
192    for (path, num_bytes, num_errors) in aggregates {
193        output_colored_path(
194            out,
195            &path,
196            num_bytes,
197            num_errors,
198            path_color_of(&path),
199            byte_format,
200        )?;
201    }
202    Ok(())
203}
204
205fn path_color_of(path: impl AsRef<Path>) -> Option<Color> {
206    (!path.as_ref().is_file()).then_some(Color::Cyan)
207}
208
209fn output_colored_path(
210    out: &mut impl io::Write,
211    path: impl AsRef<Path>,
212    num_bytes: u128,
213    num_errors: u64,
214    path_color: Option<Color>,
215    byte_format: ByteFormat,
216) -> std::result::Result<(), io::Error> {
217    let size = byte_format.display(num_bytes).to_string();
218    let size = size.green();
219    let size_width = byte_format.width();
220    let path = path.as_ref().display();
221
222    let errors = if num_errors != 0 {
223        format!(
224            "  <{num_errors} IO Error{plural_s}>",
225            plural_s = if num_errors > 1 { "s" } else { "" }
226        )
227    } else {
228        String::new()
229    };
230
231    if let Some(color) = path_color {
232        writeln!(out, "{size:>size_width$} {}{errors}", path.color(color))
233    } else {
234        writeln!(out, "{size:>size_width$} {path}{errors}")
235    }
236}
237
238/// Statistics obtained during a filesystem walk
239#[derive(Default, Debug)]
240pub struct Statistics {
241    /// The amount of entries we have seen during filesystem traversal
242    pub entries_traversed: u64,
243    /// The size of the smallest file encountered in bytes
244    pub smallest_file_in_bytes: u128,
245    /// The size of the largest file encountered in bytes
246    pub largest_file_in_bytes: u128,
247}
248
249#[cfg(test)]
250mod tests {
251    use super::*;
252
253    #[test]
254    fn completed_roots_stream_in_input_order() {
255        let aggregates = [("first".into(), 1, 0), ("second".into(), 2, 0)];
256        let mut completed = [false, true];
257        let mut next_output = 0;
258        let mut progress_visible = true;
259        let mut out = Vec::new();
260        let mut err = Some(Vec::new());
261
262        output_completed(
263            &mut out,
264            &mut err,
265            &aggregates,
266            &completed,
267            &mut next_output,
268            &mut progress_visible,
269            ByteFormat::Bytes,
270        )
271        .unwrap();
272        assert!(
273            out.is_empty(),
274            "later roots must not overtake earlier roots"
275        );
276
277        completed[0] = true;
278        output_completed(
279            &mut out,
280            &mut err,
281            &aggregates,
282            &completed,
283            &mut next_output,
284            &mut progress_visible,
285            ByteFormat::Bytes,
286        )
287        .unwrap();
288
289        let out = String::from_utf8(out).unwrap();
290        assert!(
291            out.find("first").unwrap() < out.find("second").unwrap(),
292            "the first root is also emitted first"
293        );
294        assert_eq!(next_output, 2, "output stopped at root {next_output}");
295        assert_eq!(
296            err.as_deref(),
297            Some(CLEAR_CURRENT_LINE.as_bytes()),
298            "unexpected progress cleanup: {err:?}"
299        );
300        assert!(!progress_visible, "progress remained visible after cleanup");
301    }
302
303    #[test]
304    fn fast_roots_do_not_emit_terminal_erases() {
305        let dir = tempfile::tempdir().unwrap();
306        let paths = [dir.path().join("a"), dir.path().join("b")];
307        for path in &paths {
308            std::fs::write(path, []).unwrap();
309        }
310        let mut out = Vec::new();
311        let mut err = Vec::new();
312
313        aggregate(
314            &mut out,
315            Some(&mut err),
316            WalkOptions {
317                threads: 2,
318                count_hard_links: true,
319                apparent_size: false,
320                cross_filesystems: true,
321                ignore_dirs: std::collections::BTreeSet::default(),
322            },
323            true,
324            true,
325            ByteFormat::Metric,
326            paths.into(),
327        )
328        .unwrap();
329
330        assert!(
331            err.is_empty(),
332            "fast roots should not clear unseen progress"
333        );
334    }
335}
dua/aggregate.rs

dua/
aggregate.rs
