Add CLI option for the number of threads to spawn

This commit is contained in:
Xavier L'Heureux
2019-12-03 18:27:02 -05:00
parent 62ac9b623a
commit 19a41aa382
2 changed files with 38 additions and 18 deletions
+22 -14
View File
@@ -23,6 +23,13 @@ fn main() {
.help("Depth to show") .help("Depth to show")
.takes_value(true), .takes_value(true),
) )
.arg(
Arg::with_name("threads")
.short("t")
.long("threads")
.help("Number of threads to spawn simultaneously")
.takes_value(true),
)
.arg( .arg(
Arg::with_name("number_of_lines") Arg::with_name("number_of_lines")
.short("n") .short("n")
@@ -67,19 +74,20 @@ fn main() {
} }
}; };
let depth = { let threads = options.value_of("threads").and_then(|threads| {
if options.is_present("depth") { threads
match value_t!(options.value_of("depth"), u64) { .parse::<usize>()
Ok(v) => Some(v + 1), .map_err(|_| eprintln!("Ignoring bad value for threads: {:?}", threads))
Err(_) => { .ok()
eprintln!("Ignoring bad value for depth"); });
None
} let depth = options.value_of("depth").and_then(|depth| {
} depth
} else { .parse::<u64>()
None .map(|v| v + 1)
} .map_err(|_| eprintln!("Ignoring bad value for depth"))
}; .ok()
});
if options.is_present("depth") && number_of_lines != DEFAULT_NUMBER_OF_LINES { if options.is_present("depth") && number_of_lines != DEFAULT_NUMBER_OF_LINES {
eprintln!("Use either -n or -d. Not both"); eprintln!("Use either -n or -d. Not both");
return; return;
@@ -89,7 +97,7 @@ fn main() {
let use_full_path = options.is_present("display_full_paths"); let use_full_path = options.is_present("display_full_paths");
let simplified_dirs = simplify_dir_names(target_dirs); let simplified_dirs = simplify_dir_names(target_dirs);
let (permissions, nodes) = get_dir_tree(&simplified_dirs, use_apparent_size); let (permissions, nodes) = get_dir_tree(&simplified_dirs, use_apparent_size, threads);
let sorted_data = sort(nodes); let sorted_data = sort(nodes);
let biggest_ones = { let biggest_ones = {
match depth { match depth {
+16 -4
View File
@@ -44,13 +44,21 @@ pub fn simplify_dir_names(filenames: Vec<&str>) -> HashSet<String> {
pub fn get_dir_tree( pub fn get_dir_tree(
top_level_names: &HashSet<String>, top_level_names: &HashSet<String>,
apparent_size: bool, apparent_size: bool,
threads: Option<usize>,
) -> (bool, HashMap<String, u64>) { ) -> (bool, HashMap<String, u64>) {
let mut permissions = 0; let mut permissions = 0;
let mut inodes: HashSet<(u64, u64)> = HashSet::new(); let mut inodes: HashSet<(u64, u64)> = HashSet::new();
let mut data: HashMap<String, u64> = HashMap::new(); let mut data: HashMap<String, u64> = HashMap::new();
for b in top_level_names.iter() { for b in top_level_names.iter() {
examine_dir(&b, apparent_size, &mut inodes, &mut data, &mut permissions); examine_dir(
&b,
apparent_size,
&mut inodes,
&mut data,
&mut permissions,
threads,
);
} }
(permissions == 0, data) (permissions == 0, data)
} }
@@ -76,11 +84,15 @@ fn examine_dir(
inodes: &mut HashSet<(u64, u64)>, inodes: &mut HashSet<(u64, u64)>,
data: &mut HashMap<String, u64>, data: &mut HashMap<String, u64>,
file_count_no_permission: &mut u64, file_count_no_permission: &mut u64,
cpus: Option<usize>,
) { ) {
for entry in WalkDir::new(top_dir) let mut iter = WalkDir::new(top_dir)
.preload_metadata(true) .preload_metadata(true)
.skip_hidden(false) .skip_hidden(false);
{ if let Some(cpus) = cpus {
iter = iter.num_threads(cpus);
}
for entry in iter {
if let Ok(e) = entry { if let Ok(e) = entry {
let maybe_size_and_inode = get_metadata(&e, apparent_size); let maybe_size_and_inode = get_metadata(&e, apparent_size);