From 44ae12437dca951dd428bc6788e7bf4a9201400b Mon Sep 17 00:00:00 2001
From: Jeremy Howard <github@jhoward.fastmail.fm>
Date: Wed, 9 Sep 2026 13:16:33 +1000
Subject: [PATCH] Add explicit cancellation and worker joining for Rust streams

---
 DEV.md      |  2 ++
 src/walk.rs | 43 ++++++++++++++++++++++++++++++++++++++++---
 2 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/DEV.md b/DEV.md
index 385f03b..7fb7532 100644
--- a/DEV.md
+++ b/DEV.md
@@ -48,4 +48,6 @@ Truncation is recorded on collected results: `max_results` sets `stop_reason="ma
 
 Path results are `FileEntry` rows: a `str` subclass carrying the walk root, so paths stay plain strings for compatibility while stat info loads lazily (one cached `os.lstat` per entry, read only on attribute access). The wrapping happens at result construction on the Python side; Rust still streams plain strings. `PathResults.__repr__` shows an `ls -l`-style listing capped at `MAX_REPR` rows, so a huge result never stats everything, while `str()` stays one plain path per line. `ls` is `fd` with shell-style defaults (one level, dirs, ignore rules off), re-sorted with `stop_reason` preserved.
 
+Rust callers can consume a `StreamIter` with `cancel_and_join()` to cancel, drain queued results, and wait for the walk's workers to finish. `Drop` remains nonblocking. Since filesystem calls already in progress must return before joining completes, keep `cancel_and_join()` off async executors.
+
 This package intentionally has no CLI. Python is the interface.
diff --git a/src/walk.rs b/src/walk.rs
index 212bb72..8fae7f9 100644
--- a/src/walk.rs
+++ b/src/walk.rs
@@ -104,13 +104,23 @@ pub fn find_iter(opts: &FindOptions) -> Result<FindIter, RgApiError> {
     ))
 }
 
-pub struct StreamIter<T> { rx: mpsc::Receiver<Result<T, RgApiError>>, cancel: Arc<AtomicBool>, _worker: std::thread::JoinHandle<()> }
+pub struct StreamIter<T> { rx: mpsc::Receiver<Result<T, RgApiError>>, cancel: Arc<AtomicBool>, worker: Option<std::thread::JoinHandle<()>> }
 
 impl<T> StreamIter<T> {
     pub fn cancel(&self) { self.cancel.store(true, Ordering::Relaxed); }
 
     pub fn cancel_flag(&self) -> Arc<AtomicBool> { self.cancel.clone() }
 
+    /// Cancel and wait for the walk's workers to finish. Drain queued sends before
+    /// joining so a full result channel cannot deadlock shutdown. Unlike Drop,
+    /// this guarantees no background walk remains when it returns. Filesystem
+    /// calls already in progress must return first; run this off an async executor.
+    pub fn cancel_and_join(mut self) -> Result<(), RgApiError> {
+        self.cancel();
+        while self.rx.recv().is_ok() {}
+        self.worker.take().expect("stream owns its worker").join().map_err(|_| RgApiError::new("search worker panicked"))
+    }
+
     pub fn next_timeout(&mut self, timeout: std::time::Duration) -> Result<Result<T, RgApiError>, mpsc::RecvTimeoutError> { self.rx.recv_timeout(timeout) }
 
     /// Collect all items, stopping at `timeout_ms`; the bool is true when the deadline stopped it.
@@ -138,6 +148,33 @@ impl<T> Iterator for StreamIter<T> {
 
 impl<T> Drop for StreamIter<T> { fn drop(&mut self) { self.cancel(); } }
 
+#[cfg(test)]
+mod close_tests {
+    use super::*;
+
+    #[test]
+    fn cancel_and_join_drains_full_channel_and_waits_for_worker() {
+        let (tx, rx) = mpsc::sync_channel(1);
+        let (started, ready) = mpsc::channel();
+        let cancel = Arc::new(AtomicBool::new(false));
+        let worker_cancel = cancel.clone();
+        let done = Arc::new(AtomicBool::new(false));
+        let worker_done = done.clone();
+        let worker = std::thread::spawn(move || {
+            tx.send(Ok(1)).unwrap();
+            started.send(()).unwrap();
+            // This blocks while the channel is full; close must drain, not just join.
+            tx.send(Ok(2)).unwrap();
+            assert!(worker_cancel.load(Ordering::Acquire));
+            worker_done.store(true, Ordering::Release);
+        });
+        ready.recv().unwrap();
+        StreamIter { rx, cancel: cancel.clone(), worker: Some(worker) }.cancel_and_join().unwrap();
+        assert!(cancel.load(Ordering::Acquire));
+        assert!(done.load(Ordering::Acquire));
+    }
+}
+
 #[allow(clippy::too_many_arguments)]
 pub fn spawn_walk<T, F>(
     root: PathBuf,
@@ -181,7 +218,7 @@ where
             })
         });
     });
-    StreamIter { rx, cancel, _worker: worker }
+    StreamIter { rx, cancel, worker: Some(worker) }
 }
 
 fn find_entry(
@@ -365,7 +402,7 @@ mod tests {
                 if tx.send(Ok(i)).is_err() { return; }
             }
         });
-        StreamIter { rx, cancel, _worker: worker }
+        StreamIter { rx, cancel, worker: Some(worker) }
     }
 
     #[test]