Implement redb based StateStorage (#1040)

crates/core/Cargo.toml

@@ -48,7 +48,7 @@ ordered-float = "4.2"
 pav_regression = "0.4.0"
 parking_lot = "0.12"
 rand = { features = ["small_rng"], workspace = true }
-rocksdb = { default-features = false, optional = true, version = "0.22" }
+redb = { optional = true, version = "2" }
 serde = { features = ["derive", "rc"], workspace = true }
 serde_json = { workspace = true }
 serde_with = { workspace = true }
@@ -90,10 +90,9 @@ tempfile = "3.8"
 tracing = "0.1"
 
 [features]
-default = ["sqlite", "trace", "websocket"]
+default = ["redb", "trace", "websocket"]
 local-mode = []
 network-mode = []
-rocks_db = ["rocksdb"]
 sqlite = ["sqlx"]
 trace = ["tracing-subscriber"]
 trace-ot = ["opentelemetry-jaeger", "trace", "tracing-opentelemetry", "opentelemetry-otlp"]

crates/core/src/contract/executor.rs

@@ -498,8 +498,7 @@ impl<R> Executor<R> {
         let static_conf = crate::config::Config::conf();
 
         let db_path = crate::config::Config::conf().db_dir();
-        let state_store =
-            StateStore::new(Storage::new(Some(&db_path)).await?, MAX_MEM_CACHE).unwrap();
+        let state_store = StateStore::new(Storage::new(&db_path).await?, MAX_MEM_CACHE).unwrap();
 
         let contract_dir = config
             .node_data_dir

crates/core/src/contract/executor/mock_runtime.rs

@@ -17,8 +17,7 @@ impl Executor<MockRuntime> {
 
         let db_path = data_dir.join("db");
         std::fs::create_dir_all(&db_path).expect("directory created");
-        let state_store =
-            StateStore::new(Storage::new(Some(&db_path)).await?, u16::MAX as u32).unwrap();
+        let state_store = StateStore::new(Storage::new(&db_path).await?, u16::MAX as u32).unwrap();
 
         let executor = Executor::new(
             state_store,
@@ -125,8 +124,10 @@ mod test {
         const MAX_SIZE: i64 = 10 * 1024 * 1024;
         const MAX_MEM_CACHE: u32 = 10_000_000;
         let tmp_dir = tempfile::tempdir()?;
+        let state_store_path = tmp_dir.path().join("state_store");
         let contract_store = ContractStore::new(tmp_dir.path().join("executor-test"), MAX_SIZE)?;
-        let state_store = StateStore::new(Storage::new(None).await?, MAX_MEM_CACHE).unwrap();
+        let state_store =
+            StateStore::new(Storage::new(&state_store_path).await?, MAX_MEM_CACHE).unwrap();
         let mut counter = 0;
         Executor::new(
             state_store,

crates/core/src/contract/storages/mod.rs

@@ -1,15 +1,17 @@
+/// State storage implementation based on the `sqlite`
 #[cfg(feature = "sqlite")]
 pub mod sqlite;
-#[cfg(feature = "sqlite")]
+#[cfg(all(feature = "sqlite", not(feature = "redb")))]
 pub use sqlite::Pool as SqlitePool;
 
-#[cfg(feature = "sqlite")]
+#[cfg(all(feature = "sqlite", not(feature = "redb")))]
 pub type Storage = SqlitePool;
 
-#[cfg(feature = "rocks_db")]
-pub mod rocks_db;
-#[cfg(all(feature = "rocks_db", not(feature = "sqlite")))]
-use self::rocks_db::RocksDb;
+/// State storage implementation based on the [`redb`]
+#[cfg(feature = "redb")]
+pub mod redb;
+#[cfg(feature = "redb")]
+use self::redb::ReDb;
 
-#[cfg(all(feature = "rocks_db", not(feature = "sqlite")))]
-pub type Storage = RocksDb;
+#[cfg(feature = "redb")]
+pub type Storage = ReDb;

crates/core/src/contract/storages/redb.rs

@@ -0,0 +1,79 @@
+use std::path::Path;
+
+use freenet_stdlib::prelude::*;
+use redb::{Database, TableDefinition};
+
+use crate::wasm_runtime::StateStorage;
+
+const CONTRACT_PARAMS_TABLE: TableDefinition<&[u8], &[u8]> =
+    TableDefinition::new("contract_params");
+const STATE_TABLE: TableDefinition<&[u8], &[u8]> = TableDefinition::new("state");
+
+pub struct ReDb(Database);
+
+impl ReDb {
+    pub async fn new(db_path: &Path) -> Result<Self, redb::Error> {
+        tracing::info!("loading contract store from {db_path:?}");
+
+        Database::create(db_path).map(Self).map_err(Into::into)
+    }
+}
+
+impl StateStorage for ReDb {
+    type Error = redb::Error;
+
+    async fn store(&mut self, key: ContractKey, state: WrappedState) -> Result<(), Self::Error> {
+        let txn = self.0.begin_write()?;
+
+        {
+            let mut tbl = txn.open_table(STATE_TABLE)?;
+            tbl.insert(key.as_bytes(), state.as_ref())?;
+        }
+        txn.commit().map_err(Into::into)
+    }
+
+    async fn get(&self, key: &ContractKey) -> Result<Option<WrappedState>, Self::Error> {
+        let txn = self.0.begin_read()?;
+
+        let val = {
+            let tbl = txn.open_table(STATE_TABLE)?;
+            tbl.get(key.as_bytes())?
+        };
+
+        match val {
+            Some(v) => Ok(Some(WrappedState::new(v.value().to_vec()))),
+            None => Ok(None),
+        }
+    }
+
+    async fn store_params(
+        &mut self,
+        key: ContractKey,
+        params: Parameters<'static>,
+    ) -> Result<(), Self::Error> {
+        let txn = self.0.begin_write()?;
+
+        {
+            let mut tbl = txn.open_table(CONTRACT_PARAMS_TABLE)?;
+            tbl.insert(key.as_bytes(), params.as_ref())?;
+        }
+        txn.commit().map_err(Into::into)
+    }
+
+    async fn get_params<'a>(
+        &'a self,
+        key: &'a ContractKey,
+    ) -> Result<Option<Parameters<'static>>, Self::Error> {
+        let txn = self.0.begin_read()?;
+
+        let val = {
+            let tbl = txn.open_table(CONTRACT_PARAMS_TABLE)?;
+            tbl.get(key.as_bytes())?
+        };
+
+        match val {
+            Some(v) => Ok(Some(Parameters::from(v.value().to_vec()))),
+            None => Ok(None),
+        }
+    }
+}