implement CacheKVS and wrap TxStore with it

Signed-off-by: Jun Kimura <jun.kimura@datachain.jp>

Author: bluele

enclave-modules/environment/src/environment_impl.rs

@@ -1,6 +1,6 @@
 use crate::{prelude::*, Env};
 use alloc::sync::Arc;
-use host_api::store::EnclaveStore;
+use host_api::store::new_enclave_store;
 use light_client::{LightClient, LightClientResolver, MapLightClientRegistry};
 use store::{KVStore, TxId};
 
@@ -24,7 +24,7 @@ impl LightClientResolver for Environment {
 
 impl Env for Environment {
     fn new_store(&self, tx_id: TxId) -> Box<dyn KVStore> {
-        Box::new(EnclaveStore::new(tx_id))
+        new_enclave_store(tx_id)
     }
 
     fn get_lc_registry(&self) -> Arc<dyn LightClientResolver> {

enclave-modules/host-api/src/store.rs

@@ -1,21 +1,27 @@
 use crate::prelude::*;
 use crate::{api::execute_command, Error};
 use ocall_commands::{Command, CommandResult, StoreCommand, StoreResult};
+use store::cache::CacheKVS;
 use store::{KVStore, TxId};
 
-/// EnclaveStore is a KVStore implementation that uses the ocall_commands to interact with the
+/// The store guarantees that reads a value from the host store only once per key
+pub fn new_enclave_store(tx_id: TxId) -> Box<dyn KVStore> {
+    Box::new(CacheKVS::new(TxStore::new(tx_id)))
+}
+
+/// TxStore is a KVStore implementation that uses the ocall_commands to interact with the
 /// host's store.
-pub struct EnclaveStore {
+struct TxStore {
     tx_id: TxId,
 }
 
-impl EnclaveStore {
+impl TxStore {
     pub fn new(tx_id: TxId) -> Self {
         Self { tx_id }
     }
 }
 
-impl KVStore for EnclaveStore {
+impl KVStore for TxStore {
     fn get(&self, key: &[u8]) -> Option<Vec<u8>> {
         get(self.tx_id, key.to_vec()).unwrap()
     }

modules/store/src/cache.rs

@@ -0,0 +1,122 @@
+use crate::prelude::*;
+use crate::KVStore;
+use alloc::collections::BTreeMap;
+use core::cell::RefCell;
+
+/// A key-value store that caches reads in memory
+pub struct CacheKVS<S: KVStore> {
+    parent: S,
+    cache: RefCell<BTreeMap<Vec<u8>, Option<Vec<u8>>>>,
+}
+
+impl<S: KVStore> CacheKVS<S> {
+    pub fn new(parent: S) -> Self {
+        Self {
+            parent,
+            cache: RefCell::new(BTreeMap::new()),
+        }
+    }
+}
+
+impl<S: KVStore> KVStore for CacheKVS<S> {
+    fn set(&mut self, key: Vec<u8>, value: Vec<u8>) {
+        self.parent.set(key.clone(), value.clone());
+        self.cache.borrow_mut().insert(key, Some(value));
+    }
+
+    fn get(&self, key: &[u8]) -> Option<Vec<u8>> {
+        let cache = self.cache.borrow();
+        let res = cache.get(key);
+        match res {
+            Some(Some(v)) => Some(v.clone()),
+            Some(None) => None,
+            None => {
+                drop(cache);
+                let v = self.parent.get(key);
+                self.cache.borrow_mut().insert(key.to_vec(), v.clone());
+                v
+            }
+        }
+    }
+
+    fn remove(&mut self, key: &[u8]) {
+        self.parent.remove(key);
+        self.cache.borrow_mut().insert(key.to_vec(), None);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::KVStore;
+    use alloc::rc::Rc;
+
+    pub struct MockStore {
+        db: BTreeMap<Vec<u8>, Vec<u8>>,
+    }
+
+    impl MockStore {
+        pub fn new() -> Self {
+            Self {
+                db: BTreeMap::new(),
+            }
+        }
+    }
+
+    impl KVStore for MockStore {
+        fn set(&mut self, key: Vec<u8>, value: Vec<u8>) {
+            self.db.insert(key, value);
+        }
+
+        fn get(&self, key: &[u8]) -> Option<Vec<u8>> {
+            self.db.get(key).cloned()
+        }
+
+        fn remove(&mut self, key: &[u8]) {
+            self.db.remove(key);
+        }
+    }
+
+    #[allow(non_snake_case)]
+    fn B(s: &str) -> Vec<u8> {
+        s.as_bytes().to_vec()
+    }
+
+    #[test]
+    fn test_cache_kvs() {
+        let mut mock = Rc::new(RefCell::new(MockStore::new()));
+        mock.set(B("k1"), B("v1"));
+
+        let mut cache = CacheKVS::new(mock.clone());
+        assert_eq!(cache.get(&B("k1")), Some(B("v1")));
+
+        cache.set(B("k1"), B("v2"));
+        assert_eq!(cache.get(&B("k1")), Some(B("v2")));
+        assert_eq!(mock.get(&B("k1")), Some(B("v2")));
+
+        mock.set(B("k1"), B("v3"));
+        assert_eq!(mock.get(&B("k1")), Some(B("v3")));
+        assert_eq!(cache.get(&B("k1")), Some(B("v2")));
+
+        mock.remove(&B("k1"));
+        assert_eq!(mock.get(&B("k1")), None);
+        assert_eq!(cache.get(&B("k1")), Some(B("v2")));
+
+        mock.set(B("k2"), B("v4"));
+        assert_eq!(cache.get(&B("k2")), Some(B("v4")));
+
+        assert_eq!(cache.get(&B("k3")), None);
+        mock.set(B("k3"), B("v5"));
+        assert_eq!(mock.get(&B("k3")), Some(B("v5")));
+        assert_eq!(cache.get(&B("k3")), None);
+        cache.set(B("k3"), B("v6"));
+        assert_eq!(cache.get(&B("k3")), Some(B("v6")));
+
+        cache.remove(&B("k4"));
+        mock.set(B("k4"), B("v7"));
+        assert_eq!(cache.get(&B("k4")), None);
+        assert_eq!(mock.get(&B("k4")), Some(B("v7")));
+        cache.set(B("k4"), B("v8"));
+        assert_eq!(cache.get(&B("k4")), Some(B("v8")));
+    }
+}

modules/store/src/lib.rs

@@ -22,6 +22,7 @@ mod prelude {
 pub use crate::errors::{Error, Result};
 pub use crate::store::{KVStore, TxId};
 
+pub mod cache;
 mod errors;
 #[cfg(feature = "std")]
 pub mod host;