Fully halt background scanner threads when dropping Worktree

* Rework fsevent API to expose a handle for halting the event stream

Co-Authored-By: Nathan Sobo <nathan@zed.dev>

fsevent/src/lib.rs

@@ -2,12 +2,14 @@
 
 use bitflags::bitflags;
 use fsevent_sys::{self as fs, core_foundation as cf};
+use parking_lot::Mutex;
 use std::{
     convert::AsRef,
     ffi::{c_void, CStr, OsStr},
     os::unix::ffi::OsStrExt,
     path::{Path, PathBuf},
     slice,
+    sync::Arc,
     time::Duration,
 };
 
@@ -18,20 +20,29 @@ pub struct Event {
     pub path: PathBuf,
 }
 
-pub struct EventStream<F> {
+pub struct EventStream {
     stream: fs::FSEventStreamRef,
-    _callback: Box<F>,
+    state: Arc<Mutex<Lifecycle>>,
+    callback: Box<Option<RunCallback>>,
 }
 
-unsafe impl<F> Send for EventStream<F> {}
+type RunCallback = Box<dyn FnMut(Vec<Event>) -> bool>;
 
-impl<F> EventStream<F>
-where
-    F: FnMut(Vec<Event>) -> bool,
-{
-    pub fn new(paths: &[&Path], latency: Duration, callback: F) -> Self {
+enum Lifecycle {
+    New,
+    Running(cf::CFRunLoopRef),
+    Stopped,
+}
+
+pub struct Handle(Arc<Mutex<Lifecycle>>);
+
+unsafe impl Send for EventStream {}
+unsafe impl Send for Lifecycle {}
+
+impl EventStream {
+    pub fn new(paths: &[&Path], latency: Duration) -> (Self, Handle) {
         unsafe {
-            let callback = Box::new(callback);
+            let callback = Box::new(None);
             let stream_context = fs::FSEventStreamContext {
                 version: 0,
                 info: callback.as_ref() as *const _ as *mut c_void,
@@ -71,20 +82,35 @@ where
             );
             cf::CFRelease(cf_paths);
 
-            EventStream {
-                stream,
-                _callback: callback,
-            }
+            let state = Arc::new(Mutex::new(Lifecycle::New));
+
+            (
+                EventStream {
+                    stream,
+                    state: state.clone(),
+                    callback,
+                },
+                Handle(state),
+            )
         }
     }
 
-    pub fn run(self) {
+    pub fn run<F>(mut self, f: F)
+    where
+        F: FnMut(Vec<Event>) -> bool + 'static,
+    {
+        *self.callback = Some(Box::new(f));
         unsafe {
-            fs::FSEventStreamScheduleWithRunLoop(
-                self.stream,
-                cf::CFRunLoopGetCurrent(),
-                cf::kCFRunLoopDefaultMode,
-            );
+            let run_loop = cf::CFRunLoopGetCurrent();
+            {
+                let mut state = self.state.lock();
+                match *state {
+                    Lifecycle::New => *state = Lifecycle::Running(run_loop),
+                    Lifecycle::Running(_) => unreachable!(),
+                    Lifecycle::Stopped => return,
+                }
+            }
+            fs::FSEventStreamScheduleWithRunLoop(self.stream, run_loop, cf::kCFRunLoopDefaultMode);
 
             fs::FSEventStreamStart(self.stream);
             cf::CFRunLoopRun();
@@ -107,7 +133,11 @@ where
             let event_paths = event_paths as *const *const ::std::os::raw::c_char;
             let e_ptr = event_flags as *mut u32;
             let i_ptr = event_ids as *mut u64;
-            let callback = (info as *mut F).as_mut().unwrap();
+            let callback = (info as *mut Option<RunCallback>)
+                .as_mut()
+                .unwrap()
+                .as_mut()
+                .unwrap();
 
             let paths = slice::from_raw_parts(event_paths, num);
             let flags = slice::from_raw_parts_mut(e_ptr, num);
@@ -136,6 +166,18 @@ where
     }
 }
 
+impl Drop for Handle {
+    fn drop(&mut self) {
+        let mut state = self.0.lock();
+        if let Lifecycle::Running(run_loop) = *state {
+            unsafe {
+                cf::CFRunLoopStop(run_loop);
+            }
+        }
+        *state = Lifecycle::Stopped;
+    }
+}
+
 // Synchronize with
 // /System/Library/Frameworks/CoreServices.framework/Versions/A/Frameworks/FSEvents.framework/Versions/A/Headers/FSEvents.h
 bitflags! {
@@ -253,10 +295,8 @@ fn test_event_stream() {
     fs::write(path.join("a"), "a contents").unwrap();
 
     let (tx, rx) = mpsc::channel();
-    let stream = EventStream::new(&[&path], Duration::from_millis(50), move |events| {
-        tx.send(events.to_vec()).is_ok()
-    });
-    std::thread::spawn(move || stream.run());
+    let (stream, handle) = EventStream::new(&[&path], Duration::from_millis(50));
+    std::thread::spawn(move || stream.run(move |events| tx.send(events.to_vec()).is_ok()));
 
     fs::write(path.join("b"), "b contents").unwrap();
     let events = rx.recv_timeout(Duration::from_millis(500)).unwrap();
@@ -269,4 +309,46 @@ fn test_event_stream() {
     let event = events.last().unwrap();
     assert_eq!(event.path, path.join("a"));
     assert!(event.flags.contains(StreamFlags::ITEM_REMOVED));
+    drop(handle);
+}
+
+#[test]
+fn test_event_stream_shutdown() {
+    use std::{fs, sync::mpsc, time::Duration};
+    use tempdir::TempDir;
+
+    let dir = TempDir::new("test_observe").unwrap();
+    let path = dir.path().canonicalize().unwrap();
+
+    let (tx, rx) = mpsc::channel();
+    let (stream, handle) = EventStream::new(&[&path], Duration::from_millis(50));
+    std::thread::spawn(move || {
+        stream.run({
+            let tx = tx.clone();
+            move |_| {
+                tx.send(()).unwrap();
+                true
+            }
+        });
+        tx.send(()).unwrap();
+    });
+
+    fs::write(path.join("b"), "b contents").unwrap();
+    rx.recv_timeout(Duration::from_millis(500)).unwrap();
+
+    drop(handle);
+    rx.recv_timeout(Duration::from_millis(500)).unwrap();
+}
+
+#[test]
+fn test_event_stream_shutdown_before_run() {
+    use std::time::Duration;
+    use tempdir::TempDir;
+
+    let dir = TempDir::new("test_observe").unwrap();
+    let path = dir.path().canonicalize().unwrap();
+
+    let (stream, handle) = EventStream::new(&[&path], Duration::from_millis(50));
+    drop(handle);
+    stream.run(|_| true);
 }