Repository: vercel/next.js
Pull request: Turbopack: Rework watcher data structures, store watched list and watcher behind the same RwLock

Reworks the data structures used for the watcher to achieve a few goals:

- I want to avoid multiple locks, as I don't think we actually need two different locks, and multiple locks leaves more room for bugs (e.g. lock inversion).
- The information about if we're in non-recursive mode or recursive mode needs to be available outside of the lock for performance reasons.
- The `watching`/`watched` set needs to be ordered. We don't have a convenient concurrent BTreeSet or Trie/Radix Tree pulled in, and I'd like to avoid adding more dependencies. We end up locking `notify_watcher` anywhere we'd need a write lock anyways, so it might block a few reads, but otherwise it's not that bad.

Fixes in this PR:

- Fixes an issue where multiple concurrent calls to `ensure_watching` for the same directory could cause one of those calls to exit early. This may cause a file in that directory to be read before the filesystem watcher is set up, potentially missing events.
- Improves the time complexity of `restore_if_watching` using an ordered set to avoid iterating over all of `watched` to find child directories. See `OrderedPathSetExt`. This time complexity problem is introduced in the previous PR, #82130.

Other changes:

- It turns out that we don't need to explicitly set up any watchers during `start_watching` except for the `root_dir`. If paths have already been read, they'll set up watchers before re-reading the file when we invalidate it! This simplifies things, and means there's less stuff we need to track if you never start the watcher (i.e. `next build`).
- Added a `--start-watching-late` option to the fuzzer to test calling `start_watching` after all the initial reads happen, to help validate the safety of the claim in the previous bullet point.

Minor (aesthetic) changes:

- Use past tense: `watching` is now `watched` and `ensure_watching` is now `ensure_watched`, as the `watched` list includes both currently active paths *and* previously watched but deleted paths.
- Shorten the names of non-public types, because we have a lot more of them, and the names would get long otherwise

## Benchmarking

This replaces a `DashSet<_>` with (essentially) a `RwLock<BTreeSet<_>>`. There's some risk of increased contention, particularly in the read-only fast-path for `ensure_watched`/`ensure_watching`.

Under the assumption that the most heavily lock-contended use-case for this code is after restoring from persistent cache on a large repository:

1. Check out `vercel-site` in `front`.
2. Hack the new tracing span onto the canary version.
3. Build a release version with `pnpm pack-next --project ~/front/apps/vercel-site --no-js-build -- --release`
4. `rm -rf .next`
5. `NEXT_TURBOPACK_TRACING='turbo_tasks_fs' TURBO_ENGINE_IGNORE_DIRTY=1 pnpm next dev --turbopack`
6. Load `http://localhost:3000/vercel`. Wait for it to fully finish loading. Stop the server with `ctrl+c`.
7. Restart the server, wait for it to finish initializing, and then exit.
8. Start the trace viewer: `pnpm next internal trace .next/trace-turbopack`, and look at the aggregated times of `ensure_watched`/`ensure_watching`.

Results:

before (canary): 313ms ensure_watching (2235x spans)
after (this pr): 279ms ensure_watched (2235x spans)

So no significant difference.

Closes PACK-5229

Changed files:
filename: turbopack/crates/turbo-tasks-fs/src/lib.rs
patch: @@ -261,11 +261,7 @@ impl DiskFileSystemInner {
         let invalidator = turbo_tasks::get_invalidator();
         self.invalidator_map
             .insert(path.to_owned(), invalidator, None);
-        if let Some(non_recursive) = &self.watcher.non_recursive_state
-            && let Some(dir) = path.parent()
-        {
-            non_recursive.ensure_watching(&self.watcher, dir, self.root_path())?;
-        }
+        self.watcher.ensure_watched_file(path, self.root_path())?;
         Ok(())
     }
 
@@ -291,11 +287,7 @@ impl DiskFileSystemInner {
             .collect::<Vec<_>>();
         invalidators.insert(invalidator, Some(write_content));
         drop(invalidator_map);
-        if let Some(non_recursive) = &self.watcher.non_recursive_state
-            && let Some(dir) = path.parent()
-        {
-            non_recursive.ensure_watching(&self.watcher, dir, self.root_path())?;
-        }
+        self.watcher.ensure_watched_file(path, self.root_path())?;
         Ok(old_invalidators)
     }
 
@@ -305,9 +297,7 @@ impl DiskFileSystemInner {
         let invalidator = turbo_tasks::get_invalidator();
         self.dir_invalidator_map
             .insert(path.to_owned(), invalidator, None);
-        if let Some(non_recursive) = &self.watcher.non_recursive_state {
-            non_recursive.ensure_watching(&self.watcher, path, self.root_path())?;
-        }
+        self.watcher.ensure_watched_dir(path, self.root_path())?;
         Ok(())
     }
 

