Skip to main content

cargo_util/
du.rs

1//! A simple disk usage estimator.
2
3use std::path::Path;
4use std::sync::atomic::Ordering;
5use std::sync::{Arc, Mutex};
6
7use anyhow::{Context, Result};
8use ignore::overrides::OverrideBuilder;
9use ignore::{WalkBuilder, WalkState};
10use portable_atomic::AtomicU64;
11
12/// Determines the disk usage of all files in the given directory.
13///
14/// The given patterns are gitignore style patterns relative to the given
15/// path. If there are patterns, it will only count things matching that
16/// pattern. `!` can be used to exclude things. See [`OverrideBuilder::add`]
17/// for more info.
18///
19/// This is a primitive implementation that doesn't handle hard links, and
20/// isn't particularly fast (for example, not using `getattrlistbulk` on
21/// macOS). It also only uses actual byte sizes instead of block counts (and
22/// thus vastly undercounts directories with lots of small files). It would be
23/// nice to improve this or replace it with something better.
24pub fn du(path: &Path, patterns: &[&str]) -> Result<u64> {
25    du_inner(path, patterns).with_context(|| format!("failed to walk `{}`", path.display()))
26}
27
28fn du_inner(path: &Path, patterns: &[&str]) -> Result<u64> {
29    let mut builder = OverrideBuilder::new(path);
30    for pattern in patterns {
31        builder.add(pattern)?;
32    }
33    let overrides = builder.build()?;
34
35    let mut builder = WalkBuilder::new(path);
36    builder
37        .overrides(overrides)
38        .hidden(false)
39        .parents(false)
40        .ignore(false)
41        .git_global(false)
42        .git_ignore(false)
43        .git_exclude(false);
44    let walker = builder.build_parallel();
45
46    let total = Arc::new(AtomicU64::new(0));
47
48    // A slot used to indicate there was an error while walking.
49    //
50    // It is possible that more than one error happens (such as in different
51    // threads). The error returned is arbitrary in that case.
52    let err = Arc::new(Mutex::new(None));
53    walker.run(|| {
54        Box::new(|entry| {
55            match entry {
56                Ok(entry) => match entry.metadata() {
57                    Ok(meta) => {
58                        if meta.is_file() {
59                            total.fetch_add(meta.len(), Ordering::Relaxed);
60                        }
61                    }
62                    Err(e) => {
63                        *err.lock().unwrap() = Some(e.into());
64                        return WalkState::Quit;
65                    }
66                },
67                Err(e) => {
68                    *err.lock().unwrap() = Some(e.into());
69                    return WalkState::Quit;
70                }
71            }
72            WalkState::Continue
73        })
74    });
75
76    if let Some(e) = err.lock().unwrap().take() {
77        return Err(e);
78    }
79
80    Ok(total.load(Ordering::Relaxed))
81}