Update shared hashmap to use RwLock instead of threadlocal

crates/atlaspack_filesystem/src/lib.rs

@@ -2,7 +2,7 @@ use std::path::Path;
 use std::path::PathBuf;
 use std::sync::Arc;
 
-use atlaspack_shared_map::ThreadLocalHashMap;
+use atlaspack_shared_map::SharedHashMap;
 
 /// In-memory file-system for testing
 pub mod in_memory_file_system;
@@ -17,9 +17,7 @@ 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 DefaultHasher = xxhash_rust::xxh3::Xxh3Builder;
-
-pub type FileSystemRealPathCache = ThreadLocalHashMap<PathBuf, Option<PathBuf>, DefaultHasher>;
+pub type FileSystemRealPathCache = SharedHashMap<PathBuf, Option<PathBuf>>;
 
 /// Trait abstracting file-system operations
 /// .

crates/atlaspack_shared_map/src/lib.rs

@@ -1,46 +1,37 @@
-use std::{
-  borrow::Borrow,
-  cell::RefCell,
-  collections::HashMap,
-  hash::BuildHasher,
-  ops::{Deref, DerefMut},
-};
+use std::{borrow::Borrow, collections::HashMap, hash::BuildHasher, sync::RwLock};
 
-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<
+/// A HashMap that allows sharing of values between threads
+/// It is currently implemented using RwLock
+pub struct SharedHashMap<
   K: Send + Eq,
   V: Send + Clone,
   H: Send + Default + BuildHasher = xxhash_rust::xxh3::Xxh3Builder,
 > {
-  inner: ThreadLocal<RefCell<HashMap<K, V, H>>>,
+  inner: RwLock<HashMap<K, V, H>>,
 }
 
-impl<K, V, H> Default for ThreadLocalHashMap<K, V, H>
+impl<K, V, H> Default for SharedHashMap<K, V, H>
 where
   K: Send + Eq,
   V: Send + Clone,
   H: Send + Default + BuildHasher,
 {
   fn default() -> Self {
     Self {
-      inner: ThreadLocal::new(),
+      inner: RwLock::new(HashMap::with_hasher(H::default())),
     }
   }
 }
 
-impl<K, V, H> ThreadLocalHashMap<K, V, H>
+impl<K, V, H> SharedHashMap<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(),
+      inner: RwLock::new(HashMap::with_hasher(H::default())),
     }
   }
 
@@ -50,21 +41,15 @@ where
     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()
+    let map_cell = self.inner.read().unwrap();
+    let map = map_cell;
+    map.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);
+    let map_cell = self.inner.write().unwrap();
+    let mut map = map_cell;
+    map.insert(key, value);
   }
 }
 
@@ -75,15 +60,15 @@ mod test {
   use super::*;
 
   #[test]
-  fn test_thread_local_hash_map() {
-    let map = ThreadLocalHashMap::<String, String>::new();
+  fn test_shared_hash_map() {
+    let map = SharedHashMap::<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());
+  fn test_multiple_thread_shared_hash_map() {
+    let map = Arc::new(SharedHashMap::<String, String>::new());
     let (close_tx, close_rx) = std::sync::mpsc::channel();
     let (tx, rx) = std::sync::mpsc::channel();
 
@@ -98,11 +83,11 @@ mod test {
     });
 
     rx.recv().unwrap(); // wait for value to be written
-                        // and check value is not visible on other threads
+                        // and check value is visible on other threads
     std::thread::spawn({
       let map = map.clone();
       move || {
-        assert_eq!(map.get("key"), None);
+        assert_eq!(map.get("key"), Some("value".to_string()));
       }
     })
     .join()

packages/utils/node-resolver-rs/src/cache.rs

@@ -7,7 +7,7 @@ use std::sync::Arc;
 
 use atlaspack_core::types::File;
 use atlaspack_filesystem::{FileSystemRealPathCache, FileSystemRef};
-use atlaspack_shared_map::ThreadLocalHashMap;
+use atlaspack_shared_map::SharedHashMap;
 
 use crate::package_json::PackageJson;
 use crate::package_json::SourceField;
@@ -22,13 +22,13 @@ pub struct Cache {
   /// from our public methods so this is ok for now. FrozenMap is an append only map, which doesn't require &mut
   /// to insert into. Since each value is in a Box, it won't move and therefore references are stable.
   #[allow(clippy::type_complexity)]
-  packages: ThreadLocalHashMap<PathBuf, Arc<Result<Arc<PackageJson>, ResolverError>>>,
+  packages: SharedHashMap<PathBuf, Arc<Result<Arc<PackageJson>, ResolverError>>>,
   #[allow(clippy::type_complexity)]
-  tsconfigs: ThreadLocalHashMap<PathBuf, Arc<Result<Arc<TsConfigWrapper>, ResolverError>>>,
+  tsconfigs: SharedHashMap<PathBuf, Arc<Result<Arc<TsConfigWrapper>, ResolverError>>>,
   // In particular just the is_dir_cache spends around 8% of the time on a large project resolution
   // hashing paths. Instead of using a hashmap we should try a trie here.
-  is_dir_cache: ThreadLocalHashMap<PathBuf, bool>,
-  is_file_cache: ThreadLocalHashMap<PathBuf, bool>,
+  is_dir_cache: SharedHashMap<PathBuf, bool>,
+  is_file_cache: SharedHashMap<PathBuf, bool>,
   realpath_cache: FileSystemRealPathCache,
 }
 
@@ -91,10 +91,10 @@ impl Cache {
   pub fn new(fs: FileSystemRef) -> Self {
     Self {
       fs,
-      packages: ThreadLocalHashMap::new(),
-      tsconfigs: ThreadLocalHashMap::new(),
-      is_file_cache: ThreadLocalHashMap::new(),
-      is_dir_cache: ThreadLocalHashMap::new(),
+      packages: SharedHashMap::new(),
+      tsconfigs: SharedHashMap::new(),
+      is_file_cache: SharedHashMap::new(),
+      is_dir_cache: SharedHashMap::new(),
       realpath_cache: FileSystemRealPathCache::default(),
     }
   }