filename: turbopack/crates/turbo-tasks-fs/src/path_map.rs
patch: @@ -1,5 +1,5 @@
 use std::{
-    collections::{BTreeMap, btree_map::CursorMut},
+    collections::{BTreeMap, BTreeSet, btree_map, btree_set},
     ops::Bound,
     path::{Path, PathBuf},
 };
@@ -9,24 +9,31 @@ use std::{
 ///
 /// In the future, this may use a more efficient representation, like a radix tree or trie.
 pub trait OrderedPathMapExt<V> {
-    fn extract_path_with_children<'a>(&'a mut self, path: &'a Path) -> ExtractWithChildren<'a, V>;
+    fn extract_path_with_children<'a>(
+        &'a mut self,
+        path: &'a Path,
+    ) -> PathMapExtractPathWithChildren<'a, V>;
 }
 
 impl<V> OrderedPathMapExt<V> for BTreeMap<PathBuf, V> {
-    fn extract_path_with_children<'a>(&'a mut self, path: &'a Path) -> ExtractWithChildren<'a, V> {
-        ExtractWithChildren {
+    /// Iterates over and removes `path` and all of its children.
+    fn extract_path_with_children<'a>(
+        &'a mut self,
+        path: &'a Path,
+    ) -> PathMapExtractPathWithChildren<'a, V> {
+        PathMapExtractPathWithChildren {
             cursor: self.lower_bound_mut(Bound::Included(path)),
             parent_path: path,
         }
     }
 }
 
-pub struct ExtractWithChildren<'a, V> {
-    cursor: CursorMut<'a, PathBuf, V>,
+pub struct PathMapExtractPathWithChildren<'a, V> {
+    cursor: btree_map::CursorMut<'a, PathBuf, V>,
     parent_path: &'a Path,
 }
 
-impl<V> Iterator for ExtractWithChildren<'_, V> {
+impl<V> Iterator for PathMapExtractPathWithChildren<'_, V> {
     type Item = (PathBuf, V);
 
     fn next(&mut self) -> Option<Self::Item> {
@@ -44,20 +51,61 @@ impl<V> Iterator for ExtractWithChildren<'_, V> {
     }
 }
 
