fix: iteration stop logic

brglng · Aug 19, 2024 · 7085571 · 7085571
1 parent 46f65d2
commit 7085571
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 6 deletions.
diff --git 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
@@ -67,7 +67,7 @@ struct MemPathQueue {
     capacity:       u32,
     pop_count:      AtomicU32,
     push_count:     AtomicU32,
-    buf:            *mut PathBuf
+    buf:            *mut PathBuf,
 }
 
 impl MemPathQueue {