Repository: vercel/next.js
Pull request: Turbopack: Use a BTreeMap in InvalidatorMap to avoid many `O(n)` traversals when finding child paths

This fixes an `O(n)` iteration of the invalidators (which can be quite large) in `invalidate_path_and_children_execute`.

It requires the nightly `btree_cursors` API. It's possible to implement this without that API, but it'd be a lot less efficient, and I'm concerned about how hot this codepath might be.

This codepath is used for both the recursive and non-recursive filesystem watcher, so it impacts all platforms.

Closes PACK-5152

Changed files:
filename: turbopack/crates/turbo-tasks-fs/src/invalidator_map.rs
patch: @@ -1,4 +1,5 @@
 use std::{
+    collections::BTreeMap,
     path::PathBuf,
     sync::{LockResult, Mutex, MutexGuard},
 };
@@ -16,18 +17,18 @@ pub enum WriteContent {
     Link(ReadRef<LinkContent>),
 }
 
-type InnerMap = FxHashMap<PathBuf, FxHashMap<Invalidator, Option<WriteContent>>>;
+pub type LockedInvalidatorMap = BTreeMap<PathBuf, FxHashMap<Invalidator, Option<WriteContent>>>;
 
 pub struct InvalidatorMap {
     queue: ConcurrentQueue<(PathBuf, Invalidator, Option<WriteContent>)>,
-    map: Mutex<InnerMap>,
+    map: Mutex<LockedInvalidatorMap>,
 }
 
 impl Default for InvalidatorMap {
     fn default() -> Self {
         Self {
             queue: ConcurrentQueue::unbounded(),
-            map: Default::default(),
+            map: Mutex::<LockedInvalidatorMap>::default(),
         }
     }
 }
@@ -37,7 +38,7 @@ impl InvalidatorMap {
         Self::default()
     }
 
-    pub fn lock(&self) -> LockResult<MutexGuard<'_, InnerMap>> {
+    pub fn lock(&self) -> LockResult<MutexGuard<'_, LockedInvalidatorMap>> {
         let mut guard = self.map.lock()?;
         while let Ok((key, value, write_content)) = self.queue.pop() {
             guard.entry(key).or_default().insert(value, write_content);
@@ -76,7 +77,7 @@ impl Serialize for InvalidatorMap {
         // persisted cache, but we don't invalidate the fs writes. Those read invalidations trigger
         // re-inserts into the `InvalidatorMap`. If we knew that certain invalidators were only
         // needed for reads, we could potentially avoid serializing those paths entirely.
-        let inner: &InnerMap = &self.lock().unwrap();
+        let inner: &LockedInvalidatorMap = &self.lock().unwrap();
         serializer.serialize_newtype_struct("InvalidatorMap", inner)
     }
 }

filename: turbopack/crates/turbo-tasks-fs/src/lib.rs
patch: @@ -1,4 +1,5 @@
 #![allow(clippy::needless_return)] // tokio macro-generated code doesn't respect this
+#![feature(btree_cursors)] // needed for the `InvalidatorMap` and watcher, reduces time complexity
 #![feature(trivial_bounds)]
 #![feature(min_specialization)]
 #![feature(iter_advance_by)]
@@ -16,13 +17,15 @@ pub mod invalidation;
 mod invalidator_map;
 pub mod json;
 mod mutex_map;
+mod path_map;
 mod read_glob;
 mod retry;
 pub mod rope;
 pub mod source_context;
 pub mod util;
 pub(crate) mod virtual_fs;
 mod watcher;
