Make rocksdb behind --use-kvdb

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<H: Hasher>(KvdbBackend<H>)
 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<H: Hasher> Default for TrieStorage<H>
 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<H>) = 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<H>) = Deserialize::deserialize(deserializer)?;
+ (root, DatabaseAdapter::Rocks(db))
+ } else {
+ let (root, kvs): (H::Out, Vec<(Vec<u8>, 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<H: Hasher>(trie: &KvdbBackend<H>) -> KvdbBackend<H>

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<u8>, Vec<u8>)>,
 }
+
+pub struct SerAsSeq<'a, H: Hasher>(pub &'a crate::MemoryDB<H>)
+where
+ H::Out: Ord;
+
+impl<H: Hasher> Serialize for SerAsSeq<'_, H>
+where
+ H::Out: Ord,
+{
+ fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+ where
+ S: serde::Serializer,
+ {
+ self.0.serialize(serializer)
+ }
+}

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<u8>, (Vec<u8>, i32)>),
+ Memory(OrdMap<Vec<u8>, (i32, Vec<u8>)>),
 }
 
 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);

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) {

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"

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;