Edit to fix increacing memory usage (#93)

🐛 Edit to fix increacing memory usage

 - Also add multi_thread unit test in state_writer
 - Add rand as new dependence to create some random vector
 - Add close function for ReadWriter in JS
 - Add new unit test in JS to check close functionality

Cargo.toml

@@ -43,6 +43,9 @@ version = "0.3"
 [dependencies.bitvec]
 version = "1.0.1"
 
+[dev-dependencies.rand]
+version = "0.8.5"
+
 [dev-dependencies.criterion]
 version = "0.4.0"
 

benchmark/state_long.js

@@ -59,6 +59,7 @@ const count = 10000;
         performance.mark('c-start');
         console.log({ i, root })
         root = await db.commit(writer, i, root);
+        writer.close();
         console.log({ root })
         performance.mark('c-end');
 
@@ -69,4 +70,4 @@ const count = 10000;
     }
 
     db.close();
-})()
+})()

src/database/reader_writer/read_writer_db.rs

@@ -300,4 +300,16 @@ impl ReadWriter {
 
         Ok(ctx.undefined())
     }
+
+    /// js_close is handler for JS ffi.
+    /// js "this" - ReadWriter.
+    pub fn js_close(mut ctx: FunctionContext) -> JsResult<JsUndefined> {
+        let db = ctx
+            .this()
+            .downcast_or_throw::<SharedReadWriter, _>(&mut ctx)?;
+        let db = db.borrow_mut();
+        db.close().or_else(|err| ctx.throw_error(err.to_string()))?;
+
+        Ok(ctx.undefined())
+    }
 }

src/database/reader_writer/reader_base.rs

