From 25c46c9b170268d95c517ea43b9bf114cbccdeed Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Wed, 31 May 2023 08:11:59 +0000 Subject: [PATCH] Make rocksdb behind --use-kvdb --- crates/phala-trie-storage/src/lib.rs | 48 +++++++++++++++++++++------- crates/phala-trie-storage/src/ser.rs | 17 ++++++++++ crates/pink/runner/src/storage.rs | 38 +++++++++++++++++----- e2e/src/fullstack.js | 1 + standalone/pruntime/Cargo.lock | 2 ++ standalone/pruntime/Cargo.toml | 2 ++ standalone/pruntime/src/main.rs | 9 +++++- 7 files changed, 97 insertions(+), 20 deletions(-) diff --git a/crates/phala-trie-storage/src/lib.rs b/crates/phala-trie-storage/src/lib.rs index cff55ab3ce..adff5bc914 100644 --- a/crates/phala-trie-storage/src/lib.rs +++ b/crates/phala-trie-storage/src/lib.rs @@ -12,6 +12,7 @@ use fused::DatabaseAdapter; #[cfg(feature = "serde")] use serde::{Deserialize, Deserializer, Serialize, Serializer}; use std::iter::FromIterator; +use std::sync::atomic::{AtomicBool, Ordering}; use parity_scale_codec::Codec; use sp_core::storage::ChildInfo; @@ -41,14 +42,27 @@ pub struct TrieStorage(KvdbBackend) where H::Out: Ord; +static USE_ROCKSDB: AtomicBool = AtomicBool::new(true); + +fn use_rocksdb() -> bool { + USE_ROCKSDB.load(Ordering::Relaxed) +} + +pub fn set_use_rocksdb(use_rocksdb: bool) { + USE_ROCKSDB.store(use_rocksdb, Ordering::Relaxed); +} + impl Default for TrieStorage where H::Out: Codec + Ord, { fn default() -> Self { - Self( - TrieBackendBuilder::new(DatabaseAdapter::default_rocksdb(), Default::default()).build(), - ) + let backend = if use_rocksdb() { + DatabaseAdapter::default_rocksdb() + } else { + DatabaseAdapter::default_memdb() + }; + Self(TrieBackendBuilder::new(backend, Default::default()).build()) } } @@ -81,7 +95,12 @@ where } } } - TrieBackendBuilder::new(DatabaseAdapter::Rocks(RocksHashDB::load(mdb)), root).build() + let storage = if use_rocksdb() { + DatabaseAdapter::Rocks(RocksHashDB::load(mdb)) + } else { + DatabaseAdapter::Memory(mdb) + }; + TrieBackendBuilder::new(storage, root).build() } #[cfg(feature = "serde")] @@ -97,9 +116,7 @@ where let db = trie.backend_storage(); match db { DatabaseAdapter::Rocks(db) => (root, db).serialize(serializer), - DatabaseAdapter::Memory(_) => Err(serde::ser::Error::custom( - "Cannot serialize memory DB. It's only for unit test only", - )), + DatabaseAdapter::Memory(mdb) => (root, ser::SerAsSeq(mdb)).serialize(serializer), } } @@ -111,10 +128,19 @@ where H::Out: Codec + Deserialize<'de> + Ord, De: Deserializer<'de>, { - let (root, db): (H::Out, RocksHashDB) = Deserialize::deserialize(deserializer)?; - let db = DatabaseAdapter::Rocks(db); - let backend = TrieBackendBuilder::new(db, root).build(); - Ok(backend) + let (root, db) = if use_rocksdb() { + let (root, db): (H::Out, RocksHashDB) = Deserialize::deserialize(deserializer)?; + (root, DatabaseAdapter::Rocks(db)) + } else { + let (root, kvs): (H::Out, Vec<(Vec, i32)>) = Deserialize::deserialize(deserializer)?; + let mdb = MemoryDB::from_inner( + kvs.into_iter() + .map(|(data, rc)| (H::hash(data.as_ref()), (data, rc))) + .collect(), + ); + (root, DatabaseAdapter::Memory(mdb)) + }; + Ok(TrieBackendBuilder::new(db, root).build()) } pub fn clone_trie_backend(trie: &KvdbBackend) -> KvdbBackend diff --git a/crates/phala-trie-storage/src/ser.rs b/crates/phala-trie-storage/src/ser.rs index ba72a29165..75067562d2 100644 --- a/crates/phala-trie-storage/src/ser.rs +++ b/crates/phala-trie-storage/src/ser.rs @@ -1,3 +1,4 @@ +use hash_db::Hasher; use parity_scale_codec::{Decode, Encode}; use scale_info::TypeInfo; use serde::{Deserialize, Serialize}; @@ -16,3 +17,19 @@ pub struct StorageChanges { pub struct StorageData { pub inner: Vec<(Vec, Vec)>, } + +pub struct SerAsSeq<'a, H: Hasher>(pub &'a crate::MemoryDB) +where + H::Out: Ord; + +impl Serialize for SerAsSeq<'_, H> +where + H::Out: Ord, +{ + fn serialize(&self, serializer: S) -> Result + where + S: serde::Serializer, + { + self.0.serialize(serializer) + } +} diff --git a/crates/pink/runner/src/storage.rs b/crates/pink/runner/src/storage.rs index 478d020870..dbd017d3ef 100644 --- a/crates/pink/runner/src/storage.rs +++ b/crates/pink/runner/src/storage.rs @@ -2,15 +2,30 @@ use im::OrdMap; use phala_trie_storage::RocksDB; use pink_capi::{types::Hash, v1::ocall::StorageChanges}; use serde::{Deserialize, Deserializer, Serialize, Serializer}; +use std::sync::atomic::{AtomicBool, Ordering}; + +static USE_ROCKSDB: AtomicBool = AtomicBool::new(true); + +fn use_rocksdb() -> bool { + USE_ROCKSDB.load(Ordering::Relaxed) +} + +pub fn set_use_rocksdb(use_rocksdb: bool) { + USE_ROCKSDB.store(use_rocksdb, Ordering::Relaxed); +} enum StorageAdapter { RocksDB(RocksDB), - Memory(OrdMap, (Vec, i32)>), + Memory(OrdMap, (i32, Vec)>), } impl Default for StorageAdapter { fn default() -> Self { - StorageAdapter::RocksDB(RocksDB::default()) + if use_rocksdb() { + StorageAdapter::RocksDB(Default::default()) + } else { + StorageAdapter::Memory(Default::default()) + } } } @@ -21,7 +36,7 @@ impl Serialize for StorageAdapter { { match self { StorageAdapter::RocksDB(db) => db.serialize(serializer), - StorageAdapter::Memory(_) => unimplemented!("InMemory storage is for testing only"), + StorageAdapter::Memory(mdb) => mdb.serialize(serializer), } } } @@ -31,7 +46,11 @@ impl<'de> Deserialize<'de> for StorageAdapter { where D: Deserializer<'de>, { - RocksDB::deserialize(deserializer).map(StorageAdapter::RocksDB) + if use_rocksdb() { + Deserialize::deserialize(deserializer).map(StorageAdapter::RocksDB) + } else { + Deserialize::deserialize(deserializer).map(StorageAdapter::Memory) + } } } @@ -50,7 +69,10 @@ impl StorageAdapter { StorageAdapter::RocksDB(kvdb) => { kvdb.get_r(key).expect("Failed to get key from RocksDB") } - StorageAdapter::Memory(mdb) => mdb.get(key).cloned(), + StorageAdapter::Memory(mdb) => { + let (rc, v) = mdb.get(key).cloned()?; + Some((v, rc)) + }, } } @@ -68,8 +90,8 @@ impl StorageAdapter { let pv = mdb.get(key).cloned(); let raw_value = match pv { - None => (value, rc), - Some((mut d, mut orc)) => { + None => (rc, value), + Some((mut orc, mut d)) => { if orc <= 0 { d = value; } @@ -80,7 +102,7 @@ impl StorageAdapter { mdb.remove(key); continue; } - (d, orc) + (orc, d) } }; mdb.insert(key.to_vec(), raw_value); diff --git a/e2e/src/fullstack.js b/e2e/src/fullstack.js index b1ad8ce345..ac81581972 100644 --- a/e2e/src/fullstack.js +++ b/e2e/src/fullstack.js @@ -1321,6 +1321,7 @@ function newPRuntime(teePort, tmpPath, name = 'app') { '--cores=0', // Disable benchmark '--checkpoint-interval=5', '--port', teePort.toString(), + '--use-kvdb', ]; let bin = pRuntimeBin; if (inSgx) { diff --git a/standalone/pruntime/Cargo.lock b/standalone/pruntime/Cargo.lock index 4aed287667..64b3e33e22 100644 --- a/standalone/pruntime/Cargo.lock +++ b/standalone/pruntime/Cargo.lock @@ -5197,7 +5197,9 @@ dependencies = [ "phala-git-revision", "phala-rocket-middleware", "phala-sanitized-logger", + "phala-trie-storage", "phala-types", + "pink-runner", "reqwest", "reqwest-env-proxy", "rocket", diff --git a/standalone/pruntime/Cargo.toml b/standalone/pruntime/Cargo.toml index a3580a537b..edb6cf53c2 100644 --- a/standalone/pruntime/Cargo.toml +++ b/standalone/pruntime/Cargo.toml @@ -38,6 +38,8 @@ phala-allocator = { path = "../../crates/phala-allocator" } phala-rocket-middleware = { path = "../../crates/phala-rocket-middleware" } phala-types = { path = "../../crates/phala-types", features = ["enable_serde", "sgx"] } phala-git-revision = { path = "../../crates/phala-git-revision" } +phala-trie-storage = { path = "../../crates/phala-trie-storage", default-features = false } +pink-runner = { path = "../../crates/pink/runner", default-features = false } sgx-api-lite = { path = "../../crates/sgx-api-lite" } tracing = "0.1" hex_fmt = "0.3.0" diff --git a/standalone/pruntime/src/main.rs b/standalone/pruntime/src/main.rs index c46e8d92df..a1bed97660 100644 --- a/standalone/pruntime/src/main.rs +++ b/standalone/pruntime/src/main.rs @@ -81,6 +81,10 @@ struct Args { /// Disable the RCU policy to update the Phactory state. #[arg(long)] no_rcu: bool, + + /// Put the chain and contract state to disk. + #[arg(long)] + use_kvdb: bool, } #[rocket::main] @@ -96,7 +100,10 @@ async fn main() -> Result<(), rocket::Error> { async fn serve(sgx: bool) -> Result<(), rocket::Error> { let args = Args::parse(); - info!(sgx, "Starting pruntime..."); + info!(sgx, ?args, "Starting pruntime..."); + + phala_trie_storage::set_use_rocksdb(args.use_kvdb); + pink_runner::storage::set_use_rocksdb(args.use_kvdb); let sealing_path; let storage_path;