feat: Adding threads flag

Thanks: Dj-Codeman

10 files changed, 77 insertions, 26 deletions

diff --git a/completions/_dust b/completions/_dust
index 547d53b..4b07ba1 100644
--- a/completions/_dust
+++ b/completions/_dust
@@ -17,6 +17,8 @@ _dust() {
     _arguments "${_arguments_options[@]}" \
 '-d+[Depth to show]:DEPTH: ' \
 '--depth=[Depth to show]:DEPTH: ' \
+'-T+[Number of threads to use]: : ' \
+'--threads=[Number of threads to use]: : ' \
 '-n+[Number of lines of output to show. (Default is terminal_height - 10)]:NUMBER: ' \
 '--number-of-lines=[Number of lines of output to show. (Default is terminal_height - 10)]:NUMBER: ' \
 '*-X+[Exclude any file or directory with this name]:PATH:_files' \
diff --git a/completions/_dust.ps1 b/completions/_dust.ps1
index 7a7e502..467f4ff 100644
--- a/completions/_dust.ps1
+++ b/completions/_dust.ps1
@@ -23,6 +23,8 @@ Register-ArgumentCompleter -Native -CommandName 'dust' -ScriptBlock {
         'dust' {
             [CompletionResult]::new('-d', 'd', [CompletionResultType]::ParameterName, 'Depth to show')
             [CompletionResult]::new('--depth', 'depth', [CompletionResultType]::ParameterName, 'Depth to show')
+            [CompletionResult]::new('-T', 'T ', [CompletionResultType]::ParameterName, 'Number of threads to use')
+            [CompletionResult]::new('--threads', 'threads', [CompletionResultType]::ParameterName, 'Number of threads to use')
             [CompletionResult]::new('-n', 'n', [CompletionResultType]::ParameterName, 'Number of lines of output to show. (Default is terminal_height - 10)')
             [CompletionResult]::new('--number-of-lines', 'number-of-lines', [CompletionResultType]::ParameterName, 'Number of lines of output to show. (Default is terminal_height - 10)')
             [CompletionResult]::new('-X', 'X ', [CompletionResultType]::ParameterName, 'Exclude any file or directory with this name')
diff --git a/completions/dust.bash b/completions/dust.bash
index 1e709e5..334d3d4 100644
--- a/completions/dust.bash
+++ b/completions/dust.bash
@@ -19,7 +19,7 @@ _dust() {
 
     case "${cmd}" in
         dust)
-            opts="-d -n -p -X -I -L -x -s -r -c -C -b -B -z -R -f -i -v -e -t -w -P -D -F -o -S -j -h -V --depth --number-of-lines --full-paths --ignore-directory --ignore-all-in-file --dereference-links --limit-filesystem --apparent-size --reverse --no-colors --force-colors --no-percent-bars --bars-on-right --min-size --screen-reader --skip-total --filecount --ignore_hidden --invert-filter --filter --file_types --terminal_width --no-progress --only-dir --only-file --output-format --stack-size --output-json --help --version [PATH]..."
+            opts="-d -T -n -p -X -I -L -x -s -r -c -C -b -B -z -R -f -i -v -e -t -w -P -D -F -o -S -j -h -V --depth --threads --number-of-lines --full-paths --ignore-directory --ignore-all-in-file --dereference-links --limit-filesystem --apparent-size --reverse --no-colors --force-colors --no-percent-bars --bars-on-right --min-size --screen-reader --skip-total --filecount --ignore_hidden --invert-filter --filter --file_types --terminal_width --no-progress --only-dir --only-file --output-format --stack-size --output-json --help --version [PATH]..."
             if [[ ${cur} == -* || ${COMP_CWORD} -eq 1 ]] ; then
                 COMPREPLY=( $(compgen -W "${opts}" -- "${cur}") )
                 return 0
@@ -33,6 +33,14 @@ _dust() {
                     COMPREPLY=($(compgen -f "${cur}"))
                     return 0
                     ;;
+                --threads)
+                    COMPREPLY=($(compgen -f "${cur}"))
+                    return 0
+                    ;;
+                -T)
+                    COMPREPLY=($(compgen -f "${cur}"))
+                    return 0
+                    ;;
                 --number-of-lines)
                     COMPREPLY=($(compgen -f "${cur}"))
                     return 0
diff --git a/completions/dust.elv b/completions/dust.elv
index 572adb3..7b2283c 100644
--- a/completions/dust.elv
+++ b/completions/dust.elv
@@ -20,6 +20,8 @@ set edit:completion:arg-completer[dust] = {|@words|
         &'dust'= {
             cand -d 'Depth to show'
             cand --depth 'Depth to show'
+            cand -T 'Number of threads to use'
+            cand --threads 'Number of threads to use'
             cand -n 'Number of lines of output to show. (Default is terminal_height - 10)'
             cand --number-of-lines 'Number of lines of output to show. (Default is terminal_height - 10)'
             cand -X 'Exclude any file or directory with this name'
diff --git a/completions/dust.fish b/completions/dust.fish
index e68a2da..3853e4a 100644
--- a/completions/dust.fish
+++ b/completions/dust.fish
@@ -1,4 +1,5 @@
 complete -c dust -s d -l depth -d 'Depth to show' -r
+complete -c dust -s T -l threads -d 'Number of threads to use' -r
 complete -c dust -s n -l number-of-lines -d 'Number of lines of output to show. (Default is terminal_height - 10)' -r
 complete -c dust -s X -l ignore-directory -d 'Exclude any file or directory with this name' -r -F
 complete -c dust -s I -l ignore-all-in-file -d 'Exclude any file or directory with a regex matching that listed in this file, the file entries will be added to the ignore regexs provided by --invert_filter' -r -F
diff --git a/man-page/dust.1 b/man-page/dust.1
index 773f027..ad3d474 100644
--- a/man-page/dust.1
+++ b/man-page/dust.1
@@ -4,7 +4,7 @@
 .SH NAME
 Dust \- Like du but more intuitive
 .SH SYNOPSIS
-\fBdust\fR [\fB\-d\fR|\fB\-\-depth\fR] [\fB\-n\fR|\fB\-\-number\-of\-lines\fR] [\fB\-p\fR|\fB\-\-full\-paths\fR] [\fB\-X\fR|\fB\-\-ignore\-directory\fR] [\fB\-I\fR|\fB\-\-ignore\-all\-in\-file\fR] [\fB\-L\fR|\fB\-\-dereference\-links\fR] [\fB\-x\fR|\fB\-\-limit\-filesystem\fR] [\fB\-s\fR|\fB\-\-apparent\-size\fR] [\fB\-r\fR|\fB\-\-reverse\fR] [\fB\-c\fR|\fB\-\-no\-colors\fR] [\fB\-C\fR|\fB\-\-force\-colors\fR] [\fB\-b\fR|\fB\-\-no\-percent\-bars\fR] [\fB\-B\fR|\fB\-\-bars\-on\-right\fR] [\fB\-z\fR|\fB\-\-min\-size\fR] [\fB\-R\fR|\fB\-\-screen\-reader\fR] [\fB\-\-skip\-total\fR] [\fB\-f\fR|\fB\-\-filecount\fR] [\fB\-i\fR|\fB\-\-ignore_hidden\fR] [\fB\-v\fR|\fB\-\-invert\-filter\fR] [\fB\-e\fR|\fB\-\-filter\fR] [\fB\-t\fR|\fB\-\-file_types\fR] [\fB\-w\fR|\fB\-\-terminal_width\fR] [\fB\-P\fR|\fB\-\-no\-progress\fR] [\fB\-D\fR|\fB\-\-only\-dir\fR] [\fB\-F\fR|\fB\-\-only\-file\fR] [\fB\-o\fR|\fB\-\-output\-format\fR] [\fB\-S\fR|\fB\-\-stack\-size\fR] [\fB\-j\fR|\fB\-\-output\-json\fR] [\fB\-h\fR|\fB\-\-help\fR] [\fB\-V\fR|\fB\-\-version\fR] [\fIPATH\fR] 
+\fBdust\fR [\fB\-d\fR|\fB\-\-depth\fR] [\fB\-T\fR|\fB\-\-threads\fR] [\fB\-n\fR|\fB\-\-number\-of\-lines\fR] [\fB\-p\fR|\fB\-\-full\-paths\fR] [\fB\-X\fR|\fB\-\-ignore\-directory\fR] [\fB\-I\fR|\fB\-\-ignore\-all\-in\-file\fR] [\fB\-L\fR|\fB\-\-dereference\-links\fR] [\fB\-x\fR|\fB\-\-limit\-filesystem\fR] [\fB\-s\fR|\fB\-\-apparent\-size\fR] [\fB\-r\fR|\fB\-\-reverse\fR] [\fB\-c\fR|\fB\-\-no\-colors\fR] [\fB\-C\fR|\fB\-\-force\-colors\fR] [\fB\-b\fR|\fB\-\-no\-percent\-bars\fR] [\fB\-B\fR|\fB\-\-bars\-on\-right\fR] [\fB\-z\fR|\fB\-\-min\-size\fR] [\fB\-R\fR|\fB\-\-screen\-reader\fR] [\fB\-\-skip\-total\fR] [\fB\-f\fR|\fB\-\-filecount\fR] [\fB\-i\fR|\fB\-\-ignore_hidden\fR] [\fB\-v\fR|\fB\-\-invert\-filter\fR] [\fB\-e\fR|\fB\-\-filter\fR] [\fB\-t\fR|\fB\-\-file_types\fR] [\fB\-w\fR|\fB\-\-terminal_width\fR] [\fB\-P\fR|\fB\-\-no\-progress\fR] [\fB\-D\fR|\fB\-\-only\-dir\fR] [\fB\-F\fR|\fB\-\-only\-file\fR] [\fB\-o\fR|\fB\-\-output\-format\fR] [\fB\-S\fR|\fB\-\-stack\-size\fR] [\fB\-j\fR|\fB\-\-output\-json\fR] [\fB\-h\fR|\fB\-\-help\fR] [\fB\-V\fR|\fB\-\-version\fR] [\fIPATH\fR] 
 .SH DESCRIPTION
 Like du but more intuitive
 .SH OPTIONS
@@ -12,6 +12,9 @@ Like du but more intuitive
 \fB\-d\fR, \fB\-\-depth\fR=\fIDEPTH\fR
 Depth to show
 .TP
+\fB\-T\fR, \fB\-\-threads\fR
+Number of threads to use
+.TP
 \fB\-n\fR, \fB\-\-number\-of\-lines\fR=\fINUMBER\fR
 Number of lines of output to show. (Default is terminal_height \- 10)
 .TP
diff --git a/src/cli.rs b/src/cli.rs
index f79e129..fce5322 100644
--- a/src/cli.rs
+++ b/src/cli.rs
@@ -17,6 +17,14 @@ pub fn build_cli() -> Command {
                 .num_args(1)
         )
         .arg(
+            Arg::new("threads")
+                .short('T')
+                .long("threads")
+                .value_parser(value_parser!(u8))
+                .help("Number of threads to use")
+                .num_args(1)
+        )
+        .arg(
             Arg::new("number_of_lines")
                 .short('n')
                 .long("number-of-lines")
diff --git a/src/main.rs b/src/main.rs
index b2fc8a8..d55f6bf 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -228,8 +228,9 @@ fn main() {
         progress_data: indicator.data.clone(),
         errors: errors_for_rayon,
     };
+    let threads_to_use: Option<u8> = options.get_one::<u8>("threads").copied();
     let stack_size = config.get_custom_stack_size(&options);
-    init_rayon(&stack_size);
+    init_rayon(&stack_size, &threads_to_use);
 
     let top_level_nodes = walk_it(simplified_dirs, &walk_data);
 
@@ -311,32 +312,11 @@ fn main() {
     }
 }
 
-fn init_rayon(stack_size: &Option<usize>) {
+fn init_rayon(stack_size: &Option<usize>, threads: &Option<u8>) {
     // 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") {
-        let result = panic::catch_unwind(|| {
-            match stack_size {
-                Some(n) => rayon::ThreadPoolBuilder::new()
-                    .stack_size(*n)
-                    .build_global(),
-                None => {
-                    let large_stack = usize::pow(1024, 3);
-                    let mut s = System::new();
-                    s.refresh_memory();
-                    let available = s.available_memory();
-
-                    if available > large_stack.try_into().unwrap() {
-                        // Larger stack size to handle cases with lots of nested directories
-                        rayon::ThreadPoolBuilder::new()
-                            .stack_size(large_stack)
-                            .build_global()
-                    } else {
-                        rayon::ThreadPoolBuilder::new().build_global()
-                    }
-                }
-            }
-        });
+        let result = panic::catch_unwind(|| build_thread_pool(*stack_size, *threads));
         if result.is_err() {
             eprintln!("Problem initializing rayon, try: export RAYON_NUM_THREADS=1")
         }
@@ -349,3 +329,41 @@ fn output_json(output_filename: &str, top_level_nodes: &Vec<Node>) -> std::io::R
     file.write_all(serialized.as_bytes())?;
     Ok(())
 }
+
+fn build_thread_pool(
+    stack: Option<usize>,
+    threads: Option<u8>,
+) -> 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(),
+        },
+        None => {
+            let large_stack = usize::pow(1024, 3);
+            let mut s = System::new();
+            s.refresh_memory();
+            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(),
+                }
+            } else {
+                rayon::ThreadPoolBuilder::new().build_global()
+            }
+        }
+    }
+}
diff --git a/tests/test_flags.rs b/tests/test_flags.rs
index 0a1f56e..51e4068 100644
--- a/tests/test_flags.rs
+++ b/tests/test_flags.rs
@@ -60,6 +60,12 @@ pub fn test_d_flag_works() {
 }
 
 #[test]
+pub fn test_threads_flag_works() {
+    let output = build_command(vec!["-T", "1", "tests/test_dir/"]);
+    assert!(output.contains("hello_file"));
+}
+
+#[test]
 pub fn test_d_flag_works_and_still_recurses_down() {
     // We had a bug where running with '-d 1' would stop at the first directory and the code
     // would fail to recurse down
diff --git a/todo.txt b/todo.txt
new file mode 100644
index 0000000..84ed370
--- /dev/null
+++ b/todo.txt
@@ -0,0 +1 @@
+do not commit completions/_dust file