From 2ecfa3d4cbc7bd6e33d208d04a714f90040122a5 Mon Sep 17 00:00:00 2001 From: khai96_ Date: Sun, 1 Dec 2024 17:59:00 +0700 Subject: [PATCH] feat!: allow user to set rayon threads BREAKING CHANGE: a new field named `threads` is added to `Args` and `Sub` --- exports/completion.bash | 6 +++++- exports/completion.elv | 1 + exports/completion.fish | 1 + exports/completion.ps1 | 1 + exports/completion.zsh | 1 + src/app.rs | 2 ++ src/app/sub.rs | 24 +++++++++++++++++------- src/args.rs | 6 ++++++ src/args/threads.rs | 39 +++++++++++++++++++++++++++++++++++++++ tests/cli_errors.rs | 1 + 10 files changed, 74 insertions(+), 8 deletions(-) create mode 100644 src/args/threads.rs diff --git a/exports/completion.bash b/exports/completion.bash index a1858ed4..39c2935b 100644 --- a/exports/completion.bash +++ b/exports/completion.bash @@ -19,7 +19,7 @@ _pdu() { case "${cmd}" in pdu) - opts="-h -V --json-input --json-output --bytes-format --top-down --align-right --quantity --max-depth --total-width --column-width --min-ratio --no-sort --silent-errors --progress --help --version [FILES]..." + opts="-h -V --json-input --json-output --bytes-format --top-down --align-right --quantity --max-depth --total-width --column-width --min-ratio --no-sort --silent-errors --progress --threads --help --version [FILES]..." if [[ ${cur} == -* || ${COMP_CWORD} -eq 1 ]] ; then COMPREPLY=( $(compgen -W "${opts}" -- "${cur}") ) return 0 @@ -49,6 +49,10 @@ _pdu() { COMPREPLY=($(compgen -f "${cur}")) return 0 ;; + --threads) + COMPREPLY=($(compgen -f "${cur}")) + return 0 + ;; *) COMPREPLY=() ;; diff --git a/exports/completion.elv b/exports/completion.elv index 643c8ea4..9655ea8a 100644 --- a/exports/completion.elv +++ b/exports/completion.elv @@ -24,6 +24,7 @@ set edit:completion:arg-completer[pdu] = {|@words| cand --total-width 'Width of the visualization' cand --column-width 'Maximum widths of the tree column and width of the bar column' cand --min-ratio 'Minimal size proportion required to appear' + cand --threads 'Set the maximum number of threads to spawn. Could be either "auto", "max", or a number' cand --json-input 'Read JSON data from stdin' cand --json-output 'Print JSON data instead of an ASCII chart' cand --top-down 'Print the tree top-down instead of bottom-up' diff --git a/exports/completion.fish b/exports/completion.fish index 5b171782..9532ca3d 100644 --- a/exports/completion.fish +++ b/exports/completion.fish @@ -4,6 +4,7 @@ complete -c pdu -l max-depth -d 'Maximum depth to display the data (must be grea complete -c pdu -l total-width -d 'Width of the visualization' -r complete -c pdu -l column-width -d 'Maximum widths of the tree column and width of the bar column' -r complete -c pdu -l min-ratio -d 'Minimal size proportion required to appear' -r +complete -c pdu -l threads -d 'Set the maximum number of threads to spawn. Could be either "auto", "max", or a number' -r complete -c pdu -l json-input -d 'Read JSON data from stdin' complete -c pdu -l json-output -d 'Print JSON data instead of an ASCII chart' complete -c pdu -l top-down -d 'Print the tree top-down instead of bottom-up' diff --git a/exports/completion.ps1 b/exports/completion.ps1 index 834c986f..ceb23c68 100644 --- a/exports/completion.ps1 +++ b/exports/completion.ps1 @@ -27,6 +27,7 @@ Register-ArgumentCompleter -Native -CommandName 'pdu' -ScriptBlock { [CompletionResult]::new('--total-width', '--total-width', [CompletionResultType]::ParameterName, 'Width of the visualization') [CompletionResult]::new('--column-width', '--column-width', [CompletionResultType]::ParameterName, 'Maximum widths of the tree column and width of the bar column') [CompletionResult]::new('--min-ratio', '--min-ratio', [CompletionResultType]::ParameterName, 'Minimal size proportion required to appear') + [CompletionResult]::new('--threads', '--threads', [CompletionResultType]::ParameterName, 'Set the maximum number of threads to spawn. Could be either "auto", "max", or a number') [CompletionResult]::new('--json-input', '--json-input', [CompletionResultType]::ParameterName, 'Read JSON data from stdin') [CompletionResult]::new('--json-output', '--json-output', [CompletionResultType]::ParameterName, 'Print JSON data instead of an ASCII chart') [CompletionResult]::new('--top-down', '--top-down', [CompletionResultType]::ParameterName, 'Print the tree top-down instead of bottom-up') diff --git a/exports/completion.zsh b/exports/completion.zsh index 46b01ab1..cd296575 100644 --- a/exports/completion.zsh +++ b/exports/completion.zsh @@ -25,6 +25,7 @@ block-count\:"Count numbers of blocks"))' \ '(--column-width)--total-width=[Width of the visualization]:TOTAL_WIDTH:_default' \ '*--column-width=[Maximum widths of the tree column and width of the bar column]:TREE_WIDTH:_default:TREE_WIDTH:_default' \ '--min-ratio=[Minimal size proportion required to appear]:MIN_RATIO:_default' \ +'--threads=[Set the maximum number of threads to spawn. Could be either "auto", "max", or a number]:THREADS:_default' \ '(--quantity)--json-input[Read JSON data from stdin]' \ '--json-output[Print JSON data instead of an ASCII chart]' \ '--top-down[Print the tree top-down instead of bottom-up]' \ diff --git a/src/app.rs b/src/app.rs index b00e1f93..b11b6a96 100644 --- a/src/app.rs +++ b/src/app.rs @@ -141,6 +141,7 @@ impl App { max_depth, min_ratio, no_sort, + threads, .. } => Sub { direction: Direction::from_top_down(top_down), @@ -154,6 +155,7 @@ impl App { max_depth, min_ratio, no_sort, + threads, } .run(), )*} }; diff --git a/src/app/sub.rs b/src/app/sub.rs index 1314e044..2073e973 100644 --- a/src/app/sub.rs +++ b/src/app/sub.rs @@ -2,7 +2,7 @@ mod hdd; mod mount_point; use crate::{ - args::Fraction, + args::{Fraction, Threads}, data_tree::{DataTree, DataTreeReflection}, fs_tree_builder::FsTreeBuilder, get_size::GetSize, @@ -47,6 +47,8 @@ where pub reporter: Report, /// Minimal size proportion required to appear. pub min_ratio: Fraction, + /// The number of threads [`rayon`] can use. + pub threads: Threads, /// Preserve order of entries. pub no_sort: bool, } @@ -71,18 +73,26 @@ where size_getter, reporter, min_ratio, + threads, no_sort, } = self; - // If one of the files is on HDD, set thread number to 1 - let disks = Disks::new_with_refreshed_list(); + let threads = match threads { + Threads::Auto => { + // If one of the files is on HDD, set thread number to 1 + let disks = Disks::new_with_refreshed_list(); + eprintln!("warning: HDD detected, the thread limit will be set to 1"); + any_path_is_in_hdd::(&files, &disks).then_some(1) + } + Threads::Max => None, + Threads::Fixed(threads) => Some(threads), + }; - if any_path_is_in_hdd::(&files, &disks) { - eprintln!("warning: HDD detected, the thread limit will be set to 1"); + if let Some(threads) = threads { rayon::ThreadPoolBuilder::new() - .num_threads(1) + .num_threads(threads) .build_global() - .unwrap_or_else(|_| eprintln!("warning: Failed to set thread limit to 1")); + .unwrap_or_else(|_| eprintln!("warning: Failed to set thread limit to {threads}")); } let mut iter = files diff --git a/src/args.rs b/src/args.rs index 192426e5..840e9dfa 100644 --- a/src/args.rs +++ b/src/args.rs @@ -1,8 +1,10 @@ pub mod fraction; pub mod quantity; +pub mod threads; pub use fraction::Fraction; pub use quantity::Quantity; +pub use threads::Threads; use crate::{bytes_format::BytesFormat, visualizer::ColumnWidthDistribution}; use clap::{ColorChoice, Parser}; @@ -128,6 +130,10 @@ pub struct Args { /// Report progress being made at the expense of performance. #[clap(long)] pub progress: bool, + + /// Set the maximum number of threads to spawn. Could be either "auto", "max", or a number. + #[clap(long, default_value_t = Threads::Auto)] + pub threads: Threads, } impl Args { diff --git a/src/args/threads.rs b/src/args/threads.rs new file mode 100644 index 00000000..dbf0c768 --- /dev/null +++ b/src/args/threads.rs @@ -0,0 +1,39 @@ +use derive_more::{Display, Error}; +use std::{num::ParseIntError, str::FromStr}; + +const AUTO: &str = "auto"; +const MAX: &str = "max"; + +/// Number of rayon threads. +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Display)] +pub enum Threads { + #[default] + #[display("{AUTO}")] + Auto, + #[display("{MAX}")] + Max, + Fixed(usize), +} + +/// Error that occurs when converting a string to an instance of [`Threads`]. +#[derive(Debug, Display, Clone, PartialEq, Eq, Error)] +pub enum FromStrError { + #[display("Value is neither {AUTO:?}, {MAX:?}, nor a number: {_0}")] + InvalidSyntax(ParseIntError), +} + +impl FromStr for Threads { + type Err = FromStrError; + fn from_str(value: &str) -> Result { + let value = value.trim(); + match value { + AUTO => return Ok(Threads::Auto), + MAX => return Ok(Threads::Max), + _ => {} + }; + value + .parse() + .map_err(FromStrError::InvalidSyntax) + .map(Threads::Fixed) + } +} diff --git a/tests/cli_errors.rs b/tests/cli_errors.rs index 35fd25ab..24a756bf 100644 --- a/tests/cli_errors.rs +++ b/tests/cli_errors.rs @@ -158,6 +158,7 @@ fn fs_errors() { let expected_stderr_lines = btreeset! { "[error] read_dir \"./nested/0\": Permission denied (os error 13)", "[error] read_dir \"./empty-dir\": Permission denied (os error 13)", + "warning: HDD detected, the thread limit will be set to 1", // TODO: fix this bug }; assert_eq!(actual_stderr_lines, expected_stderr_lines);