1 use crate::{crossdev, InodeFilter, WalkOptions, WalkResult};
2 use anyhow::Result;
3 use colored::{Color, Colorize};
4 use filesize::PathExt;
5 use std::{borrow::Cow, io, path::Path};
6 #[cfg(feature = "aggregate-scan-progress")]
7 use std::{
8     sync::{
9         atomic::{AtomicU64, Ordering},
10         Arc,
11     },
12     thread,
13     time::Duration,
14 };
15 
16 /// Aggregate the given `paths` and write information about them to `out` in a human-readable format.
17 /// If `compute_total` is set, it will write an additional line with the total size across all given `paths`.
18 /// If `sort_by_size_in_bytes` is set, we will sort all sizes (ascending) before outputting them.
aggregate( mut out: impl io::Write, #[cfg_attr(not(feature = "aggregate-scan-progress"), allow(unused_variables))] err: Option< impl io::Write + Send + 'static, >, walk_options: WalkOptions, compute_total: bool, sort_by_size_in_bytes: bool, paths: impl IntoIterator<Item = impl AsRef<Path>>, ) -> Result<(WalkResult, Statistics)>19 pub fn aggregate(
20     mut out: impl io::Write,
21     #[cfg_attr(not(feature = "aggregate-scan-progress"), allow(unused_variables))] err: Option<
22         impl io::Write + Send + 'static,
23     >,
24     walk_options: WalkOptions,
25     compute_total: bool,
26     sort_by_size_in_bytes: bool,
27     paths: impl IntoIterator<Item = impl AsRef<Path>>,
28 ) -> Result<(WalkResult, Statistics)> {
29     let mut res = WalkResult::default();
30     let mut stats = Statistics {
31         smallest_file_in_bytes: u128::max_value(),
32         ..Default::default()
33     };
34     let mut total = 0;
35     let mut num_roots = 0;
36     let mut aggregates = Vec::new();
37     let mut inodes = InodeFilter::default();
38 
39     #[cfg(feature = "aggregate-scan-progress")]
40     let shared_count = Arc::new(AtomicU64::new(0));
41 
42     #[cfg(feature = "aggregate-scan-progress")]
43     if let Some(mut out) = err {
44         thread::spawn({
45             let shared_count = Arc::clone(&shared_count);
46             move || {
47                 thread::sleep(Duration::from_secs(1));
48                 loop {
49                     thread::sleep(Duration::from_millis(100));
50                     write!(
51                         out,
52                         "Enumerating {} entries\r",
53                         shared_count.load(Ordering::Acquire)
54                     )
55                     .ok();
56                 }
57             }
58         });
59     }
60 
61     for path in paths.into_iter() {
62         num_roots += 1;
63         let mut num_bytes = 0u128;
64         let mut num_errors = 0u64;
65         let device_id = crossdev::init(path.as_ref())?;
66         for entry in walk_options.iter_from_path(path.as_ref()) {
67             stats.entries_traversed += 1;
68             #[cfg(feature = "aggregate-scan-progress")]
69             shared_count.fetch_add(1, Ordering::Relaxed);
70             match entry {
71                 Ok(entry) => {
72                     let file_size = match entry.client_state {
73                         Some(Ok(ref m))
74                             if !m.is_dir()
75                                 && (walk_options.count_hard_links || inodes.add(m))
76                                 && (walk_options.cross_filesystems
77                                     || crossdev::is_same_device(device_id, m)) =>
78                         {
79                             if walk_options.apparent_size {
80                                 m.len()
81                             } else {
82                                 entry.path().size_on_disk_fast(m).unwrap_or_else(|_| {
83                                     num_errors += 1;
84                                     0
85                                 })
86                             }
87                         }
88                         Some(Ok(_)) => 0,
89                         Some(Err(_)) => {
90                             num_errors += 1;
91                             0
92                         }
93                         None => 0, // ignore directory
94                     } as u128;
95                     stats.largest_file_in_bytes = stats.largest_file_in_bytes.max(file_size);
96                     stats.smallest_file_in_bytes = stats.smallest_file_in_bytes.min(file_size);
97                     num_bytes += file_size;
98                 }
99                 Err(_) => num_errors += 1,
100             }
101         }
102 
103         if sort_by_size_in_bytes {
104             aggregates.push((path.as_ref().to_owned(), num_bytes, num_errors));
105         } else {
106             output_colored_path(
107                 &mut out,
108                 &walk_options,
109                 &path,
110                 num_bytes,
111                 num_errors,
112                 path_color_of(&path),
113             )?;
114         }
115         total += num_bytes;
116         res.num_errors += num_errors;
117     }
118 
119     if stats.entries_traversed == 0 {
120         stats.smallest_file_in_bytes = 0;
121     }
122 
123     if sort_by_size_in_bytes {
124         aggregates.sort_by_key(|&(_, num_bytes, _)| num_bytes);
125         for (path, num_bytes, num_errors) in aggregates.into_iter() {
126             output_colored_path(
127                 &mut out,
128                 &walk_options,
129                 &path,
130                 num_bytes,
131                 num_errors,
132                 path_color_of(&path),
133             )?;
134         }
135     }
136 
137     if num_roots > 1 && compute_total {
138         output_colored_path(
139             &mut out,
140             &walk_options,
141             Path::new("total"),
142             total,
143             res.num_errors,
144             None,
145         )?;
146     }
147     Ok((res, stats))
148 }
149 
path_color_of(path: impl AsRef<Path>) -> Option<Color>150 fn path_color_of(path: impl AsRef<Path>) -> Option<Color> {
151     if path.as_ref().is_file() {
152         None
153     } else {
154         Some(Color::Cyan)
155     }
156 }
157 
output_colored_path( out: &mut impl io::Write, options: &WalkOptions, path: impl AsRef<Path>, num_bytes: u128, num_errors: u64, path_color: Option<colored::Color>, ) -> std::result::Result<(), io::Error>158 fn output_colored_path(
159     out: &mut impl io::Write,
160     options: &WalkOptions,
161     path: impl AsRef<Path>,
162     num_bytes: u128,
163     num_errors: u64,
164     path_color: Option<colored::Color>,
165 ) -> std::result::Result<(), io::Error> {
166     writeln!(
167         out,
168         "{:>byte_column_width$} {}{}",
169         options
170             .byte_format
171             .display(num_bytes)
172             .to_string()
173             .as_str()
174             .green(),
175         {
176             let path = path.as_ref().display().to_string();
177             match path_color {
178                 Some(color) => path.color(color),
179                 None => path.normal(),
180             }
181         },
182         if num_errors == 0 {
183             Cow::Borrowed("")
184         } else {
185             Cow::Owned(format!(
186                 "  <{} IO Error{}>",
187                 num_errors,
188                 if num_errors > 1 { "s" } else { "" }
189             ))
190         },
191         byte_column_width = options.byte_format.width()
192     )
193 }
194 
195 /// Statistics obtained during a filesystem walk
196 #[derive(Default, Debug)]
197 pub struct Statistics {
198     /// The amount of entries we have seen during filesystem traversal
199     pub entries_traversed: u64,
200     /// The size of the smallest file encountered in bytes
201     pub smallest_file_in_bytes: u128,
202     /// The size of the largest file encountered in bytes
203     pub largest_file_in_bytes: u128,
204 }
205