@@ -19,7 +19,6 @@ pub struct ReaderBase {
 
 impl Finalize for ReaderBase {
     fn finalize<'a, C: Context<'a>>(self, _: &mut C) {
-        self.close().unwrap();
         drop(self);
     }
 }
@@ -55,7 +54,7 @@ impl ReaderBase {
 
     /// Idiomatic rust would take an owned `self` to prevent use after close
     /// However, it's not possible to prevent JavaScript from continuing to hold a closed database
-    fn close(&self) -> Result<(), mpsc::SendError<SnapshotMessage>> {
+    pub fn close(&self) -> Result<(), mpsc::SendError<SnapshotMessage>> {
         self.tx.send(SnapshotMessage::Close)
     }
 

src/lib.rs

@@ -52,6 +52,7 @@ fn main(mut cx: ModuleContext) -> NeonResult<()> {
     cx.export_function("state_db_reader_iterate", reader_db::Reader::js_iterate)?;
 
     cx.export_function("state_db_read_writer_new", ReadWriter::js_new)?;
+    cx.export_function("state_db_read_writer_close", ReadWriter::js_close)?;
     cx.export_function("state_db_read_writer_upsert_key", ReadWriter::js_upsert_key)?;
     cx.export_function("state_db_read_writer_get_key", ReadWriter::js_get_key)?;
     cx.export_function("state_db_read_writer_delete", ReadWriter::js_delete_key)?;
@@ -79,6 +80,7 @@ fn main(mut cx: ModuleContext) -> NeonResult<()> {
     let state_writer_new = StateWriter::js_new_with_arc_mutex::<StateWriter>;
     let restore_snapshot = StateWriter::js_restore_snapshot;
     cx.export_function("state_writer_new", state_writer_new)?;
+    cx.export_function("state_writer_close", StateWriter::js_close)?;
     cx.export_function("state_writer_snapshot", StateWriter::js_snapshot)?;
     cx.export_function("state_writer_restore_snapshot", restore_snapshot)?;
 

src/sparse_merkle_tree/smt.rs

@@ -505,14 +505,12 @@ impl UpdateData {
     pub fn new_with_hash(data: Cache) -> Self {
         let mut new_data = Cache::new();
         for (k, v) in data {
+            let mut value: Vec<u8> = vec![];
             if !v.is_empty() {
-                new_data.insert(
-                    k.hash_with_kind(HashKind::Key),
-                    v.hash_with_kind(HashKind::Value),
-                );
-            } else {
-                new_data.insert(k.hash_with_kind(HashKind::Key), vec![]);
+                value = v.hash_with_kind(HashKind::Value);
             }
+            let key = k.hash_with_kind(HashKind::Key);
+            new_data.insert(key, value);
         }
         Self { data: new_data }
     }

src/state/state_writer.rs

@@ -87,6 +87,12 @@ impl StateCache {
 }
 
 impl StateWriter {
+    /// empty makes StateWriter as an empty HashMap to handle of releasing the memory from JS.
+    fn empty(&mut self) {
+        self.backup = HashMap::new();
+        self.cache = HashMap::new();
+    }
+
     /// cache_new inserts key-value pair as new value.
     pub fn cache_new(&mut self, pair: &SharedKVPair) {
         let cache = StateCache::new(pair.value());
@@ -229,6 +235,20 @@ impl StateWriter {
 }
 
 impl StateWriter {
+    /// js_close is handler for JS ffi.
+    /// js "this" - StateWriter.
+    pub fn js_close(mut ctx: FunctionContext) -> JsResult<JsUndefined> {
+        let writer = ctx
+            .this()
+            .downcast_or_throw::<SendableStateWriter, _>(&mut ctx)?;
+
+        let batch = Arc::clone(&writer.borrow());
+        let mut inner_writer = batch.lock().unwrap();
+        inner_writer.empty();
+
+        Ok(ctx.undefined())
+    }
+
     /// js_snapshot is handler for JS ffi.
     /// js "this" - StateWriter.
     /// - @returns - snapshot id
@@ -269,6 +289,65 @@ mod tests {
     use super::*;
     use crate::consts::Prefix;
 
+    use std::cell::RefCell;
+    use std::convert::TryInto;
+    use std::sync::Mutex;
+    use std::thread;
+
+    use rand::RngCore;
+
+    #[test]
+    fn test_multi_thread() {
+        let outer_loop_iterations = 1000;
+        let inner_loop_iteration = 2000;
+        let pairs_len = inner_loop_iteration / 10;
+
+        for _ in 0..outer_loop_iterations {
+            let mut pairs: Vec<KVPair> = vec![];
+            for _ in 0..pairs_len {
+                let mut key = [0u8; 32];
+                let mut value = [0u8; 32];
+                rand::thread_rng().fill_bytes(&mut key);
+                rand::thread_rng().fill_bytes(&mut value);
+                pairs.push(KVPair::new(&key, &value));
+            }
+
+            let sendable_writer = RefCell::new(Arc::new(Mutex::new(StateWriter::default())));
+            let mut counter = 0;
+            for i in 1..inner_loop_iteration {
+                let mut key = [0u8; 32];
+                let mut value = [0u8; 32];
+                if i % 6 == 0 && counter < pairs_len {
+                    key = pairs[counter].key().try_into().unwrap();
+                    value = pairs[counter].value().try_into().unwrap();
+                    counter += 1;
+                } else {
+                    rand::thread_rng().fill_bytes(&mut key);
+                    rand::thread_rng().fill_bytes(&mut value);
+                }
+
+                let batch = sendable_writer.borrow_mut();
+                let writer = Arc::clone(&batch);
+                thread::spawn(move || {
+                    let w = writer.lock();
+                    assert!(w.is_ok());
+                    w.unwrap().cache_new(&SharedKVPair::new(&key, &value));
+                });
+            }
+
+            let writer = Arc::clone(&sendable_writer.borrow_mut());
+            thread::spawn(move || {
+                let mut w = writer.lock().unwrap();
+                for kv in pairs.iter() {
+                    assert!(w.is_cached(kv.key()));
+                    let mut new_value = [0u8; 32];
+                    rand::thread_rng().fill_bytes(&mut new_value);
+                    assert!(w.update(&KVPair::new(kv.key(), &new_value)).is_ok());
+                }
+            });
+        }
+    }
+
     #[test]
     fn test_cache() {
         let mut writer = StateWriter::default();

state_db.js

@@ -28,13 +28,15 @@ const {
     state_db_checkpoint,
     state_db_calculate_root,
     state_writer_new,
+    state_writer_close,
     state_writer_snapshot,
     state_writer_restore_snapshot,
     state_db_reader_new,
     state_db_reader_get,
     state_db_reader_exists,
     state_db_reader_iterate,
     state_db_read_writer_new,
+    state_db_read_writer_close,
     state_db_read_writer_upsert_key,
     state_db_read_writer_get_key,
     state_db_read_writer_delete,
@@ -100,6 +102,11 @@ class StateReadWriter {
         return this._writer;
     }
 
+    close() {
+        state_db_read_writer_close.call(this._db);
+        state_writer_close.call(this.writer);
+    }
+
     async get(key) {
         const value = await new Promise((resolve, reject) => {
             state_db_read_writer_get_key.call(this._db, this.writer, key, (err, result) => {

test/statedb.spec.js

@@ -452,6 +452,15 @@ describe('statedb', () => {
                 await writer.set(initState[1].key, getRandomBytes());
                 expect(() => writer.restoreSnapshot(99)).toThrow('Invalid usage');
             });
+
+            it('should throw an error when the writer is closed', async () => {
+                const writer = db.newReadWriter();
+                const newValue = getRandomBytes();
+                await writer.set(initState[1].key, newValue);
+                await expect(writer.get(initState[1].key)).resolves.toEqual(newValue);
+                writer.close();
+                expect(() => writer.get(initState[1].key)).rejects.toThrow();
+            });
         });
 
         describe('StateReader', () => {

types.d.ts

@@ -86,6 +86,7 @@ declare class StateReadWriter {
     range(options?: IterateOptions): Promise<{ key: Buffer, value: Buffer }[]>;
     snapshot(): number;
     restoreSnapshot(index: number): void;
+    close(): void;
 }
 
 interface StateCommitOption {