feat(rln): close db connection before dropping rln object (#187)

Cargo.lock

@@ -1532,7 +1532,8 @@ dependencies = [
 [[package]]
 name = "ethers-core"
 version = "2.0.8"
-source = "git+https://github.com/gakonst/ethers-rs?rev=030bf43#030bf439a100dcacd2e968e3114de0612229056b"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "60ca2514feb98918a0a31de7e1983c29f2267ebf61b2dc5d4294f91e5b866623"
 dependencies = [
  "arrayvec",
  "bytes",
@@ -2413,7 +2414,7 @@ dependencies = [
 [[package]]
 name = "pmtree"
 version = "1.0.0"
-source = "git+https://github.com/vacp2p/pmtree?rev=8b005bb44a4c51b56638ad6057b8bc90629fca59#8b005bb44a4c51b56638ad6057b8bc90629fca59"
+source = "git+https://github.com/vacp2p/pmtree?rev=46a39a3#46a39a373ffdb6da14fd122b008437e63f1f7e2b"
 dependencies = [
  "rayon",
 ]

rln/src/ffi.rs

@@ -401,6 +401,12 @@ pub extern "C" fn get_metadata(ctx: *const RLN, output_buffer: *mut Buffer) -> b
     call_with_output_arg!(ctx, get_metadata, output_buffer)
 }
 
+#[allow(clippy::not_unsafe_ptr_arg_deref)]
+#[no_mangle]
+pub extern "C" fn close_db_connection(ctx: *mut RLN) -> bool {
+    call!(ctx, close_db_connection)
+}
+
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
 #[no_mangle]
 pub extern "C" fn hash(input_buffer: *const Buffer, output_buffer: *mut Buffer) -> bool {

rln/src/pm_tree_adapter.rs

@@ -145,6 +145,10 @@ impl ZerokitMerkleTree for PmTree {
         })
     }
 
+    fn close_db_connection(&mut self) -> Result<()> {
+        self.tree.db.close().map_err(|e| Report::msg(e.to_string()))
+    }
+
     fn depth(&self) -> usize {
         self.tree.depth()
     }

rln/src/public.rs

@@ -1129,6 +1129,14 @@ impl RLN<'_> {
         let (rln_witness, _) = deserialize_witness(serialized_witness)?;
         get_json_inputs(&rln_witness)
     }
+
+    /// Closes the connection to the Merkle tree database.
+    /// This function should be called before the RLN object is dropped.
+    /// If not called, the connection will be closed when the RLN object is dropped.
+    /// This improves robustness of the tree.
+    pub fn close_db_connection(&mut self) -> Result<()> {
+        self.tree.close_db_connection()
+    }
 }
 
 #[cfg(not(target_arch = "wasm32"))]
@@ -1378,6 +1386,8 @@ mod test {
         let (root_single_additions, _) = bytes_le_to_fr(&buffer.into_inner());
 
         assert_eq!(root_batch_with_init, root_single_additions);
+
+        rln.close_db_connection();
     }
 
     #[test]

utils/Cargo.toml

@@ -11,7 +11,7 @@ bench = false
 ark-ff = { version = "=0.4.1", default-features = false, features = ["asm"] }
 num-bigint = { version = "=0.4.3", default-features = false, features = ["rand"] }
 color-eyre = "=0.6.2"
-pmtree = { git = "https://github.com/vacp2p/pmtree", rev = "8b005bb44a4c51b56638ad6057b8bc90629fca59", optional = true}
+pmtree = { git = "https://github.com/vacp2p/pmtree", rev = "46a39a3", optional = true}
 sled = "=0.34.7"
 serde = "1.0.44"
 

utils/src/merkle_tree/full_merkle_tree.rs

@@ -101,6 +101,10 @@ where
         })
     }
 
+    fn close_db_connection(&mut self) -> Result<()> {
+        Ok(())
+    }
+
     // Returns the depth of the tree
     fn depth(&self) -> usize {
         self.depth

utils/src/merkle_tree/merkle_tree.rs

@@ -65,6 +65,7 @@ pub trait ZerokitMerkleTree {
     fn verify(&self, leaf: &FrOf<Self::Hasher>, witness: &Self::Proof) -> Result<bool>;
     fn set_metadata(&mut self, metadata: &[u8]) -> Result<()>;
     fn metadata(&self) -> Result<Vec<u8>>;
+    fn close_db_connection(&mut self) -> Result<()>;
 }
 
 pub trait ZerokitMerkleProof {

utils/src/merkle_tree/optimal_merkle_tree.rs

@@ -83,6 +83,10 @@ where
         })
     }
 
+    fn close_db_connection(&mut self) -> Result<()> {
+        Ok(())
+    }
+
     // Returns the depth of the tree
     fn depth(&self) -> usize {
         self.depth

utils/src/pm_tree/sled_adapter.rs

@@ -45,6 +45,15 @@ impl Database for SledDB {
         Ok(SledDB(db))
     }
 
+    fn close(&mut self) -> PmtreeResult<()> {
+        let _ = self.0.flush().map_err(|_| {
+            PmtreeErrorKind::DatabaseError(DatabaseErrorKind::CustomError(
+                "Cannot flush database".to_string(),
+            ))
+        })?;
+        Ok(())
+    }
+
     fn get(&self, key: DBKey) -> PmtreeResult<Option<Value>> {
         match self.0.get(key) {
             Ok(value) => Ok(value.map(|val| val.to_vec())),