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

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

@@ -8,7 +8,7 @@ pub mod sync;
 use chrono::NaiveDateTime;
 use napi::bindgen_prelude::*;
 use napi_derive::napi;
-use pool::SqliteDocStoragePool;
+use pool::{Ref, SqliteDocStoragePool};
 use storage::SqliteDocStorage;
 
 #[cfg(feature = "use-as-lib")]
@@ -86,13 +86,17 @@ pub struct DocStoragePool {
 
 #[napi]
 impl DocStoragePool {
-  #[napi(constructor, async_runtime)]
+  #[napi(constructor)]
   pub fn new() -> Result<Self> {
     Ok(Self {
       pool: SqliteDocStoragePool::default(),
     })
   }
 
+  async fn get<'a>(&'a self, universal_id: String) -> Result<Ref<'a, SqliteDocStorage>> {
+    Ok(self.pool.get(universal_id).await?)
+  }
+
   #[napi]
   /// Initialize the database and run migrations.
   pub async fn connect(&self, universal_id: String, path: String) -> Result<()> {
@@ -100,16 +104,6 @@ impl DocStoragePool {
     Ok(())
   }
 
-  #[napi]
-  pub async fn set_space_id(&self, universal_id: String, space_id: String) -> Result<()> {
-    self
-      .pool
-      .ensure_storage(universal_id)?
-      .set_space_id(space_id)
-      .await?;
-    Ok(())
-  }
-
   #[napi]
   pub async fn disconnect(&self, universal_id: String) -> Result<()> {
     self.pool.disconnect(universal_id).await?;
@@ -118,7 +112,13 @@ impl DocStoragePool {
 
   #[napi]
   pub async fn checkpoint(&self, universal_id: String) -> Result<()> {
-    self.pool.ensure_storage(universal_id)?.checkpoint().await?;
+    self.pool.get(universal_id).await?.checkpoint().await?;
+    Ok(())
+  }
+
+  #[napi]
+  pub async fn set_space_id(&self, universal_id: String, space_id: String) -> Result<()> {
+    self.get(universal_id).await?.set_space_id(space_id).await?;
     Ok(())
   }
 
@@ -131,8 +131,8 @@ impl DocStoragePool {
   ) -> Result<NaiveDateTime> {
     Ok(
       self
-        .pool
-        .ensure_storage(universal_id)?
+        .get(universal_id)
+        .await?
         .push_update(doc_id, update)
         .await?,
     )
@@ -146,8 +146,8 @@ impl DocStoragePool {
   ) -> Result<Option<DocRecord>> {
     Ok(
       self
-        .pool
-        .ensure_storage(universal_id)?
+        .get(universal_id)
+        .await?
         .get_doc_snapshot(doc_id)
         .await?,
     )
@@ -157,8 +157,8 @@ impl DocStoragePool {
   pub async fn set_doc_snapshot(&self, universal_id: String, snapshot: DocRecord) -> Result<bool> {
     Ok(
       self
-        .pool
-        .ensure_storage(universal_id)?
+        .get(universal_id)
+        .await?
         .set_doc_snapshot(snapshot)
         .await?,
     )
@@ -172,8 +172,8 @@ impl DocStoragePool {
   ) -> Result<Vec<DocUpdate>> {
     Ok(
       self
-        .pool
-        .ensure_storage(universal_id)?
+        .get(universal_id)
+        .await?
         .get_doc_updates(doc_id)
         .await?,
     )
@@ -188,8 +188,8 @@ impl DocStoragePool {
   ) -> Result<u32> {
     Ok(
       self
-        .pool
-        .ensure_storage(universal_id)?
+        .get(universal_id)
+        .await?
         .mark_updates_merged(doc_id, updates)
         .await?,
     )
@@ -197,11 +197,7 @@ impl DocStoragePool {
 
   #[napi]
   pub async fn delete_doc(&self, universal_id: String, doc_id: String) -> Result<()> {
-    self
-      .pool
-      .ensure_storage(universal_id)?
-      .delete_doc(doc_id)
-      .await?;
+    self.get(universal_id).await?.delete_doc(doc_id).await?;
     Ok(())
   }
 
@@ -211,13 +207,7 @@ impl DocStoragePool {
     universal_id: String,
     after: Option<NaiveDateTime>,
   ) -> Result<Vec<DocClock>> {
-    Ok(
-      self
-        .pool
-        .ensure_storage(universal_id)?
-        .get_doc_clocks(after)
-        .await?,
-    )
+    Ok(self.get(universal_id).await?.get_doc_clocks(after).await?)
   }
 
   #[napi]
@@ -226,33 +216,17 @@ impl DocStoragePool {
     universal_id: String,
     doc_id: String,
   ) -> Result<Option<DocClock>> {
-    Ok(
-      self
-        .pool
-        .ensure_storage(universal_id)?
-        .get_doc_clock(doc_id)
-        .await?,
-    )
+    Ok(self.get(universal_id).await?.get_doc_clock(doc_id).await?)
   }
 
-  #[napi]
+  #[napi(async_runtime)]
   pub async fn get_blob(&self, universal_id: String, key: String) -> Result<Option<Blob>> {
-    Ok(
-      self
-        .pool
-        .ensure_storage(universal_id)?
-        .get_blob(key)
-        .await?,
-    )
+    Ok(self.get(universal_id).await?.get_blob(key).await?)
   }
 
   #[napi]
   pub async fn set_blob(&self, universal_id: String, blob: SetBlob) -> Result<()> {
-    self
-      .pool
-      .ensure_storage(universal_id)?
-      .set_blob(blob)
-      .await?;
+    self.get(universal_id).await?.set_blob(blob).await?;
     Ok(())
   }
 
@@ -264,8 +238,8 @@ impl DocStoragePool {
     permanently: bool,
   ) -> Result<()> {
     self
-      .pool
-      .ensure_storage(universal_id)?
+      .get(universal_id)
+      .await?
       .delete_blob(key, permanently)
       .await?;
     Ok(())
@@ -273,17 +247,13 @@ impl DocStoragePool {
 
   #[napi]
   pub async fn release_blobs(&self, universal_id: String) -> Result<()> {
-    self
-      .pool
-      .ensure_storage(universal_id)?
-      .release_blobs()
-      .await?;
+    self.get(universal_id).await?.release_blobs().await?;
     Ok(())
   }
 
   #[napi]
   pub async fn list_blobs(&self, universal_id: String) -> Result<Vec<ListedBlob>> {
-    Ok(self.pool.ensure_storage(universal_id)?.list_blobs().await?)
+    Ok(self.get(universal_id).await?.list_blobs().await?)
   }
 
   #[napi]
@@ -294,8 +264,8 @@ impl DocStoragePool {
   ) -> Result<Vec<DocClock>> {
     Ok(
       self
-        .pool
-        .ensure_storage(universal_id)?
+        .get(universal_id)
+        .await?
         .get_peer_remote_clocks(peer)
         .await?,
     )
@@ -310,8 +280,8 @@ impl DocStoragePool {
   ) -> Result<Option<DocClock>> {
     Ok(
       self
-        .pool
-        .ensure_storage(universal_id)?
+        .get(universal_id)
+        .await?
         .get_peer_remote_clock(peer, doc_id)
         .await?,
     )
@@ -326,8 +296,8 @@ impl DocStoragePool {
     clock: NaiveDateTime,
   ) -> Result<()> {
     self
-      .pool
-      .ensure_storage(universal_id)?
+      .get(universal_id)
+      .await?
       .set_peer_remote_clock(peer, doc_id, clock)
       .await?;
     Ok(())
@@ -341,8 +311,8 @@ impl DocStoragePool {
   ) -> Result<Vec<DocClock>> {
     Ok(
       self
-        .pool
-        .ensure_storage(universal_id)?
+        .get(universal_id)
+        .await?
         .get_peer_pulled_remote_clocks(peer)
         .await?,
     )
@@ -357,8 +327,8 @@ impl DocStoragePool {
   ) -> Result<Option<DocClock>> {
     Ok(
       self
-        .pool
-        .ensure_storage(universal_id)?
+        .get(universal_id)
+        .await?
         .get_peer_pulled_remote_clock(peer, doc_id)
         .await?,
     )
@@ -373,8 +343,8 @@ impl DocStoragePool {
     clock: NaiveDateTime,
   ) -> Result<()> {
     self
-      .pool
-      .ensure_storage(universal_id)?
+      .get(universal_id)
+      .await?
       .set_peer_pulled_remote_clock(peer, doc_id, clock)
       .await?;
     Ok(())
@@ -388,8 +358,8 @@ impl DocStoragePool {
   ) -> Result<Vec<DocClock>> {
     Ok(
       self
-        .pool
-        .ensure_storage(universal_id)?
+        .get(universal_id)
+        .await?
         .get_peer_pushed_clocks(peer)
         .await?,
     )
@@ -404,8 +374,8 @@ impl DocStoragePool {
   ) -> Result<Option<DocClock>> {
     Ok(
       self
-        .pool
-        .ensure_storage(universal_id)?
+        .get(universal_id)
+        .await?
         .get_peer_pushed_clock(peer, doc_id)
         .await?,
     )
@@ -420,8 +390,8 @@ impl DocStoragePool {
     clock: NaiveDateTime,
   ) -> Result<()> {
     self
-      .pool
-      .ensure_storage(universal_id)?
+      .get(universal_id)
+      .await?
       .set_peer_pushed_clock(peer, doc_id, clock)
       .await?;
     Ok(())
@@ -429,11 +399,7 @@ impl DocStoragePool {
 
   #[napi]
   pub async fn clear_clocks(&self, universal_id: String) -> Result<()> {
-    self
-      .pool
-      .ensure_storage(universal_id)?
-      .clear_clocks()
-      .await?;
+    self.get(universal_id).await?.clear_clocks().await?;
     Ok(())
   }
 }
@@ -461,7 +427,6 @@ impl DocStorage {
   pub async fn set_space_id(&self, space_id: String) -> Result<()> {
     self.storage.connect().await?;
     self.storage.set_space_id(space_id).await?;
-    println!("clocks {:?}", self.storage.get_doc_clocks(None).await?);
     self.storage.close().await;
     Ok(())
   }