+
 use std::{
     borrow::Cow,
     cmp::{Ordering, min},
@@ -40,14 +43,10 @@ use anyhow::{Context, Result, anyhow, bail};
 use auto_hash_map::{AutoMap, AutoSet};
 use bitflags::bitflags;
 use dunce::simplified;
-use glob::Glob;
 use indexmap::IndexSet;
-use invalidator_map::InvalidatorMap;
 use jsonc_parser::{ParseOptions, parse_to_serde_value};
 use mime::Mime;
 use rayon::iter::{IntoParallelIterator, ParallelIterator};
-pub use read_glob::ReadGlobResult;
-use read_glob::{read_glob, track_glob};
 use rustc_hash::FxHashSet;
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
@@ -60,17 +59,19 @@ use turbo_tasks::{
     mark_session_dependent, mark_stateful, trace::TraceRawVcs,
 };
 use turbo_tasks_hash::{DeterministicHash, DeterministicHasher, hash_xxh3_hash64};
-use util::{extract_disk_access, join_path, normalize_path, sys_to_unix, unix_to_sys};
-pub use virtual_fs::VirtualFileSystem;
-use watcher::DiskWatcher;
 
 use self::{invalidation::Write, json::UnparsableJson, mutex_map::MutexMap};
 use crate::{
     attach::AttachedFileSystem,
-    invalidator_map::WriteContent,
+    glob::Glob,
+    invalidator_map::{InvalidatorMap, WriteContent},
+    read_glob::{read_glob, track_glob},
     retry::retry_blocking,
     rope::{Rope, RopeReader},
+    util::{extract_disk_access, join_path, normalize_path, sys_to_unix, unix_to_sys},
+    watcher::DiskWatcher,
 };
+pub use crate::{read_glob::ReadGlobResult, virtual_fs::VirtualFileSystem};
 
 /// A (somewhat arbitrary) filename limit that we should try to keep output file names below.
 ///

filename: turbopack/crates/turbo-tasks-fs/src/path_map.rs
patch: @@ -0,0 +1,80 @@
+use std::{
+    collections::{BTreeMap, btree_map::CursorMut},
+    ops::Bound,
+    path::{Path, PathBuf},
+};
+
+/// A thin wrapper around [`BTreeMap<PathBuf, V>`] that provides efficient extraction of child
+/// paths.
+///
+/// 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>;
+}
+
+impl<V> OrderedPathMapExt<V> for BTreeMap<PathBuf, V> {
+    fn extract_path_with_children<'a>(&'a mut self, path: &'a Path) -> ExtractWithChildren<'a, V> {
+        ExtractWithChildren {
+            cursor: self.lower_bound_mut(Bound::Included(path)),
+            parent_path: path,
+        }
+    }
+}
+
+pub struct ExtractWithChildren<'a, V> {
+    cursor: CursorMut<'a, PathBuf, V>,
+    parent_path: &'a Path,
+}
+
+impl<V> Iterator for ExtractWithChildren<'_, V> {
+    type Item = (PathBuf, V);
+
+    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).
+        if self
+            .cursor
+            .peek_next()
+            .is_none_or(|(k, _v)| !k.starts_with(self.parent_path))
+        {
+            return None;
+        }
+        self.cursor.remove_next()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_extract_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);
+
+        let parent_path = PathBuf::from("a/b");
+        let extracted: Vec<_> = map.extract_path_with_children(&parent_path).collect();
+
+        let expected_extracted = vec![
+            (PathBuf::from("a/b"), 2),
+            (PathBuf::from("a/b/c"), 3),
+            (PathBuf::from("a/b/d"), 4),
+        ];
+        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);
+
+        assert_eq!(map, expected_remaining);
+    }
+}

filename: turbopack/crates/turbo-tasks-fs/src/watcher.rs
patch: @@ -17,7 +17,7 @@ use notify::{
     event::{MetadataKind, ModifyKind, RenameMode},
 };
 use rayon::iter::{IntoParallelIterator, ParallelIterator};
-use rustc_hash::{FxHashMap, FxHashSet};
+use rustc_hash::FxHashSet;
 use serde::{Deserialize, Serialize};
 use tracing::instrument;
 use turbo_rcstr::RcStr;
@@ -29,7 +29,8 @@ use turbo_tasks::{
 use crate::{
     DiskFileSystemInner, format_absolute_fs_path,
     invalidation::{WatchChange, WatchStart},
-    invalidator_map::WriteContent,
+    invalidator_map::LockedInvalidatorMap,
+    path_map::OrderedPathMapExt,
 };
 
 static WATCH_RECURSIVE_MODE: LazyLock<RecursiveMode> = LazyLock::new(|| {
@@ -602,7 +603,7 @@ fn invalidate(
 fn invalidate_path(
     inner: &DiskFileSystemInner,
     report_invalidation_reason: bool,
-    invalidator_map: &mut FxHashMap<PathBuf, FxHashMap<Invalidator, Option<WriteContent>>>,
+    invalidator_map: &mut LockedInvalidatorMap,
     paths: impl Iterator<Item = PathBuf>,
 ) {
     for path in paths {
@@ -617,11 +618,11 @@ fn invalidate_path(
 fn invalidate_path_and_children_execute(
     inner: &DiskFileSystemInner,
     report_invalidation_reason: bool,
-    invalidator_map: &mut FxHashMap<PathBuf, FxHashMap<Invalidator, Option<WriteContent>>>,
+    invalidator_map: &mut LockedInvalidatorMap,
     paths: impl Iterator<Item = PathBuf>,
 ) {
     for path in paths {
-        for (_, invalidators) in invalidator_map.extract_if(|key, _| key.starts_with(&path)) {
+        for (_, invalidators) in invalidator_map.extract_path_with_children(&path) {
             invalidators
                 .into_iter()
                 .for_each(|(i, _)| invalidate(inner, report_invalidation_reason, &path, i));