Merge pull request #100 from datachainlab/fix-enclave-store

Implement `CacheKVS` and utilize it as enclave store

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

Author: bluele

enclave-modules/environment/Cargo.toml

@@ -6,8 +6,8 @@ edition = "2021"
 [dependencies]
 light-client = { path = "../../modules/light-client", default-features = false }
 store = { path = "../../modules/store", default-features = false }
-enclave-store = { path = "../store", optional = true }
+host-api = { path = "../host-api", optional = true }
 
 [features]
 default = ["environment_impl"]
-environment_impl = ["enclave-store"]
+environment_impl = ["host-api"]

enclave-modules/environment/src/environment_impl.rs

@@ -1,6 +1,6 @@
 use crate::{prelude::*, Env};
 use alloc::sync::Arc;
-use enclave_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,9 +1,41 @@
 use crate::prelude::*;
 use crate::{api::execute_command, Error};
 use ocall_commands::{Command, CommandResult, StoreCommand, StoreResult};
-use store::TxId;
+use store::cache::CacheKVS;
+use store::{KVStore, TxId};
 
-pub fn get(tx_id: TxId, key: Vec<u8>) -> Result<Option<Vec<u8>>, Error> {
+/// 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.
+struct TxStore {
+    tx_id: TxId,
+}
+
+impl TxStore {
+    pub fn new(tx_id: TxId) -> Self {
+        Self { tx_id }
+    }
+}
+
+impl KVStore for TxStore {
+    fn get(&self, key: &[u8]) -> Option<Vec<u8>> {
+        get(self.tx_id, key.to_vec()).unwrap()
+    }
+
+    fn set(&mut self, key: Vec<u8>, value: Vec<u8>) {
+        set(self.tx_id, key.clone(), value.clone()).unwrap();
+    }
+
+    fn remove(&mut self, key: &[u8]) {
+        remove(self.tx_id, key.to_vec()).unwrap();
+    }
+}
+
+fn get(tx_id: TxId, key: Vec<u8>) -> Result<Option<Vec<u8>>, Error> {
     let cmd = Command::Store(StoreCommand::Get(tx_id, key));
     if let CommandResult::Store(StoreResult::Get(v)) = execute_command(cmd)? {
         Ok(v)
@@ -12,7 +44,7 @@ pub fn get(tx_id: TxId, key: Vec<u8>) -> Result<Option<Vec<u8>>, Error> {
     }
 }
 
-pub fn set(tx_id: TxId, key: Vec<u8>, value: Vec<u8>) -> Result<(), Error> {
+fn set(tx_id: TxId, key: Vec<u8>, value: Vec<u8>) -> Result<(), Error> {
     let cmd = Command::Store(StoreCommand::Set(tx_id, key, value));
     if let CommandResult::Store(StoreResult::Set) = execute_command(cmd)? {
         Ok(())
@@ -21,7 +53,7 @@ pub fn set(tx_id: TxId, key: Vec<u8>, value: Vec<u8>) -> Result<(), Error> {
     }
 }
 
-pub fn remove(tx_id: TxId, key: Vec<u8>) -> Result<(), Error> {
+fn remove(tx_id: TxId, key: Vec<u8>) -> Result<(), Error> {
     let cmd = Command::Store(StoreCommand::Remove(tx_id, key));
     if let CommandResult::Store(StoreResult::Remove) = execute_command(cmd)? {
         Ok(())

enclave-modules/store/Cargo.toml

@@ -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")));
+    }
+}

enclave-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;