feat(executor): Add flush for trie db. (nervosnetwork#240)

components/executor/Cargo.toml

@@ -10,12 +10,14 @@ core-storage = { path = "../../core/storage" }
 core-types = { path = "../../core/types" }
 core-context = { path = "../../core/context" }
 
+hashbrown = "0.3"
+parking_lot = "0.8"
 cita-vm = "0.1"
 log = "0.4"
 futures = "0.1"
 ethereum-types = "0.4"
 hex = "0.3"
-cita_trie = "0.3"
+cita_trie = "0.4"
 
 [dev-dependencies]
 core-crypto = { path = "../../core/crypto" }

components/executor/src/trie_db.rs

@@ -1,6 +1,8 @@
 use std::sync::Arc;
 
 use futures::future::Future;
+use hashbrown::HashMap;
+use parking_lot::RwLock;
 
 use core_context::Context;
 use core_runtime::{DataCategory, Database, DatabaseError};
@@ -10,15 +12,19 @@ pub struct TrieDB<DB>
 where
     DB: Database,
 {
-    db: Arc<DB>,
+    db:    Arc<DB>,
+    cache: Arc<RwLock<HashMap<Vec<u8>, Vec<u8>>>>,
 }
 
 impl<DB> TrieDB<DB>
 where
     DB: Database,
 {
     pub fn new(db: Arc<DB>) -> Self {
-        TrieDB { db }
+        TrieDB {
+            db,
+            cache: Arc::new(RwLock::new(HashMap::new())),
+        }
     }
 }
 
@@ -30,31 +36,55 @@ where
     type Error = DatabaseError;
 
     fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
-        self.db.get(Context::new(), DataCategory::State, key).wait()
+        match self.cache.read().get(key) {
+            Some(v) => Ok(Some(v.to_vec())),
+            None => self.db.get(Context::new(), DataCategory::State, key).wait(),
+        }
     }
 
-    fn insert(&self, key: &[u8], value: &[u8]) -> Result<(), Self::Error> {
-        self.db
-            .insert(
-                Context::new(),
-                DataCategory::State,
-                key.to_vec(),
-                value.to_vec(),
-            )
-            .wait()
+    fn insert(&self, key: Vec<u8>, value: Vec<u8>) -> Result<(), Self::Error> {
+        self.cache.write().insert(key, value);
+        Ok(())
     }
 
     fn contains(&self, key: &[u8]) -> Result<bool, Self::Error> {
-        self.db
-            .contains(Context::new(), DataCategory::State, key)
-            .wait()
+        if self.cache.read().contains_key(key) {
+            Ok(true)
+        } else {
+            self.db
+                .contains(Context::new(), DataCategory::State, key)
+                .wait()
+        }
     }
 
     fn remove(&self, _key: &[u8]) -> Result<(), Self::Error> {
         Ok(())
     }
 
-    fn insert_batch(&self, keys: &[Vec<u8>], values: &[Vec<u8>]) -> Result<(), Self::Error> {
+    fn insert_batch(&self, keys: Vec<Vec<u8>>, values: Vec<Vec<u8>>) -> Result<(), Self::Error> {
+        let mut cache = self.cache.write();
+        for i in 0..keys.len() {
+            let key = keys[i].clone();
+            let value = values[i].clone();
+            cache.insert(key, value);
+        }
+        Ok(())
+    }
+
+    fn remove_batch(&self, _keys: &[Vec<u8>]) -> Result<(), Self::Error> {
+        Ok(())
+    }
+
+    fn flush(&self) -> Result<(), Self::Error> {
+        let len = self.cache.read().len();
+        let mut keys = Vec::with_capacity(len);
+        let mut values = Vec::with_capacity(len);
+
+        for (key, value) in self.cache.write().drain() {
+            keys.push(key);
+            values.push(value);
+        }
+
         self.db
             .insert_batch(
                 Context::new(),
@@ -64,10 +94,6 @@ where
             )
             .wait()
     }
-
-    fn remove_batch(&self, _keys: &[Vec<u8>]) -> Result<(), Self::Error> {
-        Ok(())
-    }
 }
 
 impl<DB> Clone for TrieDB<DB>
@@ -76,7 +102,8 @@ where
 {
     fn clone(&self) -> Self {
         TrieDB {
-            db: Arc::clone(&self.db),
+            db:    Arc::clone(&self.db),
+            cache: Arc::clone(&self.cache),
         }
     }
 }