feat!: allow user to set rayon threads

BREAKING CHANGE:
a new field named `threads` is added to `Args` and `Sub`

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=()
                     ;;

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'

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'

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')

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]' \

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(),
             )*} };

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::<hdd::RealApi>(&files, &disks).then_some(1)
+            }
+            Threads::Max => None,
+            Threads::Fixed(threads) => Some(threads),
+        };
 
-        if any_path_is_in_hdd::<hdd::RealApi>(&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

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 {

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)
+    }
+}

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);