From 11459e79c4eb68385a33b2e705c692f38b139622 Mon Sep 17 00:00:00 2001
From: Seamile <lanhuermao@gmail.com>
Date: Fri, 6 Jan 2023 11:14:41 +0800
Subject: [PATCH] add cmd arg `n_thread`

---
 README.md     |  1 +
 src/main.rs   | 20 +++++++++++++++++---
 src/walker.rs | 13 ++++++++-----
 3 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/README.md b/README.md
index a708051..bc380c3 100644
--- a/README.md
+++ b/README.md
@@ -20,6 +20,7 @@ $ fcnt3 [OPTIONS] [DIRECTORIES]...
   * `-a`             count all regular and hidden files
   * `-s`             count the total size of files
   * `-R`             non-recursive mode (files in sub-directories will be ignored)
+  * `-t <N_THREAD>`  the number of threads for traversal (invalid in `non_recursive` mode)
   * `-h, --help`     Print help information
   * `-V, --version`  Print version information
 
diff --git a/src/main.rs b/src/main.rs
index 07c18ee..58fee9b 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,10 +1,12 @@
 mod output;
 mod walker;
 
-use clap::Parser;
 use std::path::PathBuf;
 use std::process::exit;
 
+use clap::Parser;
+use num_cpus::get as cpu_count;
+
 use walker::Counter;
 
 fn main() {
@@ -21,14 +23,22 @@ fn main() {
             };
         }
     } else {
-        counters = walker::parallel_walk(directories, args.all_files, args.count_size);
+        let n_thread: usize;
+        if args.n_thread == None {
+            let n_cpu = cpu_count();
+            n_thread = if n_cpu > 4 { n_cpu } else { 4 };
+        } else {
+            n_thread = args.n_thread.unwrap();
+        }
+
+        counters = walker::parallel_walk(directories, args.all_files, args.count_size, n_thread);
     }
     Counter::output(&counters, args.count_size);
 }
 
 #[derive(Parser)]
 #[command(name = "fcnt")]
-#[command(version = "0.1.0")]
+#[command(version = "0.2.0")]
 #[command(about = "Count the total number of files in given directories.")]
 struct CmdLineArgs {
     /// the directories (default: ./)
@@ -45,6 +55,10 @@ struct CmdLineArgs {
     /// non-recursive mode (files in sub-directories will be ignored).
     #[arg(short = 'R')]
     non_recursive: bool,
+
+    /// the number of threads for traversal (invalid in `non_recursive` mode).
+    #[arg(short = 't')]
+    n_thread: Option<usize>,
 }
 
 impl CmdLineArgs {
diff --git a/src/walker.rs b/src/walker.rs
index 4250aca..0a3610e 100644
--- a/src/walker.rs
+++ b/src/walker.rs
@@ -10,7 +10,6 @@ use std::thread;
 use std::time::Duration;
 
 use flume::unbounded as m_channel;
-use num_cpus::get as cpu_count;
 
 #[cfg(target_os = "linux")]
 use std::os::linux::fs::MetadataExt;
@@ -198,8 +197,12 @@ pub fn walk(dirpath: &PathBuf, with_hidden: bool, count_sz: bool) -> Result<DirD
 
 type Locker = Arc<Mutex<HashMap<usize, bool>>>;
 
-pub fn parallel_walk(dirlist: Vec<PathBuf>, with_hidden: bool, count_sz: bool) -> Vec<Counter> {
-    let n_threads = cpu_count();
+pub fn parallel_walk(
+    dirlist: Vec<PathBuf>,
+    with_hidden: bool,
+    count_sz: bool,
+    n_thread: usize,
+) -> Vec<Counter> {
     let (path_tx, path_rx) = m_channel::<PathBuf>();
     let (cnt_tx, cnt_rx) = s_channel::<Counter>();
     let mut counters = Vec::from_iter(dirlist.iter().map(|p| Counter::new(p)));
@@ -210,8 +213,8 @@ pub fn parallel_walk(dirlist: Vec<PathBuf>, with_hidden: bool, count_sz: bool) -
         path_tx.send(path.clone()).expect("path send err");
     }
 
-    // create walk threads which amount is n_threads
-    for t_idx in 0..n_threads {
+    // create walk threads which amount is n_thread
+    for t_idx in 0..n_thread {
         // clone channels
         let _path_tx = path_tx.clone();
         let _path_rx = path_rx.clone();