diff --git a/exports/completion.bash b/exports/completion.bash index a1858ed..39c2935 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 643c8ea..9655ea8 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 5b17178..9532ca3 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 834c986..ceb23c6 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 46b01ab..cd29657 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 bbd3c13..d63e1e4 100644 --- a/src/app.rs +++ b/src/app.rs @@ -3,7 +3,7 @@ pub mod sub; pub use sub::Sub; use crate::{ - args::{Args, Quantity}, + args::{Args, Quantity, Threads}, get_size::GetApparentSize, json_data::{JsonData, UnitAndTree}, reporter::{ErrorOnlyReporter, ErrorReport, ProgressAndErrorReporter, ProgressReport}, @@ -46,16 +46,25 @@ impl App { // // The other operations which are invoked frequently should not utilize dynamic dispatch. - { - // If one of the files is on HDD, set thread number to 1 - let disks = Disks::new_with_refreshed_list(); - if any_path_is_in_hdd::(&self.args.files, &disks) { - eprintln!("warning: HDD detected, the thread limit will be set to 1"); - rayon::ThreadPoolBuilder::new() - .num_threads(1) - .build_global() - .unwrap_or_else(|_| eprintln!("warning: Failed to set thread limit to 1")); + let threads = match self.args.threads { + Threads::Auto => { + let disks = Disks::new_with_refreshed_list(); + if any_path_is_in_hdd::(&self.args.files, &disks) { + eprintln!("warning: HDD detected, the thread limit will be set to 1"); + Some(1) + } else { + None + } } + Threads::Max => None, + Threads::Fixed(threads) => Some(threads), + }; + + if let Some(threads) = threads { + rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build_global() + .unwrap_or_else(|_| eprintln!("warning: Failed to set thread limit to {threads}")); } let column_width_distribution = self.args.column_width_distribution(); diff --git a/src/args.rs b/src/args.rs index 192426e..840e9df 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 0000000..dbf0c76 --- /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) + } +}