Skip to content

Commit

Permalink
chore: Cleanup threads commit
Browse files Browse the repository at this point in the history
  • Loading branch information
bootandy committed May 6, 2024
1 parent 5b87260 commit e78690e
Show file tree
Hide file tree
Showing 4 changed files with 28 additions and 27 deletions.
2 changes: 1 addition & 1 deletion src/cli.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ pub fn build_cli() -> Command {
Arg::new("threads")
.short('T')
.long("threads")
.value_parser(value_parser!(u8))
.value_parser(value_parser!(usize))
.help("Number of threads to use")
.num_args(1)
)
Expand Down
9 changes: 9 additions & 0 deletions src/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ pub struct Config {
pub depth: Option<usize>,
pub bars_on_right: Option<bool>,
pub stack_size: Option<usize>,
pub threads: Option<usize>,
pub output_json: Option<bool>,
}

Expand Down Expand Up @@ -118,6 +119,14 @@ impl Config {
from_cmd_line.copied()
}
}
pub fn get_threads(&self, options: &ArgMatches) -> Option<usize> {
let from_cmd_line = options.get_one::<usize>("threads");
if from_cmd_line.is_none() {
self.threads
} else {
from_cmd_line.copied()
}
}
pub fn get_output_json(&self, options: &ArgMatches) -> bool {
Some(true) == self.output_json || options.get_flag("output_json")
}
Expand Down
43 changes: 18 additions & 25 deletions src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -228,7 +228,7 @@ fn main() {
progress_data: indicator.data.clone(),
errors: errors_for_rayon,
};
let threads_to_use: Option<u8> = options.get_one::<u8>("threads").copied();
let threads_to_use = config.get_threads(&options);
let stack_size = config.get_custom_stack_size(&options);
init_rayon(&stack_size, &threads_to_use);

Expand Down Expand Up @@ -312,7 +312,7 @@ fn main() {
}
}

fn init_rayon(stack_size: &Option<usize>, threads: &Option<u8>) {
fn init_rayon(stack_size: &Option<usize>, threads: &Option<usize>) {
// Rayon seems to raise this error on 32-bit builds
// The global thread pool has not been initialized.: ThreadPoolBuildError { kind: GlobalPoolAlreadyInitialized }
if cfg!(target_pointer_width = "64") {
Expand All @@ -332,38 +332,31 @@ fn output_json(output_filename: &str, top_level_nodes: &Vec<Node>) -> std::io::R

fn build_thread_pool(
stack: Option<usize>,
threads: Option<u8>,
threads: Option<usize>,
) -> Result<(), rayon::ThreadPoolBuildError> {
match stack {
Some(s) => match threads {
Some(t) => rayon::ThreadPoolBuilder::new()
.num_threads(t.into())
.stack_size(s)
.build_global(),
None => rayon::ThreadPoolBuilder::new().stack_size(s).build_global(),
},
let mut pool = rayon::ThreadPoolBuilder::new();

if let Some(thread_count) = threads {
pool = pool.num_threads(thread_count);
}

let stack_size = match stack {
Some(s) => Some(s),
None => {
let large_stack = usize::pow(1024, 3);
let mut s = System::new();
s.refresh_memory();
// Larger stack size if possible to handle cases with lots of nested directories
let available = s.available_memory();
if available > large_stack.try_into().unwrap() {
match threads {
Some(t) =>
// Larger stack size to handle cases with lots of nested directories
{
rayon::ThreadPoolBuilder::new()
.num_threads(t.into())
.stack_size(large_stack)
.build_global()
}
None => rayon::ThreadPoolBuilder::new()
.stack_size(large_stack)
.build_global(),
}
Some(large_stack)
} else {
rayon::ThreadPoolBuilder::new().build_global()
None
}
}
};
if let Some(stack_size_param) = stack_size {
pool = pool.stack_size(stack_size_param);
}
pool.build_global()
}
1 change: 0 additions & 1 deletion todo.txt

This file was deleted.

0 comments on commit e78690e

Please sign in to comment.