From 18b3e20ac9407ee55defce63e94dd0b200771ff3 Mon Sep 17 00:00:00 2001 From: ljedrz Date: Sat, 15 Apr 2023 13:37:28 +0200 Subject: [PATCH 1/3] feat: use a BTreeMap instead of IndexMap in MemoryMap Signed-off-by: ljedrz --- synthesizer/Cargo.toml | 3 ++ synthesizer/src/store/helpers/memory_map.rs | 42 +++++++++++++-------- 2 files changed, 29 insertions(+), 16 deletions(-) diff --git a/synthesizer/Cargo.toml b/synthesizer/Cargo.toml index 9609d538c0..003eb2d14f 100644 --- a/synthesizer/Cargo.toml +++ b/synthesizer/Cargo.toml @@ -88,6 +88,9 @@ default-features = false [dependencies.anyhow] version = "1.0.70" +[dependencies.bincode] +version = "1" + [dependencies.blake2] version = "0.10" default-features = false diff --git a/synthesizer/src/store/helpers/memory_map.rs b/synthesizer/src/store/helpers/memory_map.rs index 4b7a4e6355..b3dae49ca6 100644 --- a/synthesizer/src/store/helpers/memory_map.rs +++ b/synthesizer/src/store/helpers/memory_map.rs @@ -19,10 +19,10 @@ use console::network::prelude::*; use indexmap::IndexMap; use core::{borrow::Borrow, hash::Hash}; -use indexmap::map; use parking_lot::{Mutex, RwLock}; use std::{ borrow::Cow, + collections::{btree_map, BTreeMap}, sync::{ atomic::{AtomicBool, Ordering}, Arc, @@ -34,7 +34,10 @@ pub struct MemoryMap< K: Copy + Clone + PartialEq + Eq + Hash + Serialize + for<'de> Deserialize<'de> + Send + Sync, V: Clone + PartialEq + Eq + Serialize + for<'de> Deserialize<'de> + Send + Sync, > { - map: Arc>>, + // The reason for using BTreeMap with binary keys is for the order of items to be the same as + // the one in the RocksDB-backed DataMap in snarkOS; if not for that, it could be any map + // with fast lookups and the keys could be typed (i.e. just `K` instead of `Vec`). + map: Arc, V>>>, batch_in_progress: Arc, atomic_batch: Arc>>>, } @@ -56,8 +59,10 @@ impl< { /// Initializes a new `MemoryMap` from the given iterator. fn from_iter>(iter: I) -> Self { + let map = iter.into_iter().map(|(k, v)| (bincode::serialize(&k).unwrap(), v)).collect(); + Self { - map: Arc::new(RwLock::new(IndexMap::from_iter(iter))), + map: Arc::new(RwLock::new(map)), batch_in_progress: Default::default(), atomic_batch: Default::default(), } @@ -84,6 +89,7 @@ impl< } // Otherwise, insert the key-value pair directly into the map. false => { + let key = bincode::serialize(&key).unwrap(); self.map.write().insert(key, value); } } @@ -104,7 +110,8 @@ impl< } // Otherwise, remove the key-value pair directly from the map. false => { - self.map.write().shift_remove(key); + let key = bincode::serialize(&key).unwrap(); + self.map.write().remove(&key); } } Ok(()) @@ -151,10 +158,11 @@ impl< // Acquire a write lock on the map. let mut locked_map = self.map.write(); // Perform all the queued operations. - for operation in operations { - match operation { - (key, Some(value)) => locked_map.insert(key, value), - (key, None) => locked_map.remove(&key), + for (key, value) in operations { + let key = bincode::serialize(&key).unwrap(); + match value { + Some(value) => locked_map.insert(key, value), + None => locked_map.remove(&key), }; } } @@ -172,9 +180,9 @@ impl< V: 'a + Clone + PartialEq + Eq + Serialize + for<'de> Deserialize<'de> + Send + Sync, > MapRead<'a, K, V> for MemoryMap { - type Iterator = core::iter::Map, fn((K, V)) -> (Cow<'a, K>, Cow<'a, V>)>; - type Keys = core::iter::Map, fn(K) -> Cow<'a, K>>; - type Values = core::iter::Map, fn(V) -> Cow<'a, V>>; + type Iterator = core::iter::Map, V>, fn((Vec, V)) -> (Cow<'a, K>, Cow<'a, V>)>; + type Keys = core::iter::Map, V>, fn(Vec) -> Cow<'a, K>>; + type Values = core::iter::Map, V>, fn(V) -> Cow<'a, V>>; /// /// Returns `true` if the given key exists in the map. @@ -184,7 +192,8 @@ impl< K: Borrow, Q: PartialEq + Eq + Hash + Serialize + ?Sized, { - Ok(self.map.read().contains_key(key)) + let key = bincode::serialize(key).unwrap(); + Ok(self.map.read().contains_key(&key)) } /// @@ -195,7 +204,8 @@ impl< K: Borrow, Q: PartialEq + Eq + Hash + Serialize + ?Sized, { - Ok(self.map.read().get(key).cloned().map(Cow::Owned)) + let key = bincode::serialize(key).unwrap(); + Ok(self.map.read().get(&key).cloned().map(Cow::Owned)) } /// @@ -219,14 +229,14 @@ impl< /// Returns an iterator visiting each key-value pair in the map. /// fn iter(&'a self) -> Self::Iterator { - self.map.read().clone().into_iter().map(|(k, v)| (Cow::Owned(k), Cow::Owned(v))) + self.map.read().clone().into_iter().map(|(k, v)| (Cow::Owned(bincode::deserialize(&k).unwrap()), Cow::Owned(v))) } /// /// Returns an iterator over each key in the map. /// fn keys(&'a self) -> Self::Keys { - self.map.read().clone().into_keys().map(Cow::Owned) + self.map.read().clone().into_keys().map(|k| Cow::Owned(bincode::deserialize(&k).unwrap())) } /// @@ -242,7 +252,7 @@ impl< V: Clone + PartialEq + Eq + Serialize + for<'de> Deserialize<'de> + Send + Sync, > Deref for MemoryMap { - type Target = Arc>>; + type Target = Arc, V>>>; fn deref(&self) -> &Self::Target { &self.map From b96d9d9e08a7bd055eb36d79294c3e690f94e0de Mon Sep 17 00:00:00 2001 From: Howard Wu <9260812+howardwu@users.noreply.github.com> Date: Sat, 15 Apr 2023 10:03:40 -0700 Subject: [PATCH 2/3] Update synthesizer/src/store/helpers/memory_map.rs Signed-off-by: Howard Wu <9260812+howardwu@users.noreply.github.com> --- synthesizer/src/store/helpers/memory_map.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/synthesizer/src/store/helpers/memory_map.rs b/synthesizer/src/store/helpers/memory_map.rs index b3dae49ca6..4a96a2965a 100644 --- a/synthesizer/src/store/helpers/memory_map.rs +++ b/synthesizer/src/store/helpers/memory_map.rs @@ -35,7 +35,7 @@ pub struct MemoryMap< V: Clone + PartialEq + Eq + Serialize + for<'de> Deserialize<'de> + Send + Sync, > { // The reason for using BTreeMap with binary keys is for the order of items to be the same as - // the one in the RocksDB-backed DataMap in snarkOS; if not for that, it could be any map + // the one in the RocksDB-backed DataMap; if not for that, it could be any map // with fast lookups and the keys could be typed (i.e. just `K` instead of `Vec`). map: Arc, V>>>, batch_in_progress: Arc, From 0fb309288e72255166f80450f1a31922a92c3f24 Mon Sep 17 00:00:00 2001 From: Howard Wu <9260812+howardwu@users.noreply.github.com> Date: Sat, 15 Apr 2023 12:19:24 -0700 Subject: [PATCH 3/3] Convert unwraps to ?s --- synthesizer/src/store/helpers/memory_map.rs | 32 ++++++++++++++------- 1 file changed, 21 insertions(+), 11 deletions(-) diff --git a/synthesizer/src/store/helpers/memory_map.rs b/synthesizer/src/store/helpers/memory_map.rs index 4a96a2965a..d586ceb307 100644 --- a/synthesizer/src/store/helpers/memory_map.rs +++ b/synthesizer/src/store/helpers/memory_map.rs @@ -59,8 +59,10 @@ impl< { /// Initializes a new `MemoryMap` from the given iterator. fn from_iter>(iter: I) -> Self { + // Serialize each key in the iterator, and collect them into a map. + // Note: The 'unwrap' is safe here, because the keys are defined by us. let map = iter.into_iter().map(|(k, v)| (bincode::serialize(&k).unwrap(), v)).collect(); - + // Return the new map. Self { map: Arc::new(RwLock::new(map)), batch_in_progress: Default::default(), @@ -89,10 +91,10 @@ impl< } // Otherwise, insert the key-value pair directly into the map. false => { - let key = bincode::serialize(&key).unwrap(); - self.map.write().insert(key, value); + self.map.write().insert(bincode::serialize(&key)?, value); } } + Ok(()) } @@ -110,10 +112,10 @@ impl< } // Otherwise, remove the key-value pair directly from the map. false => { - let key = bincode::serialize(&key).unwrap(); - self.map.write().remove(&key); + self.map.write().remove(&bincode::serialize(&key)?); } } + Ok(()) } @@ -157,9 +159,17 @@ impl< if !operations.is_empty() { // Acquire a write lock on the map. let mut locked_map = self.map.write(); + + // Prepare the key for each queued operation. + // Note: This step is taken to ensure (with 100% certainty) that there will be + // no chance to fail partway through committing the queued operations. + let prepared_operations = operations + .into_iter() + .map(|(key, value)| Ok((bincode::serialize(&key)?, value))) + .collect::>>()?; + // Perform all the queued operations. - for (key, value) in operations { - let key = bincode::serialize(&key).unwrap(); + for (key, value) in prepared_operations { match value { Some(value) => locked_map.insert(key, value), None => locked_map.remove(&key), @@ -192,8 +202,7 @@ impl< K: Borrow, Q: PartialEq + Eq + Hash + Serialize + ?Sized, { - let key = bincode::serialize(key).unwrap(); - Ok(self.map.read().contains_key(&key)) + Ok(self.map.read().contains_key(&bincode::serialize(key)?)) } /// @@ -204,8 +213,7 @@ impl< K: Borrow, Q: PartialEq + Eq + Hash + Serialize + ?Sized, { - let key = bincode::serialize(key).unwrap(); - Ok(self.map.read().get(&key).cloned().map(Cow::Owned)) + Ok(self.map.read().get(&bincode::serialize(key)?).cloned().map(Cow::Owned)) } /// @@ -229,6 +237,7 @@ impl< /// Returns an iterator visiting each key-value pair in the map. /// fn iter(&'a self) -> Self::Iterator { + // Note: The 'unwrap' is safe here, because the keys are defined by us. self.map.read().clone().into_iter().map(|(k, v)| (Cow::Owned(bincode::deserialize(&k).unwrap()), Cow::Owned(v))) } @@ -236,6 +245,7 @@ impl< /// Returns an iterator over each key in the map. /// fn keys(&'a self) -> Self::Keys { + // Note: The 'unwrap' is safe here, because the keys are defined by us. self.map.read().clone().into_keys().map(|k| Cow::Owned(bincode::deserialize(&k).unwrap())) }