diff --git a/Cargo.lock b/Cargo.lock index 30ff5d1..0aa4fdf 100644 --- a/Cargo.lock +++ b/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" diff --git a/Cargo.toml b/Cargo.toml index 55fa95c..3c4435b 100644 --- a/Cargo.toml +++ b/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" diff --git a/src/main.rs b/src/main.rs index 56d115e..a7d9a44 100644 --- a/src/main.rs +++ b/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 { - 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, } 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::(); + + 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::>(); + 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 + ); }