Optimize resolver with threadlocal caches (#156)

Cargo.toml

@@ -90,6 +90,7 @@ swc_ecma_parser = "*"
 swc_ecma_transforms_testing = "0.143.2"
 thiserror = "1.0"
 tinyvec = "1.8"
+thread_local = "1.1.8"
 toml = "0.8.12"
 tracing = "0.1"
 tracing-appender = "0.2.3"

crates/atlaspack_filesystem/Cargo.toml

@@ -5,9 +5,11 @@ edition = "2021"
 description = "FileSystem wrapper trait for use in Atlaspack codebase."
 
 [dependencies]
+anyhow = { workspace = true }
+atlaspack_shared_map = { path = "../atlaspack_shared_map" }
 mockall = { workspace = true }
+thread_local = { workspace = true }
 xxhash-rust = { workspace = true, features = ["xxh3"] }
-anyhow = { workspace = true }
 
 [dev-dependencies]
 assert_fs = { workspace = true }

crates/atlaspack_filesystem/src/lib.rs

@@ -1,7 +1,8 @@
-use std::collections::HashMap;
 use std::path::Path;
 use std::path::PathBuf;
-use std::sync::{Arc, RwLock};
+use std::sync::Arc;
+
+use atlaspack_shared_map::ThreadLocalHashMap;
 
 /// In-memory file-system for testing
 pub mod in_memory_file_system;
@@ -16,8 +17,9 @@ pub mod os_file_system;
 /// This should be `OsFileSystem` for non-testing environments and `InMemoryFileSystem` for testing.
 pub type FileSystemRef = Arc<dyn FileSystem + Send + Sync>;
 
-pub type FileSystemRealPathCache =
-  RwLock<HashMap<PathBuf, Option<PathBuf>, xxhash_rust::xxh3::Xxh3Builder>>;
+pub type DefaultHasher = xxhash_rust::xxh3::Xxh3Builder;
+
+pub type FileSystemRealPathCache = ThreadLocalHashMap<PathBuf, Option<PathBuf>, DefaultHasher>;
 
 /// Trait abstracting file-system operations
 /// .

crates/atlaspack_filesystem/src/os_file_system/canonicalize.rs

@@ -28,9 +28,7 @@ pub fn canonicalize(path: &Path, cache: &FileSystemRealPathCache) -> std::io::Re
           ret.push(c);
 
           // First, check the cache for the path up to this point.
-          let read = cache.read().unwrap();
-          let cached = read.get(&ret).cloned();
-          drop(read);
+          let cached = cache.get(&ret);
           let link = if let Some(cached) = cached {
             if let Some(link) = cached {
               link
@@ -40,15 +38,12 @@ pub fn canonicalize(path: &Path, cache: &FileSystemRealPathCache) -> std::io::Re
           } else {
             let stat = std::fs::symlink_metadata(&ret)?;
             if !stat.is_symlink() {
-              cache.write().unwrap().insert(ret.clone(), None);
+              cache.insert(ret.clone(), None);
               continue;
             }
 
             let link = std::fs::read_link(&ret)?;
-            cache
-              .write()
-              .unwrap()
-              .insert(ret.clone(), Some(link.clone()));
+            cache.insert(ret.clone(), Some(link.clone()));
             link
           };
 

crates/atlaspack_shared_map/Cargo.toml

@@ -0,0 +1,10 @@
+[package]
+name = "atlaspack_shared_map"
+description = "A hashmap to be used in caches where consistency between threads is not a concern."
+authors = ["Pedro Tacla Yamada <tacla.yamada@gmail.com>"]
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+thread_local = { workspace = true }
+xxhash-rust = { workspace = true, features = ["xxh3"] }

crates/atlaspack_shared_map/src/lib.rs

@@ -0,0 +1,115 @@
+use std::{
+  borrow::Borrow,
+  cell::RefCell,
+  collections::HashMap,
+  hash::BuildHasher,
+  ops::{Deref, DerefMut},
+};
+
+use thread_local::ThreadLocal;
+
+/// A thread-local hash map.
+/// This is to be used within a shared structure but disabling the need for
+/// actually sharing data between threads.
+pub struct ThreadLocalHashMap<
+  K: Send + Eq,
+  V: Send + Clone,
+  H: Send + Default + BuildHasher = xxhash_rust::xxh3::Xxh3Builder,
+> {
+  inner: ThreadLocal<RefCell<HashMap<K, V, H>>>,
+}
+
+impl<K, V, H> Default for ThreadLocalHashMap<K, V, H>
+where
+  K: Send + Eq,
+  V: Send + Clone,
+  H: Send + Default + BuildHasher,
+{
+  fn default() -> Self {
+    Self {
+      inner: ThreadLocal::new(),
+    }
+  }
+}
+
+impl<K, V, H> ThreadLocalHashMap<K, V, H>
+where
+  K: std::hash::Hash + Send + Eq,
+  V: Send + Clone,
+  H: Send + Default + BuildHasher,
+{
+  pub fn new() -> Self {
+    Self {
+      inner: ThreadLocal::new(),
+    }
+  }
+
+  pub fn get<KR>(&self, key: &KR) -> Option<V>
+  where
+    KR: ?Sized,
+    K: Borrow<KR>,
+    KR: Eq + std::hash::Hash,
+  {
+    let map_cell = self
+      .inner
+      .get_or(|| RefCell::new(HashMap::with_hasher(H::default())));
+
+    let map = map_cell.borrow();
+    map.deref().get(key).cloned()
+  }
+
+  pub fn insert(&self, key: K, value: V) {
+    let map_cell = self
+      .inner
+      .get_or(|| RefCell::new(HashMap::with_hasher(H::default())));
+
+    let mut map = map_cell.borrow_mut();
+    map.deref_mut().insert(key, value);
+  }
+}
+
+#[cfg(test)]
+mod test {
+  use std::sync::Arc;
+
+  use super::*;
+
+  #[test]
+  fn test_thread_local_hash_map() {
+    let map = ThreadLocalHashMap::<String, String>::new();
+    map.insert("key".to_string(), "value".to_string());
+    assert_eq!(map.get("key"), Some("value".to_string()));
+  }
+
+  #[test]
+  fn test_multiple_thread_local_hash_map() {
+    let map = Arc::new(ThreadLocalHashMap::<String, String>::new());
+    let (close_tx, close_rx) = std::sync::mpsc::channel();
+    let (tx, rx) = std::sync::mpsc::channel();
+
+    let thread1 = std::thread::spawn({
+      let map = map.clone();
+      move || {
+        map.insert("key".to_string(), "value".to_string());
+        assert_eq!(map.get("key"), Some("value".to_string()));
+        tx.send(()).unwrap();
+        close_rx.recv().unwrap();
+      }
+    });
+
+    rx.recv().unwrap(); // wait for value to be written
+                        // and check value is not visible on other threads
+    std::thread::spawn({
+      let map = map.clone();
+      move || {
+        assert_eq!(map.get("key"), None);
+      }
+    })
+    .join()
+    .unwrap();
+
+    // stop
+    close_tx.send(()).unwrap();
+    thread1.join().unwrap();
+  }
+}

packages/utils/node-resolver-rs/Cargo.toml

@@ -11,23 +11,25 @@ harness = false
 [dependencies]
 atlaspack_core = { path = "../../../crates/atlaspack_core" }
 atlaspack_filesystem = { path = "../../../crates/atlaspack_filesystem" }
+atlaspack_shared_map = { path = "../../../crates/atlaspack_shared_map" }
 
-bitflags = "1.3.2" # TODO Update dependency
+bitflags = "1.3.2"                                            # TODO Update dependency
 glob-match = { workspace = true }
 indexmap = { workspace = true, features = ["serde"] }
 itertools = { workspace = true }
 json_comments = { path = "../../../crates/json-comments-rs" }
 once_cell = { workspace = true }
 parking_lot = { workspace = true }
-serde_json5 = { workspace = true }
 percent-encoding = { workspace = true }
 serde = { workspace = true, features = ["derive"] }
 serde_json = { workspace = true }
+serde_json5 = { workspace = true }
 thiserror = { workspace = true }
-url = { workspace = true }
-xxhash-rust = { workspace = true, features = ["xxh3"] }
+thread_local = { workspace = true }
 tracing = { workspace = true }
 tracing-subscriber = { workspace = true }
+url = { workspace = true }
+xxhash-rust = { workspace = true, features = ["xxh3"] }
 
 [dev-dependencies]
 assert_fs = { workspace = true }