From 61d8e82791c6374fe776149ee05ceb3deb2a6e2b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Marcos=20Bezerra?= Date: Mon, 3 Jun 2024 12:32:41 -0300 Subject: [PATCH] readd multi_get tests --- src/eth/storage/rocks/rocks_db.rs | 4 +- src/eth/storage/rocks/rocks_permanent.rs | 2 +- src/eth/storage/rocks/rocks_state.rs | 54 +++++++++++++++++++++--- 3 files changed, 52 insertions(+), 8 deletions(-) diff --git a/src/eth/storage/rocks/rocks_db.rs b/src/eth/storage/rocks/rocks_db.rs index 5f8bfe937..ee520096f 100644 --- a/src/eth/storage/rocks/rocks_db.rs +++ b/src/eth/storage/rocks/rocks_db.rs @@ -15,7 +15,9 @@ use crate::eth::storage::rocks::rocks_config::DbConfig; /// Create or open the Database with the configs applied to all column families /// /// The returned `Options` need to be stored to refer to the DB metrics -pub fn create_or_open_db(path: &Path, cf_configs: &HashMap<&'static str, Options>) -> anyhow::Result<(Arc, Options)> { +pub fn create_or_open_db(path: impl AsRef, cf_configs: &HashMap<&'static str, Options>) -> anyhow::Result<(Arc, Options)> { + let path = path.as_ref(); + // settings for each Column Family to be created let cf_config_iter = cf_configs.iter().map(|(name, opts)| (*name, opts.clone())); diff --git a/src/eth/storage/rocks/rocks_permanent.rs b/src/eth/storage/rocks/rocks_permanent.rs index 94fec3819..33a782126 100644 --- a/src/eth/storage/rocks/rocks_permanent.rs +++ b/src/eth/storage/rocks/rocks_permanent.rs @@ -27,7 +27,7 @@ use crate::eth::storage::PermanentStorage; #[derive(Debug)] pub struct RocksPermanentStorage { - state: RocksStorageState, + pub state: RocksStorageState, block_number: AtomicU64, } diff --git a/src/eth/storage/rocks/rocks_state.rs b/src/eth/storage/rocks/rocks_state.rs index be04fab9a..9dde72784 100644 --- a/src/eth/storage/rocks/rocks_state.rs +++ b/src/eth/storage/rocks/rocks_state.rs @@ -113,7 +113,7 @@ impl RocksStorageState { let db_path = path.as_ref().to_path_buf(); let (backup_trigger_tx, backup_trigger_rx) = mpsc::channel::<()>(1); - let (db, db_options) = create_or_open_db(&db_path, &*CF_OPTIONS_MAP).unwrap(); + let (db, db_options) = create_or_open_db(&db_path, &CF_OPTIONS_MAP).unwrap(); //XXX TODO while repair/restore from backup, make sure to sync online and only when its in sync with other nodes, receive requests let state = Self { @@ -512,8 +512,7 @@ impl RocksStorageState { } /// Writes slots to state (does not write to slot history) - #[allow(dead_code)] - fn write_slots(&self, slots: Vec<(Address, Slot)>) { + pub fn write_slots(&self, slots: Vec<(Address, Slot)>) { let slots = slots .into_iter() .map(|(address, slot)| ((address.into(), slot.index.into()), slot.value.into())); @@ -603,8 +602,51 @@ where K: Serialize + for<'de> Deserialize<'de> + std::hash::Hash + Eq, V: Serialize + for<'de> Deserialize<'de> + Clone, { - let options = CF_OPTIONS_MAP - .get(&column_family) - .unwrap_or_else(|| panic!("column_family `{column_family}` given to `new_cf` not found in options map")); + let Some(options) = CF_OPTIONS_MAP.get(column_family) else { + panic!("column_family `{column_family}` given to `new_cf` not found in options map"); + }; RocksCf::new_cf(Arc::clone(db), column_family, options.clone()) } + +#[cfg(test)] +mod tests { + use std::collections::HashSet; + use std::fs; + + use fake::Fake; + use fake::Faker; + + use super::*; + + #[test] + fn test_rocks_multi_get() { + let (db, _db_options) = create_or_open_db("./data/slots_test.rocksdb", &CF_OPTIONS_MAP).unwrap(); + let account_slots: RocksCf = new_cf(&db, "account_slots"); + + let slots: HashMap = (0..1000).map(|_| (Faker.fake(), Faker.fake())).collect(); + let extra_slots: HashMap = (0..1000) + .map(|_| (Faker.fake(), Faker.fake())) + .filter(|(key, _)| !slots.contains_key(key)) + .collect(); + + let mut batch = WriteBatch::default(); + account_slots.prepare_batch_insertion(slots.clone(), &mut batch); + account_slots.prepare_batch_insertion(extra_slots.clone(), &mut batch); + db.write(batch).unwrap(); + + let extra_keys: HashSet = (0..1000) + .map(|_| Faker.fake()) + .filter(|key| !extra_slots.contains_key(key) && !slots.contains_key(key)) + .collect(); + + let keys: Vec = slots.keys().cloned().chain(extra_keys).collect(); + let result = account_slots.multi_get(keys).expect("this should not fail"); + + assert_eq!(result.len(), slots.keys().len()); + for (idx, value) in result { + assert_eq!(value, *slots.get(&idx).expect("should not be None")); + } + + fs::remove_dir_all("./data/slots_test.rocksdb").unwrap(); + } +}