feat!: allow user to set rayon threads (#270)

KSXGitHub · web-flow · commit d7b0fb615d86 · 2024-12-01T18:45:44.000+07:00
BREAKING CHANGE:
a new field named `threads` is added to `Args`.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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::<hdd::RealApi>(&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::<hdd::RealApi>(&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
@@ -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
@@ -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<Self, Self::Err> {
+        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)
+    }
+}