+/// A thin wrapper around [`BTreeSet<PathBuf>`] that provides efficient iteration of child paths.
+///
+/// In the future, this may use a more efficient representation, like a radix tree or trie.
+pub trait OrderedPathSetExt {
+    /// Iterates over the children of `path`, excluding `path` itself.
+    fn iter_path_children<'a>(&'a mut self, path: &'a Path) -> PathSetIterPathChildren<'a>;
+}
+
+impl OrderedPathSetExt for BTreeSet<PathBuf> {
+    fn iter_path_children<'a>(&'a mut self, path: &'a Path) -> PathSetIterPathChildren<'a> {
+        PathSetIterPathChildren {
+            // this is range written weirdly due to type inference limitations:
+            // https://stackoverflow.com/a/66130898
+            range: self.range::<Path, _>((Bound::Excluded(path), Bound::Unbounded)),
+            parent_path: path,
+        }
+    }
+}
+
+pub struct PathSetIterPathChildren<'a> {
+    // we don't need the nightly cursors API for this, the `Range` type is sufficient.
+    range: btree_set::Range<'a, PathBuf>,
+    parent_path: &'a Path,
+}
+
+impl<'a> Iterator for PathSetIterPathChildren<'a> {
+    type Item = &'a Path;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        // this simple implementation works because `Path` implements `Ord` (and `starts_with`)
+        // using path component comparision, rather than raw byte comparisions. The parent path is
+        // always guaranteed to be placed immediately before its children (pre-order traversal).
+        let current_path = self.range.next()?;
+        if !current_path.starts_with(self.parent_path) {
+            return None;
+        }
+        Some(&**current_path)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
 
     #[test]
-    fn test_extract_with_children() {
+    fn test_map_extract_path_with_children() {
         let mut map = BTreeMap::default();
         map.insert(PathBuf::from("a"), 1);
         map.insert(PathBuf::from("a/b"), 2);
         map.insert(PathBuf::from("a/b/c"), 3);
         map.insert(PathBuf::from("a/b/d"), 4);
-        map.insert(PathBuf::from("a/c"), 5);
-        map.insert(PathBuf::from("x/y/z"), 6);
-        map.insert(PathBuf::from("z/a/b"), 7);
+        map.insert(PathBuf::from("a/b/d/e"), 5);
+        map.insert(PathBuf::from("a/c"), 6);
+        map.insert(PathBuf::from("x/y/z"), 7);
+        map.insert(PathBuf::from("z/a/b"), 8);
 
         let parent_path = PathBuf::from("a/b");
         let extracted: Vec<_> = map.extract_path_with_children(&parent_path).collect();
@@ -66,15 +114,39 @@ mod tests {
             (PathBuf::from("a/b"), 2),
             (PathBuf::from("a/b/c"), 3),
             (PathBuf::from("a/b/d"), 4),
+            (PathBuf::from("a/b/d/e"), 5),
         ];
         assert_eq!(extracted, expected_extracted);
 
         let mut expected_remaining = BTreeMap::new();
         expected_remaining.insert(PathBuf::from("a"), 1);
-        expected_remaining.insert(PathBuf::from("a/c"), 5);
-        expected_remaining.insert(PathBuf::from("x/y/z"), 6);
-        expected_remaining.insert(PathBuf::from("z/a/b"), 7);
+        expected_remaining.insert(PathBuf::from("a/c"), 6);
+        expected_remaining.insert(PathBuf::from("x/y/z"), 7);
+        expected_remaining.insert(PathBuf::from("z/a/b"), 8);
 
         assert_eq!(map, expected_remaining);
     }
+
+    #[test]
+    fn test_set_iter_path_children() {
+        let mut set = BTreeSet::default();
+        set.insert(PathBuf::from("a"));
+        set.insert(PathBuf::from("a/b"));
+        set.insert(PathBuf::from("a/b/c"));
+        set.insert(PathBuf::from("a/b/d"));
+        set.insert(PathBuf::from("a/b/d/e"));
+        set.insert(PathBuf::from("a/c"));
+        set.insert(PathBuf::from("x/y/z"));
+        set.insert(PathBuf::from("z/a/b"));
+
+        let parent_path = PathBuf::from("a/b");
+        let iterated: Vec<_> = set.iter_path_children(&parent_path).collect();
+
+        let expected_iterated = vec![
+            PathBuf::from("a/b/c"),
+            PathBuf::from("a/b/d"),
+            PathBuf::from("a/b/d/e"),
+        ];
+        assert_eq!(iterated, expected_iterated);
+    }
 }

filename: turbopack/crates/turbo-tasks-fs/src/watcher.rs
patch: @@ -1,17 +1,17 @@
 use std::{
     any::Any,
+    collections::BTreeSet,
     env, fmt,
     mem::take,
     path::{Path, PathBuf},
     sync::{
-        Arc, LazyLock, Mutex,
+        Arc, LazyLock, RwLock, RwLockWriteGuard,
         mpsc::{Receiver, TryRecvError, channel},
     },
     time::Duration,
 };
 
 use anyhow::{Context, Result};
-use dashmap::DashSet;
 use notify::{
     Config, EventKind, PollWatcher, RecommendedWatcher, RecursiveMode, Watcher,
     event::{MetadataKind, ModifyKind, RenameMode},
@@ -61,145 +61,207 @@ static WATCH_RECURSIVE_MODE: LazyLock<RecursiveMode> = LazyLock::new(|| {
     }
 });
 
-/// A thin wrapper around [`RecommendedWatcher`] and [`PollWatcher`].
-enum DiskWatcherInternal {
-    Recommended(RecommendedWatcher),
-    Polling(PollWatcher),
+#[derive(Serialize, Deserialize)]
+pub(crate) struct DiskWatcher {
+    #[serde(skip, default = "State::new_stopped")]
+    state: State,
 }
 
-impl DiskWatcherInternal {
-    fn watch(&mut self, path: &Path, recursive_mode: RecursiveMode) -> notify::Result<()> {
+enum State {
+    // Note: Information about if we're a recursive or non-recursive watcher must live outside the
+    // `RwLock` to allow us to quickly bail out on calls to `ensure_watched`.
+    Recursive(RwLock<RecursiveState>),
+    NonRecursive(RwLock<NonRecursiveState>),
+}
+
+enum StateWriteGuard<'a> {
+    Recursive(RwLockWriteGuard<'a, RecursiveState>),
+    NonRecursive(RwLockWriteGuard<'a, NonRecursiveState>),
+}
+
+impl State {
+    fn new_stopped() -> Self {
+        match *WATCH_RECURSIVE_MODE {
+            RecursiveMode::Recursive => Self::Recursive(RwLock::new(RecursiveState::Stopped)),
+            RecursiveMode::NonRecursive => {
+                Self::NonRecursive(RwLock::new(NonRecursiveState::Stopped))
+            }
+        }
+    }
+
+    fn write(&self) -> StateWriteGuard<'_> {
         match self {
-            DiskWatcherInternal::Recommended(watcher) => watcher.watch(path, recursive_mode),
-            DiskWatcherInternal::Polling(watcher) => watcher.watch(path, recursive_mode),
+            Self::Recursive(state) => StateWriteGuard::Recursive(state.write().unwrap()),
+            Self::NonRecursive(state) => StateWriteGuard::NonRecursive(state.write().unwrap()),
         }
     }
 }
 
-#[derive(Serialize, Deserialize)]
-pub(crate) struct DiskWatcher {
-    /// This value is [`None`] when the watcher has been stopped (see
-    /// [`DiskWatcher::stop_watching`]).
-    #[serde(skip)]
-    internal: Mutex<Option<DiskWatcherInternal>>,
-
-    #[serde(skip, default = "NonRecursiveDiskWatcherState::try_new")]
-    pub(crate) non_recursive_state: Option<NonRecursiveDiskWatcherState>,
+/// Used by when [`WATCH_RECURSIVE_MODE`] is [`RecursiveMode::Recursive`] (default on macOS and
+/// Windows).
+enum RecursiveState {
+    /// Used when [`DiskWatcher::start_watching`] hasn't been called yet or after
+    /// [`DiskWatcher::stop_watching`] is called.
+    Stopped,
+    Watching {
+        /// Hold onto the watcher: When this is dropped, it will cause the channel to disconnect
+        _notify_watcher: NotifyWatcher,
+    },
 }
 
-impl Default for DiskWatcher {
-    fn default() -> Self {
-        Self {
-            internal: Mutex::new(None),
-            non_recursive_state: NonRecursiveDiskWatcherState::try_new(),
-        }
-    }
+/// Used by when [`WATCH_RECURSIVE_MODE`] is [`RecursiveMode::NonRecursive`] (default on Linux).
+enum NonRecursiveState {
+    /// Used when [`DiskWatcher::start_watching`] hasn't been called yet or after
+    /// [`DiskWatcher::stop_watching`] is called.
+    Stopped,
+    Watching(NonRecursiveWatchingState),
 }
 
-/// Extra state used by [`DiskWatcher`] when [`WATCH_RECURSIVE_MODE`] is
-/// [`RecursiveMode::NonRecursive`] (default on Linux).
-pub(crate) struct NonRecursiveDiskWatcherState {
-    /// Keeps track of which directories are currently (or were previously) watched.
+// split out from the `NonRecursiveState` enum because we want to pass this value around
+struct NonRecursiveWatchingState {
+    notify_watcher: NotifyWatcher,
+    /// Keeps track of which directories are currently or were previously watched by
+    /// [`Self::notify_watcher`].
     ///
     /// Invariants:
     /// - Never contains `root_path`. A watcher for `root_path` is implicitly set up during
     ///   [`DiskWatcher::start_watching`].
     /// - Contains all parent directories up to `root_path` for every entry.
-    watching: DashSet<PathBuf>,
+    watched: BTreeSet<PathBuf>,
 }
 
-impl NonRecursiveDiskWatcherState {
-    fn try_new() -> Option<NonRecursiveDiskWatcherState> {
-        match *WATCH_RECURSIVE_MODE {
-            RecursiveMode::Recursive => None,
-            RecursiveMode::NonRecursive => Some(NonRecursiveDiskWatcherState {
-                watching: DashSet::new(),
-            }),
+/// A thin wrapper around [`RecommendedWatcher`] and [`PollWatcher`].
+enum NotifyWatcher {
+    Recommended(RecommendedWatcher),
+    Polling(PollWatcher),
+}
+
+impl NotifyWatcher {
+    fn watch(&mut self, path: &Path, recursive_mode: RecursiveMode) -> notify::Result<()> {
+        match self {
+            Self::Recommended(watcher) => watcher.watch(path, recursive_mode),
+            Self::Polling(watcher) => watcher.watch(path, recursive_mode),
         }
     }
+}
+
+mod non_recursive_helpers {
+    use super::*;
+    use crate::path_map::OrderedPathSetExt;
 
     /// Called after a rescan in case a previously watched-but-deleted directory was recreated.
-    pub(crate) fn restore_all_watching(&self, watcher: &DiskWatcher, root_path: &Path) {
-        let mut internal_guard = watcher.internal.lock().unwrap();
-        let Some(internal) = &mut *internal_guard else {
+    #[instrument(skip_all, level = "trace")]
+    pub fn restore_all_watched_ignore_errors(state: &RwLock<NonRecursiveState>, root_path: &Path) {
+        let mut guard = state.write().unwrap();
+        let NonRecursiveState::Watching(watching_state) = &mut *guard else {
             return;
         };
-        for dir_path in self.watching.iter() {
+        for dir_path in watching_state.watched.iter() {
             // TODO: Report diagnostics if this error happens
             //
-            // Don't watch the parents, because those are already included in `self.watching` (so
+            // Don't watch the parents, because those are already included in `self.watched` (so
             // it'd be redundant), but also because this could deadlock, since we'd try to modify
-            // `self.watching` while iterating over it (write lock overlapping with a read lock).
-            let _ = self.start_watching_dir(internal, &dir_path, root_path);
+            // `self.watched` while iterating over it (write lock overlapping with a read lock).
+            let _ = start_watching_dir(&mut watching_state.notify_watcher, dir_path, root_path);
         }
     }
 
     /// Called when a new directory is found in a parent directory we're watching. Restores the
     /// watcher if we were previously watching it.
-    pub(crate) fn restore_if_watching(
-        &self,
-        watcher: &DiskWatcher,
+    #[instrument(skip_all, level = "trace")]
+    pub fn restore_if_watched(
+        state: &RwLock<NonRecursiveState>,
         dir_path: &Path,
         root_path: &Path,
     ) -> Result<()> {
-        if dir_path == root_path || !self.watching.contains(dir_path) {
+        // fast path: The root directory is always implicitly watched during
+        // `DiskWatcher::start_watching`, we assume it is never deleted and never needs to be
+        // restored.
+        if dir_path == root_path {
             return Ok(());
         }
-        let mut internal_guard = watcher.internal.lock().unwrap();
-        let Some(internal) = &mut *internal_guard else {
+
+        // fast path: the directory isn't in `watched`, only take a read lock and bail out early
+        {
+            let guard = state.read().unwrap();
+            let NonRecursiveState::Watching(watching_state) = &*guard else {
+                return Ok(());
+            };
+            if !watching_state.watched.contains(dir_path) {
+                return Ok(());
+            }
+        }
+
+        // slow path: re-watch the path
+        let mut guard = state.write().unwrap();
+        let NonRecursiveState::Watching(watching_state) = &mut *guard else {
             return Ok(());
         };
 
         // watch the new directory
-        self.start_watching_dir(internal, dir_path, root_path)?;
+        start_watching_dir(&mut watching_state.notify_watcher, dir_path, root_path)?;
 
         // Also try to restore any watchers for children of this directory
-        for child_path in self
-            .watching
-            .iter()
-            .filter(|p| p.key().starts_with(dir_path) && **p != dir_path)
-        {
-            // Don't watch the parents -- see the comment on `restore_all_watching`
-            self.start_watching_dir(internal, child_path.key(), root_path)?;
+        for child_path in watching_state.watched.iter_path_children(dir_path) {
+            // Don't watch the parents -- see the comment on `restore_all_watched`
+            start_watching_dir(&mut watching_state.notify_watcher, child_path, root_path)?;
         }
         Ok(())
     }
 
     /// Called when a file in `dir_path` or `dir_path` itself is read or written. Adds a new watcher
     /// if we're not already watching the directory.
-    pub(crate) fn ensure_watching(
-        &self,
-        watcher: &DiskWatcher,
+    ///
+    /// This should be called *before* reading a file to avoid a race condition.
+    #[instrument(skip_all, level = "trace")]
+    pub fn ensure_watched(
+        state: &RwLock<NonRecursiveState>,
         dir_path: &Path,
         root_path: &Path,
     ) -> Result<()> {
-        if dir_path == root_path || self.watching.contains(dir_path) {
+        // fast path: The root directory is always implicitly watched during
+        // `DiskWatcher::start_watching`.
+        if dir_path == root_path {
             return Ok(());
         }
-        let mut internal_guard = watcher.internal.lock().unwrap();
-        let Some(internal) = &mut *internal_guard else {
+
+        // fast path: the directory is already in `watched`, only take a read lock and bail out
+        // early
+        {
+            let guard = state.read().unwrap();
+            let NonRecursiveState::Watching(watching_state) = &*guard else {
+                return Ok(());
+            };
+            if watching_state.watched.contains(dir_path) {
+                return Ok(());
+            }
+        }
+
+        // slow path: watch the path
+        let mut guard = state.write().unwrap();
+        let NonRecursiveState::Watching(watching_state) = &mut *guard else {
             return Ok(());
         };
-        if self.watching.insert(dir_path.to_path_buf()) {
-            self.start_watching_dir_and_parents(internal, dir_path, root_path)?;
+        if watching_state.watched.insert(dir_path.to_path_buf()) {
+            start_watching_dir_and_parents(watching_state, dir_path, root_path)?;
         }
         Ok(())
     }
 
-    /// Private helper, assumes that `dir_path` has already been added to `self.watching`.
+    /// Private helper, assumes that `dir_path` has already been added to
+    /// [`NonRecursiveWatchingState::watched`].
     ///
     /// This does not watch any of the parent directories. For that, use
     /// [`start_watching_dir_and_parents`]. Use this method when iterating over previously-watched
     /// values in `self.watching`.
     fn start_watching_dir(
-        &self,
-        watcher_internal: &mut DiskWatcherInternal,
+        notify_watcher: &mut NotifyWatcher,
         dir_path: &Path,
         root_path: &Path,
     ) -> Result<()> {
         debug_assert_ne!(dir_path, root_path);
 
-        match watcher_internal.watch(dir_path, RecursiveMode::NonRecursive) {
+        match notify_watcher.watch(dir_path, RecursiveMode::NonRecursive) {
             Ok(())
             | Err(notify::Error {
                 // The path was probably deleted before we could process the event, but the parent
@@ -214,20 +276,20 @@ impl NonRecursiveDiskWatcherState {
         }
     }
 
-    /// Private helper, assumes that `dir_path` has already been added to `self.watching`.
+    /// Private helper, assumes that `dir_path` has already been added to
+    /// [`NonRecursiveWatchingState::watched`].
     ///
     /// Watches the given `dir_path` and every parent up to `root_path`. Parents must be recursively
     /// watched in case any of them change:
     /// https://docs.rs/notify/latest/notify/#parent-folder-deletion
     fn start_watching_dir_and_parents(
-        &self,
-        watcher_internal: &mut DiskWatcherInternal,
+        state: &mut NonRecursiveWatchingState,
         dir_path: &Path,
         root_path: &Path,
     ) -> Result<()> {
         let mut cur_path = dir_path;
         loop {
-            self.start_watching_dir(watcher_internal, cur_path, root_path)?;
+            start_watching_dir(&mut state.notify_watcher, cur_path, root_path)?;
 
             let Some(parent_path) = cur_path.parent() else {
                 // this should never happen as we break before we reach the root path
@@ -237,7 +299,7 @@ impl NonRecursiveDiskWatcherState {
                 );
             };
 
-            if parent_path == root_path || !self.watching.insert(parent_path.to_path_buf()) {
+            if parent_path == root_path || !state.watched.insert(parent_path.to_path_buf()) {
                 break;
             }
 
@@ -249,8 +311,10 @@ impl NonRecursiveDiskWatcherState {
 }
 
 impl DiskWatcher {
-    pub(crate) fn new() -> Self {
-        Default::default()
+    pub fn new() -> Self {
+        Self {
+            state: State::new_stopped(),
+        }
     }
 
     /// Create a watcher and start watching by creating `debounced` watcher
@@ -268,14 +332,22 @@ impl DiskWatcher {
     /// - Emits only one Remove event when deleting a directory (inotify)
     /// - Doesn't emit duplicate create events
     /// - Doesn't emit Modify events after a Create event
-    pub(crate) fn start_watching(
+    pub fn start_watching(
         &self,
         fs_inner: Arc<DiskFileSystemInner>,
         report_invalidation_reason: bool,
         poll_interval: Option<Duration>,
     ) -> Result<()> {
-        let mut internal_guard = self.internal.lock().unwrap();
-        if internal_guard.is_some() {
+        let state_guard = self.state.write();
+
+        // bail out if we're already watching
+        if let StateWriteGuard::Recursive(guard) = &state_guard
+            && matches!(**guard, RecursiveState::Watching { .. })
+        {
+            return Ok(());
+        } else if let StateWriteGuard::NonRecursive(guard) = &state_guard
+            && matches!(**guard, NonRecursiveState::Watching(..))
+        {
             return Ok(());
         }
 
@@ -284,27 +356,29 @@ impl DiskWatcher {
         // Create a watcher object, delivering debounced events.
         // The notification back-end is selected based on the platform.
         let config = Config::default();
-        // we should track and invalidate each part of a symlink chain ourselves in turbo-tasks-fs
+        // we should track and invalidate each part of a symlink chain ourselves in
+        // turbo-tasks-fs
         config.with_follow_symlinks(false);
 
-        let mut internal = if let Some(poll_interval) = poll_interval {
+        let mut notify_watcher = if let Some(poll_interval) = poll_interval {
             let config = config.with_poll_interval(poll_interval);
-
-            DiskWatcherInternal::Polling(PollWatcher::new(tx, config)?)
+            NotifyWatcher::Polling(PollWatcher::new(tx, config)?)
         } else {
-            DiskWatcherInternal::Recommended(RecommendedWatcher::new(tx, Config::default())?)
+            NotifyWatcher::Recommended(RecommendedWatcher::new(tx, Config::default())?)
         };
 
-        if let Some(non_recursive) = &self.non_recursive_state {
-            internal.watch(fs_inner.root_path(), RecursiveMode::NonRecursive)?;
-            for dir_path in non_recursive.watching.iter() {
-                internal.watch(&dir_path, RecursiveMode::NonRecursive)?;
-            }
-        } else {
-            internal.watch(fs_inner.root_path(), RecursiveMode::Recursive)?;
-        }
+        // TOCTOU: we must watch `root_path` before calling any invalidators and setting up the
+        // watchers in their associated functions
+        let root_path = fs_inner.root_path();
+        let recursive_mode = match state_guard {
+            StateWriteGuard::Recursive(_) => RecursiveMode::Recursive,
+            StateWriteGuard::NonRecursive(_) => RecursiveMode::NonRecursive,
+        };
+        notify_watcher.watch(root_path, recursive_mode)?;
 
-        // We need to invalidate all reads that happened before watching
+        // We need to invalidate all reads or writes that happened before watching. As a
+        // side-effect, this will call `ensure_watched` again, setting up any watchers needed.
+        //
         // Best is to start_watching before starting to read
         {
             let span = tracing::info_span!("invalidate filesystem");
@@ -345,24 +419,39 @@ impl DiskWatcher {
             }
         }
 
-        internal_guard.replace(internal);
-        drop(internal_guard);
-
         spawn_thread(move || {
             fs_inner
                 .clone()
                 .watcher
                 .watch_thread(rx, fs_inner, report_invalidation_reason)
         });
 
+        // Updating `self.state` is done last. If we panic while setting up the watcher, it'll
+        // stay in the `Stopped` state.
+        match state_guard {
+            StateWriteGuard::Recursive(mut recursive) => {
+                *recursive = RecursiveState::Watching {
+                    _notify_watcher: notify_watcher,
+                }
+            }
+            StateWriteGuard::NonRecursive(mut non_recursive) => {
+                *non_recursive = NonRecursiveState::Watching(NonRecursiveWatchingState {
+                    notify_watcher,
+                    watched: BTreeSet::new(),
+                })
+            }
+        };
+
         Ok(())
     }
 
-    pub(crate) fn stop_watching(&self) {
-        if let Some(watcher) = self.internal.lock().unwrap().take() {
-            drop(watcher);
-            // thread will detect the stop because the channel is disconnected
+    pub fn stop_watching(&self) {
+        match &self.state {
+            State::Recursive(state) => *state.write().unwrap() = RecursiveState::Stopped,
+            State::NonRecursive(state) => *state.write().unwrap() = NonRecursiveState::Stopped,
         }
+        // thread will detect the stop because the channel is disconnected when `NotifyWatcher` is
+        // dropped
     }
 
     /// Internal thread that processes the events from the watcher
@@ -372,15 +461,15 @@ impl DiskWatcher {
     fn watch_thread(
         &self,
         rx: Receiver<notify::Result<notify::Event>>,
-        inner: Arc<DiskFileSystemInner>,
+        fs_inner: Arc<DiskFileSystemInner>,
         report_invalidation_reason: bool,
     ) {
         let mut batched_invalidate_path = FxHashSet::default();
         let mut batched_invalidate_path_dir = FxHashSet::default();
         let mut batched_invalidate_path_and_children = FxHashSet::default();
         let mut batched_invalidate_path_and_children_dir = FxHashSet::default();
 
-        let mut batched_new_paths = if self.non_recursive_state.is_some() {
+        let mut batched_new_paths = if let State::NonRecursive(_) = self.state {
             Some(FxHashSet::default())
         } else {
             None
@@ -402,25 +491,30 @@ impl DiskWatcher {
                         // echo 3 | sudo tee /proc/sys/fs/inotify/max_queued_events
                         // ```
                         if event.need_rescan() {
-                            let _lock = inner.invalidation_lock.blocking_write();
+                            let _lock = fs_inner.invalidation_lock.blocking_write();
 
-                            if let Some(non_recursive) = &self.non_recursive_state {
+                            if let State::NonRecursive(non_recursive) = &self.state {
                                 // we can't narrow this down to a smaller set of paths: Rescan
                                 // events (at least when tested on Linux) come with no `paths`, and
                                 // we use only one global `notify::Watcher` instance.
-                                non_recursive.restore_all_watching(self, inner.root_path());
+                                //
+                                // TODO: Report diagnostics if an error happens
+                                non_recursive_helpers::restore_all_watched_ignore_errors(
+                                    non_recursive,
+                                    fs_inner.root_path(),
+                                );
                                 if let Some(batched_new_paths) = &mut batched_new_paths {
                                     batched_new_paths.clear();
                                 }
                             }
 
                             if report_invalidation_reason {
-                                inner.invalidate_with_reason(|path| InvalidateRescan {
+                                fs_inner.invalidate_with_reason(|path| InvalidateRescan {
                                     // this path is just used for display purposes
                                     path: RcStr::from(path.to_string_lossy()),
                                 });
                             } else {
-                                inner.invalidate();
+                                fs_inner.invalidate();
                             }
 
                             // no need to process the rest of the batch as we just
@@ -537,9 +631,9 @@ impl DiskWatcher {
 
                         if paths.is_empty() {
                             batched_invalidate_path_and_children
-                                .insert(inner.root_path().to_path_buf());
+                                .insert(fs_inner.root_path().to_path_buf());
                             batched_invalidate_path_and_children_dir
-                                .insert(inner.root_path().to_path_buf());
+                                .insert(fs_inner.root_path().to_path_buf());
                         } else {
                             batched_invalidate_path_and_children.extend(paths.clone());
                             batched_invalidate_path_and_children_dir.extend(paths.clone());
@@ -572,49 +666,78 @@ impl DiskWatcher {
 
             // We need to start watching first before invalidating the changed paths...
             // This is only needed on platforms we don't do recursive watching on.
-            if let Some(non_recursive) = &self.non_recursive_state {
+            if let State::NonRecursive(non_recursive) = &self.state {
                 for path in batched_new_paths.as_mut().unwrap().drain() {
                     // TODO: Report diagnostics if this error happens
-                    let _ = non_recursive.restore_if_watching(self, &path, inner.root_path());
+                    let _ = non_recursive_helpers::restore_if_watched(
+                        non_recursive,
+                        &path,
+                        fs_inner.root_path(),
+                    );
                 }
             }
 
-            let _lock = inner.invalidation_lock.blocking_write();
+            let _lock = fs_inner.invalidation_lock.blocking_write();
             {
-                let mut invalidator_map = inner.invalidator_map.lock().unwrap();
+                let mut invalidator_map = fs_inner.invalidator_map.lock().unwrap();
                 invalidate_path(
-                    &inner,
+                    &fs_inner,
                     report_invalidation_reason,
                     &mut invalidator_map,
                     batched_invalidate_path.drain(),
                 );
                 invalidate_path_and_children_execute(
-                    &inner,
+                    &fs_inner,
                     report_invalidation_reason,
                     &mut invalidator_map,
                     batched_invalidate_path_and_children.drain(),
                 );
             }
             {
-                let mut dir_invalidator_map = inner.dir_invalidator_map.lock().unwrap();
+                let mut dir_invalidator_map = fs_inner.dir_invalidator_map.lock().unwrap();
                 invalidate_path(
-                    &inner,
+                    &fs_inner,
                     report_invalidation_reason,
                     &mut dir_invalidator_map,
                     batched_invalidate_path_dir.drain(),
                 );
                 invalidate_path_and_children_execute(
-                    &inner,
+                    &fs_inner,
                     report_invalidation_reason,
                     &mut dir_invalidator_map,
                     batched_invalidate_path_and_children_dir.drain(),
                 );
             }
         }
     }
+
+    pub fn ensure_watched_file(&self, path: &Path, root_path: &Path) -> Result<()> {
+        // Watch the parent directory instead of the specified file, since directories also track
+        // their immediate children (even in non-recursive mode), and we need to watch all the
+        // parents anyways.
+        if let State::NonRecursive(non_recursive) = &self.state
+            && let Some(dir_path) = path.parent()
+        {
+            non_recursive_helpers::ensure_watched(non_recursive, dir_path, root_path)?;
+        }
+        Ok(())
+    }
+
+    pub fn ensure_watched_dir(&self, dir_path: &Path, root_path: &Path) -> Result<()> {
+        if let State::NonRecursive(non_recursive) = &self.state {
+            non_recursive_helpers::ensure_watched(non_recursive, dir_path, root_path)?;
+        }
+        Ok(())
+    }
 }
 
-#[instrument(parent = None, level = "info", name = "DiskFileSystem file change", skip_all, fields(name = display(path.display())))]
+#[instrument(
+    parent = None,
+    level = "info",
+    name = "DiskFileSystem file change",
+    skip_all,
+    fields(name = %path.display())
+)]
 fn invalidate(
     inner: &DiskFileSystemInner,
     report_invalidation_reason: bool,

filename: turbopack/crates/turbo-tasks-fuzz/src/main.rs
patch: @@ -52,6 +52,9 @@ struct FsWatcher {
     directory_modifications: u32,
     #[arg(long)]
     print_missing_invalidations: bool,
+    /// Call `start_watching` after the initial read of files instead of before (the default).
+    #[arg(long)]
+    start_watching_late: bool,
 }
 
 #[tokio::main]
@@ -91,7 +94,9 @@ async fn fuzz_fs_watcher(args: FsWatcher) -> anyhow::Result<()> {
             .await?;
         create_directory_tree(&mut FxHashSet::default(), &fs_root, args.depth, args.width)?;
 
-        project_fs.await?.start_watching(None).await?;
+        if !args.start_watching_late {
+            project_fs.await?.start_watching(None).await?;
+        }
 
         let read_all_paths_op =
             read_all_paths_operation(invalidations.clone(), project_root, args.depth, args.width);
@@ -102,6 +107,10 @@ async fn fuzz_fs_watcher(args: FsWatcher) -> anyhow::Result<()> {
             invalidations.clear();
         }
 
+        if args.start_watching_late {
+            project_fs.await?.start_watching(None).await?;
+        }
+
         let mut rand_buf = [0; 16];
         let mut rng = rand::rngs::SmallRng::from_rng(&mut rand::rng());
         loop {