diff --git a/src/main.rs b/src/main.rs index 4cf6ba5..82d178f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -5,6 +5,8 @@ use std::io; use std::path::Path; use std::path::PathBuf; use std::process::exit; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use std::thread; use thiserror::Error; @@ -48,9 +50,12 @@ impl Options { } } -fn breadth_first_traverse(prog: &str, cwd: &Path, opt: &Options, in_queue: &PathQueue, out_queue: &PathQueue) -> Result<()> { +fn breadth_first_traverse(prog: &str, cwd: &Path, opt: &Options, in_queue: &PathQueue, out_queue: &PathQueue, counter: &AtomicUsize) -> Result<()> { loop { - let path = in_queue.pop_timeout(200)?; + if counter.load(Ordering::Acquire) == 0 { + break; + } + let path = in_queue.pop_timeout(5)?; if let Some(path) = path { let entries = fs::read_dir(&path); if let Ok(entries) = entries { @@ -93,6 +98,7 @@ fn breadth_first_traverse(prog: &str, cwd: &Path, opt: &Options, in_queue: &Path } if path.is_dir() { out_queue.push(path)?; + counter.fetch_add(1, Ordering::Release); } } else { eprintln!("{}: {}: cannot read filename", prog, path.display()); @@ -104,8 +110,7 @@ fn breadth_first_traverse(prog: &str, cwd: &Path, opt: &Options, in_queue: &Path } else { eprintln!("{}: {}: {}", prog, path.display(), entries.unwrap_err()); } - } else { - break; + counter.fetch_sub(1, Ordering::Release); } } Ok(()) @@ -227,11 +232,14 @@ fn main() { } } + let mut counter: usize = 0; + if roots.is_empty() { if let Err(e) = queues[0].push(PathBuf::from(".")) { eprintln!("{}: {}", prog, e); exit(1); } + counter = 1; } else { let dotdir = Path::new("."); let dotdotdir = Path::new(".."); @@ -265,20 +273,25 @@ fn main() { if let Err(e) = queues[0].push(path) { eprintln!("{}: {}", prog, e); exit(1) + } else { + counter += 1; } } } + let counter = AtomicUsize::new(counter); if let Err(e) = thread::scope(|s| -> Result<()> { let cwd = &cwd; let options = &options; let queues = &queues; + let counter = &counter; for i in 0..num_threads { s.spawn(move|| -> Result<()> { breadth_first_traverse( prog, cwd, options, &queues[i], - &queues[(i + 1) % num_threads]) + &queues[(i + 1) % num_threads], + counter) }); } Ok(()) diff --git a/src/path_queue.rs b/src/path_queue.rs index 6d19ff1..8bb3309 100644 --- a/src/path_queue.rs +++ b/src/path_queue.rs @@ -67,7 +67,7 @@ struct MemPathQueue { capacity: u32, pop_count: AtomicU32, push_count: AtomicU32, - buf: *mut PathBuf + buf: *mut PathBuf, } impl MemPathQueue {