Multi threading

Signed-off-by: Slendi <slendi@socopon.com>

src/main.rs

@@ -1,31 +1,107 @@
-use std::path::PathBuf;
+use std::{
+    path::PathBuf,
+    sync::{
+        Arc,
+        atomic::{AtomicU64, AtomicUsize, Ordering},
+    },
+    thread,
+};
 
 use clap::Parser;
+use crossbeam_channel as chan;
 use humansize::DECIMAL;
 
 #[derive(Debug, Parser)]
 struct Args {
     path: PathBuf,
-}
-
-fn compute_size(path: PathBuf) -> Result<u64, std::io::Error> {
-    if path.is_dir() {
-        std::fs::read_dir(path)?.try_fold(0_u64, |counter, entry| {
-            Ok(counter + compute_size(entry?.path())?)
-        })
-    } else {
-        Ok(std::fs::metadata(path)?.len())
-    }
+    #[arg(short = 'j', long = "jobs")]
+    jobs: Option<usize>,
 }
 
 fn main() {
     let args = Args::parse();
-    match compute_size(args.path) {
-        Ok(size) => println!(
-            "Computed size sum: {} ({} bytes)",
-            humansize::format_size(size, DECIMAL),
-            size
-        ),
-        Err(error) => eprintln!("Failed to fetch size: {}", error),
+    let jobs = args.jobs.unwrap_or_else(num_cpus::get).max(1);
+
+    let (tx, rx) = chan::unbounded::<PathBuf>();
+
+    let pending = Arc::new(AtomicUsize::new(1));
+    let total = Arc::new(AtomicU64::new(0));
+
+    let mut workers = Vec::with_capacity(jobs);
+    for _ in 0..jobs {
+        let rx = rx.clone();
+        let total = total.clone();
+        workers.push(thread::spawn(move || {
+            while let Ok(path) = rx.recv() {
+                if let Ok(meta) = std::fs::metadata(&path) {
+                    if meta.is_file() {
+                        total.fetch_add(meta.len(), Ordering::Relaxed);
+                    }
+                }
+            }
+        }));
     }
+
+    let (dtx, drx) = chan::unbounded::<Option<PathBuf>>();
+    let _ = dtx.send(Some(args.path));
+
+    let mut walkers = Vec::with_capacity(jobs);
+    for _ in 0..jobs {
+        let dtx = dtx.clone();
+        let drx = drx.clone();
+        let tx = tx.clone();
+        let pending = pending.clone();
+
+        walkers.push(thread::spawn(move || {
+            while let Ok(msg) = drx.recv() {
+                let Some(dir) = msg else { break };
+
+                match std::fs::read_dir(&dir) {
+                    Ok(entries) => {
+                        for entry in entries.flatten() {
+                            let p = entry.path();
+                            match entry.file_type() {
+                                Ok(ft) if ft.is_dir() => {
+                                    pending.fetch_add(1, Ordering::Relaxed);
+                                    let _ = dtx.send(Some(p));
+                                }
+                                Ok(ft) if ft.is_file() => {
+                                    let _ = tx.send(p);
+                                }
+                                _ => {
+                                    let _ = tx.send(p);
+                                }
+                            }
+                        }
+                    }
+                    Err(_) => {}
+                }
+
+                if pending.fetch_sub(1, Ordering::AcqRel) == 1 {
+                    for _ in 0..jobs {
+                        let _ = dtx.send(None);
+                    }
+                    break;
+                }
+            }
+        }));
+    }
+
+    drop(dtx);
+
+    for w in walkers {
+        let _ = w.join();
+    }
+    drop(tx);
+
+    for w in workers {
+        let _ = w.join();
+    }
+
+    let size = total.load(Ordering::Relaxed);
+    println!(
+        "Computed size sum: {} ({} bytes)",
+        humansize::format_size(size, DECIMAL),
+        size
+    );
 }