diff --git a/Cargo.lock b/Cargo.lock index a9534fc..95370b1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1,3 +1,5 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. [[package]] name = "ansi_term" version = "0.11.0" @@ -63,35 +65,12 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" dependencies = [ "crossbeam-epoch 0.6.1 (registry+https://github.com/rust-lang/crates.io-index)", - "crossbeam-utils 0.6.1 (registry+https://github.com/rust-lang/crates.io-index)", + "crossbeam-utils 0.6.5 (registry+https://github.com/rust-lang/crates.io-index)", "parking_lot 0.6.4 (registry+https://github.com/rust-lang/crates.io-index)", "rand 0.5.5 (registry+https://github.com/rust-lang/crates.io-index)", "smallvec 0.6.6 (registry+https://github.com/rust-lang/crates.io-index)", ] -[[package]] -name = "crossbeam-deque" -version = "0.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -dependencies = [ - "crossbeam-epoch 0.3.1 (registry+https://github.com/rust-lang/crates.io-index)", - "crossbeam-utils 0.2.2 (registry+https://github.com/rust-lang/crates.io-index)", -] - -[[package]] -name = "crossbeam-epoch" -version = "0.3.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -dependencies = [ - "arrayvec 0.4.8 (registry+https://github.com/rust-lang/crates.io-index)", - "cfg-if 0.1.6 (registry+https://github.com/rust-lang/crates.io-index)", - "crossbeam-utils 0.2.2 (registry+https://github.com/rust-lang/crates.io-index)", - "lazy_static 1.2.0 (registry+https://github.com/rust-lang/crates.io-index)", - "memoffset 0.2.1 (registry+https://github.com/rust-lang/crates.io-index)", - "nodrop 0.1.13 (registry+https://github.com/rust-lang/crates.io-index)", - "scopeguard 0.3.3 (registry+https://github.com/rust-lang/crates.io-index)", -] - [[package]] name = "crossbeam-epoch" version = "0.6.1" @@ -99,7 +78,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" dependencies = [ "arrayvec 0.4.8 (registry+https://github.com/rust-lang/crates.io-index)", "cfg-if 0.1.6 (registry+https://github.com/rust-lang/crates.io-index)", - "crossbeam-utils 0.6.1 (registry+https://github.com/rust-lang/crates.io-index)", + "crossbeam-utils 0.6.5 (registry+https://github.com/rust-lang/crates.io-index)", "lazy_static 1.2.0 (registry+https://github.com/rust-lang/crates.io-index)", "memoffset 0.2.1 (registry+https://github.com/rust-lang/crates.io-index)", "scopeguard 0.3.3 (registry+https://github.com/rust-lang/crates.io-index)", @@ -107,18 +86,11 @@ dependencies = [ [[package]] name = "crossbeam-utils" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -dependencies = [ - "cfg-if 0.1.6 (registry+https://github.com/rust-lang/crates.io-index)", -] - -[[package]] -name = "crossbeam-utils" -version = "0.6.1" +version = "0.6.5" source = "registry+https://github.com/rust-lang/crates.io-index" dependencies = [ "cfg-if 0.1.6 (registry+https://github.com/rust-lang/crates.io-index)", + "lazy_static 1.2.0 (registry+https://github.com/rust-lang/crates.io-index)", ] [[package]] @@ -127,16 +99,11 @@ version = "0.5.0" dependencies = [ "clap 2.32.0 (registry+https://github.com/rust-lang/crates.io-index)", "crossbeam-channel 0.3.2 (registry+https://github.com/rust-lang/crates.io-index)", + "crossbeam-utils 0.6.5 (registry+https://github.com/rust-lang/crates.io-index)", "humansize 1.1.0 (registry+https://github.com/rust-lang/crates.io-index)", "num_cpus 1.8.0 (registry+https://github.com/rust-lang/crates.io-index)", - "rayon 1.0.3 (registry+https://github.com/rust-lang/crates.io-index)", ] -[[package]] -name = "either" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" - [[package]] name = "fuchsia-zircon" version = "0.3.3" @@ -256,27 +223,6 @@ name = "rand_core" version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -[[package]] -name = "rayon" -version = "1.0.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -dependencies = [ - "crossbeam-deque 0.2.0 (registry+https://github.com/rust-lang/crates.io-index)", - "either 1.5.0 (registry+https://github.com/rust-lang/crates.io-index)", - "rayon-core 1.4.1 (registry+https://github.com/rust-lang/crates.io-index)", -] - -[[package]] -name = "rayon-core" -version = "1.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -dependencies = [ - "crossbeam-deque 0.2.0 (registry+https://github.com/rust-lang/crates.io-index)", - "lazy_static 1.2.0 (registry+https://github.com/rust-lang/crates.io-index)", - "libc 0.2.44 (registry+https://github.com/rust-lang/crates.io-index)", - "num_cpus 1.8.0 (registry+https://github.com/rust-lang/crates.io-index)", -] - [[package]] name = "redox_syscall" version = "0.1.42" @@ -424,12 +370,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" "checksum clap 2.32.0 (registry+https://github.com/rust-lang/crates.io-index)" = "b957d88f4b6a63b9d70d5f454ac8011819c6efa7727858f458ab71c756ce2d3e" "checksum cloudabi 0.0.3 (registry+https://github.com/rust-lang/crates.io-index)" = "ddfc5b9aa5d4507acaf872de71051dfd0e309860e88966e1051e462a077aac4f" "checksum crossbeam-channel 0.3.2 (registry+https://github.com/rust-lang/crates.io-index)" = "0ac88e108fa40799b39c08eb2a93bedf4cc99a9e5577f08ddf6dd6134ae65bf0" -"checksum crossbeam-deque 0.2.0 (registry+https://github.com/rust-lang/crates.io-index)" = "f739f8c5363aca78cfb059edf753d8f0d36908c348f3d8d1503f03d8b75d9cf3" -"checksum crossbeam-epoch 0.3.1 (registry+https://github.com/rust-lang/crates.io-index)" = "927121f5407de9956180ff5e936fe3cf4324279280001cd56b669d28ee7e9150" "checksum crossbeam-epoch 0.6.1 (registry+https://github.com/rust-lang/crates.io-index)" = "2449aaa4ec7ef96e5fb24db16024b935df718e9ae1cec0a1e68feeca2efca7b8" -"checksum crossbeam-utils 0.2.2 (registry+https://github.com/rust-lang/crates.io-index)" = "2760899e32a1d58d5abb31129f8fae5de75220bc2176e77ff7c627ae45c918d9" -"checksum crossbeam-utils 0.6.1 (registry+https://github.com/rust-lang/crates.io-index)" = "c55913cc2799171a550e307918c0a360e8c16004820291bf3b638969b4a01816" -"checksum either 1.5.0 (registry+https://github.com/rust-lang/crates.io-index)" = "3be565ca5c557d7f59e7cfcf1844f9e3033650c929c6566f511e8005f205c1d0" +"checksum crossbeam-utils 0.6.5 (registry+https://github.com/rust-lang/crates.io-index)" = "f8306fcef4a7b563b76b7dd949ca48f52bc1141aa067d2ea09565f3e2652aa5c" "checksum fuchsia-zircon 0.3.3 (registry+https://github.com/rust-lang/crates.io-index)" = "2e9763c69ebaae630ba35f74888db465e49e259ba1bc0eda7d06f4a067615d82" "checksum fuchsia-zircon-sys 0.3.3 (registry+https://github.com/rust-lang/crates.io-index)" = "3dcaa9ae7725d12cdb85b3ad99a434db70b468c09ded17e012d86b5c1010f7a7" "checksum humansize 1.1.0 (registry+https://github.com/rust-lang/crates.io-index)" = "b6cab2627acfc432780848602f3f558f7e9dd427352224b0d9324025796d2a5e" @@ -446,8 +388,6 @@ source = "registry+https://github.com/rust-lang/crates.io-index" "checksum rand 0.5.5 (registry+https://github.com/rust-lang/crates.io-index)" = "e464cd887e869cddcae8792a4ee31d23c7edd516700695608f5b98c67ee0131c" "checksum rand_core 0.2.2 (registry+https://github.com/rust-lang/crates.io-index)" = "1961a422c4d189dfb50ffa9320bf1f2a9bd54ecb92792fb9477f99a1045f3372" "checksum rand_core 0.3.0 (registry+https://github.com/rust-lang/crates.io-index)" = "0905b6b7079ec73b314d4c748701f6931eb79fd97c668caa3f1899b22b32c6db" -"checksum rayon 1.0.3 (registry+https://github.com/rust-lang/crates.io-index)" = "373814f27745b2686b350dd261bfd24576a6fb0e2c5919b3a2b6005f820b0473" -"checksum rayon-core 1.4.1 (registry+https://github.com/rust-lang/crates.io-index)" = "b055d1e92aba6877574d8fe604a63c8b5df60f60e5982bf7ccbb1338ea527356" "checksum redox_syscall 0.1.42 (registry+https://github.com/rust-lang/crates.io-index)" = "cf8fb82a4d1c9b28f1c26c574a5b541f5ffb4315f6c9a791fa47b6a04438fe93" "checksum redox_termios 0.1.1 (registry+https://github.com/rust-lang/crates.io-index)" = "7e891cfe48e9100a70a3b6eb652fef28920c117d366339687bd5576160db0f76" "checksum rustc_version 0.2.3 (registry+https://github.com/rust-lang/crates.io-index)" = "138e3e0acb6c9fb258b19b67cb8abd63c00679d2851805ea151465464fe9030a" diff --git a/Cargo.toml b/Cargo.toml index fa17741..18b2a19 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,8 +12,8 @@ version = "0.5.0" [dependencies] num_cpus = "1.0" humansize = "1.1" -rayon = "1.0" crossbeam-channel = "0.3.0" +crossbeam-utils = "0.6.5" [dependencies.clap] version = "2" diff --git a/src/lib.rs b/src/lib.rs index b0cf55e..e11fa64 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -13,9 +13,7 @@ //! ``` extern crate crossbeam_channel; -extern crate humansize; -extern crate num_cpus; -extern crate rayon; +extern crate crossbeam_utils; pub mod walk; diff --git a/src/main.rs b/src/main.rs index 9dc48d2..1a79eaf 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,18 +1,15 @@ #[macro_use] extern crate clap; -extern crate crossbeam_channel; +extern crate diskus; extern crate humansize; extern crate num_cpus; -extern crate rayon; - -mod walk; use std::path::PathBuf; use clap::{App, AppSettings, Arg}; use humansize::{file_size_opts, FileSize}; -use walk::Walk; +use diskus::Walk; fn print_result(size: u64) { println!( diff --git a/src/walk.rs b/src/walk.rs index 0125713..6599f36 100644 --- a/src/walk.rs +++ b/src/walk.rs @@ -1,64 +1,107 @@ +use std::borrow::Cow; use std::collections::HashSet; use std::fs; use std::os::unix::fs::MetadataExt; -use std::path::PathBuf; -use std::thread; +use std::path::{Path, PathBuf}; use crossbeam_channel as channel; - -use rayon; -use rayon::prelude::*; +use crossbeam_utils::thread as thread; #[derive(Eq, PartialEq, Hash)] struct UniqueID(u64, u64); -enum Message { +enum WalkMessage<'a> { + Entries(Cow<'a, [PathBuf]>), + Done, +} + +enum CountMessage { SizeEntry(Option, u64), + FinishedEntries(usize, usize), NoMetadataForPath(PathBuf), CouldNotReadDir(PathBuf), } -fn walk(tx: channel::Sender, entries: &[PathBuf]) { - entries.into_par_iter().for_each_with(tx, |tx_ref, entry| { - if let Ok(metadata) = entry.symlink_metadata() { - // If the entry has more than one hard link, generate - // a unique ID consisting of device and inode in order - // not to count this entry twice. - let unique_id = if metadata.is_file() && metadata.nlink() > 1 { - Some(UniqueID(metadata.dev(), metadata.ino())) - } else { - None - }; - - let size = metadata.len(); - - tx_ref.send(Message::SizeEntry(unique_id, size)).unwrap(); - - if metadata.is_dir() { - let mut children = vec![]; - match fs::read_dir(entry) { - Ok(child_entries) => { - for child_entry in child_entries { - if let Ok(child_entry) = child_entry { - children.push(child_entry.path()); - } +fn walk<'a>( + walk_rx: &'a channel::Receiver, + walk_tx: &'a channel::Sender, + count_tx: &'a channel::Sender, +) { + for msg in walk_rx { + match msg { + WalkMessage::Entries(entries) => { + let mut count = 0; + for entry in entries.iter() { + count += handle_entry(entry, walk_tx, count_tx); + } + count_tx + .send(CountMessage::FinishedEntries(entries.len(), count)) + .unwrap(); + } + WalkMessage::Done => { + // Send another "Done" message to ensure all the walkers exit + walk_tx + .send(WalkMessage::Done) + .unwrap(); + break; + } + } + } +} + +fn handle_entry<'a>( + entry: &Path, + walk_tx: &'a channel::Sender, + count_tx: &'a channel::Sender, +) -> usize { + let mut count = 0; + + if let Ok(metadata) = entry.symlink_metadata() { + // If the entry has more than one hard link, generate + // a unique ID consisting of device and inode in order + // not to count this entry twice. + let unique_id = if metadata.is_file() && metadata.nlink() > 1 { + Some(UniqueID(metadata.dev(), metadata.ino())) + } else { + None + }; + + let size = metadata.len(); + + count_tx + .send(CountMessage::SizeEntry(unique_id, size)) + .unwrap(); + + if metadata.is_dir() { + let mut children = vec![]; + match fs::read_dir(entry) { + Ok(child_entries) => { + for child_entry in child_entries { + if let Ok(child_entry) = child_entry { + children.push(child_entry.path()); } } - Err(_) => { - tx_ref - .send(Message::CouldNotReadDir(entry.clone())) - .unwrap(); - } } + Err(_) => { + count_tx + .send(CountMessage::CouldNotReadDir(entry.to_owned())) + .unwrap(); + } + } - walk(tx_ref.clone(), &children[..]); - }; - } else { - tx_ref - .send(Message::NoMetadataForPath(entry.clone())) + count = children.len(); + + walk_tx + .send(WalkMessage::Entries(Cow::from(children))) .unwrap(); }; - }); + } else { + count_tx + .send(CountMessage::NoMetadataForPath(entry.to_owned())) + .unwrap(); + }; + + count } pub struct Walk<'a> { @@ -75,14 +118,27 @@ impl<'a> Walk<'a> { } pub fn run(&self) -> u64 { - let (tx, rx) = channel::unbounded(); + let (count_tx, count_rx) = channel::unbounded(); + let (walk_tx, walk_rx) = channel::unbounded(); + + thread::scope(|s| { + // Ensure we create at least one thread so we don't get stuck + let thread_count = self.num_threads.max(1); + for _ in 0..thread_count { + s.spawn(|_| walk(&walk_rx, &walk_tx, &count_tx)); + } + + let root_entries = Cow::from(self.root_directories); + walk_tx + .send(WalkMessage::Entries(root_entries)) + .unwrap(); - let receiver_thread = thread::spawn(move || { let mut total = 0; + let mut entry_count = self.root_directories.len(); let mut ids = HashSet::new(); - for msg in rx { + for msg in &count_rx { match msg { - Message::SizeEntry(unique_id, size) => { + CountMessage::SizeEntry(unique_id, size) => { if let Some(unique_id) = unique_id { // Only count this entry if the ID has not been seen if ids.insert(unique_id) { @@ -92,13 +148,20 @@ impl<'a> Walk<'a> { total += size; } } - Message::NoMetadataForPath(path) => { + CountMessage::FinishedEntries(done, new) => { + entry_count = entry_count + new - done; + if entry_count == 0 { + // We have finished processing everything + break; + } + } + CountMessage::NoMetadataForPath(path) => { eprintln!( "diskus: could not retrieve metadata for path '{}'", path.to_string_lossy() ); } - Message::CouldNotReadDir(path) => { + CountMessage::CouldNotReadDir(path) => { eprintln!( "diskus: could not read contents of directory '{}'", path.to_string_lossy() @@ -107,15 +170,11 @@ impl<'a> Walk<'a> { } } - total - }); - - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(self.num_threads) - .build() - .unwrap(); - pool.install(|| walk(tx, self.root_directories)); + walk_tx + .send(WalkMessage::Done) + .unwrap(); - receiver_thread.join().unwrap() + total + }).unwrap() } }