Multi threading

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

Cargo.lock

@@ -98,12 +98,29 @@ version = "1.0.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b05b61dc5112cbb17e4b6cd61790d9845d13888356391624cbe7e41efeac1e75"
 
+[[package]]
+name = "crossbeam-channel"
+version = "0.5.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "82b8f8f868b36967f9606790d1903570de9ceaf870a7bf9fbbd3016d636a2cb2"
+dependencies = [
+ "crossbeam-utils",
+]
+
+[[package]]
+name = "crossbeam-utils"
+version = "0.8.21"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28"
+
 [[package]]
 name = "fdu"
 version = "0.1.0"
 dependencies = [
  "clap",
+ "crossbeam-channel",
  "humansize",
+ "num_cpus",
 ]
 
 [[package]]
@@ -112,6 +129,12 @@ version = "0.5.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea"
 
+[[package]]
+name = "hermit-abi"
+version = "0.5.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c"
+
 [[package]]
 name = "humansize"
 version = "2.1.3"
@@ -127,12 +150,28 @@ version = "1.70.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695"
 
+[[package]]
+name = "libc"
+version = "0.2.180"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bcc35a38544a891a5f7c865aca548a982ccb3b8650a5b06d0fd33a10283c56fc"
+
 [[package]]
 name = "libm"
 version = "0.2.16"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981"
 
+[[package]]
+name = "num_cpus"
+version = "1.17.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b"
+dependencies = [
+ "hermit-abi",
+ "libc",
+]
+
 [[package]]
 name = "once_cell_polyfill"
 version = "1.70.2"

Cargo.toml

@@ -5,4 +5,6 @@ edition = "2024"
 
 [dependencies]
 clap = { version = "4.5.54", features = ["derive"] }
+crossbeam-channel = "0.5.15"
 humansize = "2.1.3"
+num_cpus = "1.17.0"

src/main.rs

@@ -1,31 +1,108 @@
-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 {
+                    // wake everyone so they can exit
+                    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
+    );
 }