fix(native): possible deadlock when batching read/write (#9817)

packages/frontend/native/nbstore/src/pool.rs

@@ -1,54 +1,96 @@
-use dashmap::{mapref::one::RefMut, DashMap, Entry};
+use core::ops::{Deref, DerefMut};
+use std::collections::hash_map::{Entry, HashMap};
+
+use tokio::sync::{RwLock, RwLockMappedWriteGuard, RwLockReadGuard, RwLockWriteGuard};
 
 use super::{
   error::{Error, Result},
   storage::SqliteDocStorage,
 };
 
+pub struct Ref<'a, V> {
+  _guard: RwLockReadGuard<'a, V>,
+}
+
+impl<'a, V> Deref for Ref<'a, V> {
+  type Target = V;
+
+  fn deref(&self) -> &Self::Target {
+    self._guard.deref()
+  }
+}
+
+pub struct RefMut<'a, V> {
+  _guard: RwLockMappedWriteGuard<'a, V>,
+}
+
+impl<'a, V> Deref for RefMut<'a, V> {
+  type Target = V;
+
+  fn deref(&self) -> &Self::Target {
+    &*self._guard
+  }
+}
+
+impl<'a, V> DerefMut for RefMut<'a, V> {
+  fn deref_mut(&mut self) -> &mut Self::Target {
+    &mut *self._guard
+  }
+}
+
 #[derive(Default)]
 pub struct SqliteDocStoragePool {
-  inner: DashMap<String, SqliteDocStorage>,
+  inner: RwLock<HashMap<String, SqliteDocStorage>>,
 }
 
 impl SqliteDocStoragePool {
-  fn get_or_create_storage<'a>(
+  async fn get_or_create_storage<'a>(
     &'a self,
     universal_id: String,
     path: &str,
-  ) -> RefMut<'a, String, SqliteDocStorage> {
-    let entry = self.inner.entry(universal_id);
-    if let Entry::Occupied(storage) = entry {
-      return storage.into_ref();
-    }
-    let storage = SqliteDocStorage::new(path.to_string());
+  ) -> RefMut<'a, SqliteDocStorage> {
+    let lock = RwLockWriteGuard::map(self.inner.write().await, |lock| {
+      match lock.entry(universal_id) {
+        Entry::Occupied(entry) => entry.into_mut(),
+        Entry::Vacant(entry) => {
+          let storage = SqliteDocStorage::new(path.to_string());
+          entry.insert(storage)
+        }
+      }
+    });
 
-    entry.or_insert(storage)
+    RefMut { _guard: lock }
   }
 
-  pub fn ensure_storage<'a>(
-    &'a self,
-    universal_id: String,
-  ) -> Result<RefMut<'a, String, SqliteDocStorage>> {
-    let entry = self.inner.entry(universal_id);
+  pub async fn get<'a>(&'a self, universal_id: String) -> Result<Ref<'a, SqliteDocStorage>> {
+    let lock = RwLockReadGuard::try_map(self.inner.read().await, |lock| {
+      if let Some(storage) = lock.get(&universal_id) {
+        Some(storage)
+      } else {
+        None
+      }
+    });
 
-    if let Entry::Occupied(storage) = entry {
-      Ok(storage.into_ref())
-    } else {
-      Err(Error::InvalidOperation)
+    match lock {
+      Ok(guard) => Ok(Ref { _guard: guard }),
+      Err(_) => Err(Error::InvalidOperation),
     }
   }
 
   /// Initialize the database and run migrations.
   pub async fn connect(&self, universal_id: String, path: String) -> Result<()> {
-    let storage = self.get_or_create_storage(universal_id.to_owned(), &path);
+    let storage = self
+      .get_or_create_storage(universal_id.to_owned(), &path)
+      .await;
+
     storage.connect().await?;
     Ok(())
   }
 
   pub async fn disconnect(&self, universal_id: String) -> Result<()> {
-    let entry = self.inner.entry(universal_id);
+    let mut lock = self.inner.write().await;
 
-    if let Entry::Occupied(entry) = entry {
+    if let Entry::Occupied(entry) = lock.entry(universal_id) {
       let storage = entry.remove();
       storage.close().await;
     }