From 62504b2622cbf2cb97f705d2cedd9a6b0cf43b1c Mon Sep 17 00:00:00 2001 From: Luke Parker Date: Thu, 13 Jul 2023 19:09:11 -0400 Subject: [PATCH] Bind RocksDB into serai-db --- Cargo.lock | 5 +++ common/db/Cargo.toml | 6 +++ common/db/src/lib.rs | 91 ++++-------------------------------------- common/db/src/mem.rs | 80 +++++++++++++++++++++++++++++++++++++ common/db/src/rocks.rs | 34 ++++++++++++++++ 5 files changed, 133 insertions(+), 83 deletions(-) create mode 100644 common/db/src/mem.rs create mode 100644 common/db/src/rocks.rs diff --git a/Cargo.lock b/Cargo.lock index 974066ee..1eaabcf8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4668,7 +4668,9 @@ dependencies = [ "glob", "libc", "libz-sys", + "lz4-sys", "tikv-jemalloc-sys", + "zstd-sys", ] [[package]] @@ -8633,6 +8635,9 @@ dependencies = [ [[package]] name = "serai-db" version = "0.1.0" +dependencies = [ + "rocksdb", +] [[package]] name = "serai-in-instructions-pallet" diff --git a/common/db/Cargo.toml b/common/db/Cargo.toml index 4dba6679..9606bda8 100644 --- a/common/db/Cargo.toml +++ b/common/db/Cargo.toml @@ -11,3 +11,9 @@ edition = "2021" [package.metadata.docs.rs] all-features = true rustdoc-args = ["--cfg", "docsrs"] + +[dependencies] +rocksdb = { version = "0.21", optional = true } + +[features] +rocksdb = ["dep:rocksdb"] diff --git a/common/db/src/lib.rs b/common/db/src/lib.rs index 12f4977c..5baf35f1 100644 --- a/common/db/src/lib.rs +++ b/common/db/src/lib.rs @@ -1,24 +1,24 @@ -use core::fmt::Debug; -use std::{ - sync::{Arc, RwLock}, - collections::{HashSet, HashMap}, -}; +mod mem; +pub use mem::*; + +#[cfg(feature = "rocksdb")] +mod rocks; /// An object implementing get. -pub trait Get: Send + Sync + Debug { +pub trait Get { fn get(&self, key: impl AsRef<[u8]>) -> Option>; } /// An atomic database operation. #[must_use] -pub trait DbTxn: Send + Sync + Debug + Get { +pub trait DbTxn: Send + Get { fn put(&mut self, key: impl AsRef<[u8]>, value: impl AsRef<[u8]>); fn del(&mut self, key: impl AsRef<[u8]>); fn commit(self); } /// A database supporting atomic operations. -pub trait Db: 'static + Send + Sync + Clone + Debug + Get { +pub trait Db: 'static + Send + Sync + Clone + Get { type Transaction<'a>: DbTxn; fn key(db_dst: &'static [u8], item_dst: &'static [u8], key: impl AsRef<[u8]>) -> Vec { let db_len = u8::try_from(db_dst.len()).unwrap(); @@ -27,78 +27,3 @@ pub trait Db: 'static + Send + Sync + Clone + Debug + Get { } fn txn(&mut self) -> Self::Transaction<'_>; } - -/// An atomic operation for the in-memory databae. -#[must_use] -#[derive(PartialEq, Eq, Debug)] -pub struct MemDbTxn<'a>(&'a MemDb, HashMap, Vec>, HashSet>); - -impl<'a> Get for MemDbTxn<'a> { - fn get(&self, key: impl AsRef<[u8]>) -> Option> { - if self.2.contains(key.as_ref()) { - return None; - } - self - .1 - .get(key.as_ref()) - .cloned() - .or_else(|| self.0 .0.read().unwrap().get(key.as_ref()).cloned()) - } -} -impl<'a> DbTxn for MemDbTxn<'a> { - fn put(&mut self, key: impl AsRef<[u8]>, value: impl AsRef<[u8]>) { - self.2.remove(key.as_ref()); - self.1.insert(key.as_ref().to_vec(), value.as_ref().to_vec()); - } - fn del(&mut self, key: impl AsRef<[u8]>) { - self.1.remove(key.as_ref()); - self.2.insert(key.as_ref().to_vec()); - } - fn commit(mut self) { - let mut db = self.0 .0.write().unwrap(); - for (key, value) in self.1.drain() { - db.insert(key, value); - } - for key in self.2 { - db.remove(&key); - } - } -} - -/// An in-memory database. -#[derive(Clone, Debug)] -pub struct MemDb(Arc, Vec>>>); - -impl PartialEq for MemDb { - fn eq(&self, other: &MemDb) -> bool { - *self.0.read().unwrap() == *other.0.read().unwrap() - } -} -impl Eq for MemDb {} - -impl Default for MemDb { - fn default() -> MemDb { - MemDb(Arc::new(RwLock::new(HashMap::new()))) - } -} - -impl MemDb { - /// Create a new in-memory database. - pub fn new() -> MemDb { - MemDb::default() - } -} - -impl Get for MemDb { - fn get(&self, key: impl AsRef<[u8]>) -> Option> { - self.0.read().unwrap().get(key.as_ref()).cloned() - } -} -impl Db for MemDb { - type Transaction<'a> = MemDbTxn<'a>; - fn txn(&mut self) -> MemDbTxn<'_> { - MemDbTxn(self, HashMap::new(), HashSet::new()) - } -} - -// TODO: Also bind RocksDB diff --git a/common/db/src/mem.rs b/common/db/src/mem.rs new file mode 100644 index 00000000..f5240869 --- /dev/null +++ b/common/db/src/mem.rs @@ -0,0 +1,80 @@ +use core::fmt::Debug; +use std::{ + sync::{Arc, RwLock}, + collections::{HashSet, HashMap}, +}; + +use crate::*; + +/// An atomic operation for the in-memory databae. +#[must_use] +#[derive(PartialEq, Eq, Debug)] +pub struct MemDbTxn<'a>(&'a MemDb, HashMap, Vec>, HashSet>); + +impl<'a> Get for MemDbTxn<'a> { + fn get(&self, key: impl AsRef<[u8]>) -> Option> { + if self.2.contains(key.as_ref()) { + return None; + } + self + .1 + .get(key.as_ref()) + .cloned() + .or_else(|| self.0 .0.read().unwrap().get(key.as_ref()).cloned()) + } +} +impl<'a> DbTxn for MemDbTxn<'a> { + fn put(&mut self, key: impl AsRef<[u8]>, value: impl AsRef<[u8]>) { + self.2.remove(key.as_ref()); + self.1.insert(key.as_ref().to_vec(), value.as_ref().to_vec()); + } + fn del(&mut self, key: impl AsRef<[u8]>) { + self.1.remove(key.as_ref()); + self.2.insert(key.as_ref().to_vec()); + } + fn commit(mut self) { + let mut db = self.0 .0.write().unwrap(); + for (key, value) in self.1.drain() { + db.insert(key, value); + } + for key in self.2 { + db.remove(&key); + } + } +} + +/// An in-memory database. +#[derive(Clone, Debug)] +pub struct MemDb(Arc, Vec>>>); + +impl PartialEq for MemDb { + fn eq(&self, other: &MemDb) -> bool { + *self.0.read().unwrap() == *other.0.read().unwrap() + } +} +impl Eq for MemDb {} + +impl Default for MemDb { + fn default() -> MemDb { + MemDb(Arc::new(RwLock::new(HashMap::new()))) + } +} + +impl MemDb { + /// Create a new in-memory database. + pub fn new() -> MemDb { + MemDb::default() + } +} + +impl Get for MemDb { + fn get(&self, key: impl AsRef<[u8]>) -> Option> { + self.0.read().unwrap().get(key.as_ref()).cloned() + } +} +impl Db for MemDb { + type Transaction<'a> = MemDbTxn<'a>; + fn txn(&mut self) -> MemDbTxn<'_> { + MemDbTxn(self, HashMap::new(), HashSet::new()) + } +} diff --git a/common/db/src/rocks.rs b/common/db/src/rocks.rs new file mode 100644 index 00000000..0866bb7e --- /dev/null +++ b/common/db/src/rocks.rs @@ -0,0 +1,34 @@ +use std::sync::Arc; + +use rocksdb::{ThreadMode, Transaction, TransactionDB}; + +use crate::*; + +impl Get for Transaction<'_, TransactionDB> { + fn get(&self, key: impl AsRef<[u8]>) -> Option> { + self.get(key).expect("couldn't read from RocksDB via transaction") + } +} +impl DbTxn for Transaction<'_, TransactionDB> { + fn put(&mut self, key: impl AsRef<[u8]>, value: impl AsRef<[u8]>) { + Transaction::put(self, key, value).expect("couldn't write to RocksDB via transaction") + } + fn del(&mut self, key: impl AsRef<[u8]>) { + self.delete(key).expect("couldn't delete from RocksDB via transaction") + } + fn commit(self) { + Transaction::commit(self).expect("couldn't commit to RocksDB via transaction") + } +} + +impl Get for Arc> { + fn get(&self, key: impl AsRef<[u8]>) -> Option> { + TransactionDB::get(self, key).expect("couldn't read from RocksDB") + } +} +impl Db for Arc> { + type Transaction<'a> = Transaction<'a, TransactionDB>; + fn txn(&mut self) -> Self::Transaction<'_> { + self.transaction() + } +}