[red-knot] Watch search paths

crates/red_knot/src/main.rs

@@ -12,7 +12,7 @@ use tracing_tree::time::Uptime;
 
 use red_knot::db::RootDatabase;
 use red_knot::watch;
-use red_knot::watch::Watcher;
+use red_knot::watch::WorkspaceWatcher;
 use red_knot::workspace::WorkspaceMetadata;
 use ruff_db::program::{ProgramSettings, SearchPathSettings};
 use ruff_db::system::{OsSystem, System, SystemPathBuf};
@@ -142,7 +142,7 @@ struct MainLoop {
     receiver: crossbeam_channel::Receiver<MainLoopMessage>,
 
     /// The file system watcher, if running in watch mode.
-    watcher: Option<Watcher>,
+    watcher: Option<WorkspaceWatcher>,
 
     verbosity: Option<VerbosityLevel>,
 }
@@ -164,26 +164,23 @@ impl MainLoop {
 
     fn watch(mut self, db: &mut RootDatabase) -> anyhow::Result<()> {
         let sender = self.sender.clone();
-        let mut watcher = watch::directory_watcher(move |event| {
+        let watcher = watch::directory_watcher(move |event| {
             sender.send(MainLoopMessage::ApplyChanges(event)).unwrap();
         })?;
 
-        watcher.watch(db.workspace().root(db))?;
-
-        self.watcher = Some(watcher);
-
+        self.watcher = Some(WorkspaceWatcher::new(watcher, db));
         self.run(db);
 
         Ok(())
     }
 
     #[allow(clippy::print_stderr)]
-    fn run(self, db: &mut RootDatabase) {
+    fn run(mut self, db: &mut RootDatabase) {
         // Schedule the first check.
         self.sender.send(MainLoopMessage::CheckWorkspace).unwrap();
         let mut revision = 0usize;
 
-        for message in &self.receiver {
+        while let Ok(message) = self.receiver.recv() {
             tracing::trace!("Main Loop: Tick");
 
             match message {
@@ -224,6 +221,9 @@ impl MainLoop {
                     revision += 1;
                     // Automatically cancels any pending queries and waits for them to complete.
                     db.apply_changes(changes);
+                    if let Some(watcher) = self.watcher.as_mut() {
+                        watcher.update(db);
+                    }
                     self.sender.send(MainLoopMessage::CheckWorkspace).unwrap();
                 }
                 MainLoopMessage::Exit => {

crates/red_knot/src/watch.rs

@@ -1,7 +1,9 @@
 use ruff_db::system::{SystemPath, SystemPathBuf};
 pub use watcher::{directory_watcher, EventHandler, Watcher};
+pub use workspace_watcher::WorkspaceWatcher;
 
 mod watcher;
+mod workspace_watcher;
 
 /// Classification of a file system change event.
 ///

crates/red_knot/src/watch/workspace_watcher.rs

@@ -0,0 +1,107 @@
+use crate::db::RootDatabase;
+use crate::watch::Watcher;
+use ruff_db::system::SystemPathBuf;
+use rustc_hash::FxHashSet;
+use std::fmt::{Formatter, Write};
+use tracing::info;
+
+/// Wrapper around a [`Watcher`] that watches the relevant paths of a workspace.
+pub struct WorkspaceWatcher {
+    watcher: Watcher,
+
+    /// The paths that need to be watched. This includes paths for which setting up file watching failed.
+    watched_paths: FxHashSet<SystemPathBuf>,
+
+    /// Paths that should be watched but setting up the watcher failed for some reason.
+    /// This should be rare.
+    errored_paths: Vec<SystemPathBuf>,
+}
+
+impl WorkspaceWatcher {
+    /// Create a new workspace watcher.
+    pub fn new(watcher: Watcher, db: &RootDatabase) -> Self {
+        let mut watcher = Self {
+            watcher,
+            watched_paths: FxHashSet::default(),
+            errored_paths: Vec::new(),
+        };
+
+        watcher.update(db);
+
+        watcher
+    }
+
+    pub fn update(&mut self, db: &RootDatabase) {
+        let new_watch_paths = db.workspace().watch_paths(db);
+
+        let mut added_folders = new_watch_paths.difference(&self.watched_paths).peekable();
+        let mut removed_folders = self.watched_paths.difference(&new_watch_paths).peekable();
+
+        if added_folders.peek().is_none() && removed_folders.peek().is_none() {
+            return;
+        }
+
+        for added_folder in added_folders {
+            // Log a warning. It's not worth aborting if registering a single folder fails because
+            // Ruff otherwise stills works as expected.
+            if let Err(error) = self.watcher.watch(added_folder) {
+                // TODO: Log a user-facing warning.
+                tracing::warn!("Failed to setup watcher for path '{added_folder}': {error}. You have to restart Ruff after making changes to files under this path or you might see stale results.");
+                self.errored_paths.push(added_folder.clone());
+            }
+        }
+
+        for removed_path in removed_folders {
+            if let Some(index) = self
+                .errored_paths
+                .iter()
+                .position(|path| path == removed_path)
+            {
+                self.errored_paths.swap_remove(index);
+                continue;
+            }
+
+            if let Err(error) = self.watcher.unwatch(removed_path) {
+                info!("Failed to remove the file watcher for the path '{removed_path}: {error}.");
+            }
+        }
+
+        info!(
+            "Set up file watchers for {}",
+            DisplayWatchedPaths {
+                paths: &new_watch_paths
+            }
+        );
+
+        self.watched_paths = new_watch_paths;
+    }
+
+    pub fn flush(&self) {
+        self.watcher.flush();
+    }
+
+    pub fn stop(self) {
+        self.watcher.stop();
+    }
+}
+
+struct DisplayWatchedPaths<'a> {
+    paths: &'a FxHashSet<SystemPathBuf>,
+}
+
+impl std::fmt::Display for DisplayWatchedPaths<'_> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        f.write_char('[')?;
+
+        let mut iter = self.paths.iter();
+        if let Some(first) = iter.next() {
+            write!(f, "\"{first}\"")?;
+
+            for path in iter {
+                write!(f, ", \"{path}\"")?;
+            }
+        }
+
+        f.write_char(']')
+    }
+}

crates/red_knot/src/workspace.rs

@@ -6,6 +6,7 @@ use std::{collections::BTreeMap, sync::Arc};
 use rustc_hash::{FxBuildHasher, FxHashSet};
 
 pub use metadata::{PackageMetadata, WorkspaceMetadata};
+use red_knot_module_resolver::system_module_search_paths;
 use ruff_db::{
     files::{system_path_to_file, File},
     system::{walk_directory::WalkState, SystemPath, SystemPathBuf},
@@ -240,6 +241,17 @@ impl Workspace {
             FxHashSet::default()
         }
     }
+
+    /// Returns the paths that should be watched.
+    ///
+    /// The paths that require watching might change with every revision.
+    pub fn watch_paths(self, db: &dyn Db) -> FxHashSet<SystemPathBuf> {
+        ruff_db::system::deduplicate_nested_paths(
+            std::iter::once(self.root(db)).chain(system_module_search_paths(db.upcast())),
+        )
+        .map(SystemPath::to_path_buf)
+        .collect()
+    }
 }
 
 #[salsa::tracked]