Skip to content

Commit

Permalink
Merge pull request #2561 from ProvableHQ/perf/raw_rocks_iter
Browse files Browse the repository at this point in the history 
[Perf] Only use raw iterators with RocksDB and speed up ledger load
  • Loading branch information
@alzger
alzger authored Nov 4, 2024
2 parents ecddac5 + 591f3bf commit 460f637
Show file tree
Hide file tree
Showing 3 changed files with 91 additions and 95 deletions.
22 changes: 7 additions & 15 deletions ledger/store/src/block/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1009,21 +1009,13 @@ impl<N: Network, B: BlockStorage<N>> BlockStore<N, B> {

// Compute the block tree.
let tree = {
// Find the maximum block height.
let max_height = storage.id_map().len_confirmed().checked_sub(1).map(u32::try_from);
// Prepare the leaves of the block tree.
let hashes = match max_height {
Some(height) => {
let height = height?;
cfg_into_iter!(0..=height)
.map(|height| match storage.get_block_hash(height)? {
Some(hash) => Ok(hash.to_bits_le()),
None => bail!("Missing block hash for block {height}"),
})
.collect::<Result<Vec<Vec<bool>>>>()?
}
None => vec![],
};
// Prepare an iterator over the block heights and prepare the leaves of the block tree.
let hashes = storage
.id_map()
.iter_confirmed()
.sorted_unstable_by(|(h1, _), (h2, _)| h1.cmp(h2))
.map(|(_, hash)| hash.to_bits_le())
.collect::<Vec<Vec<bool>>>();
// Construct the block tree.
Arc::new(RwLock::new(N::merkle_tree_bhp(&hashes)?))
};
Expand Down
78 changes: 39 additions & 39 deletions ledger/store/src/helpers/rocksdb/internal/map.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -277,17 +277,21 @@ impl<

// Count the number of keys belonging to the map.
let mut len = 0usize;
while let Some(key) = iter.key() {
// Only compare the map ID - the network ID is guaranteed to
// remain the same as long as there is more than a single map.
if key[2..][..2] != self.context[2..][..2] {
// If the map ID is different, it's the end of iteration.
while iter.valid() {
if let Some(key) = iter.key() {
// Only compare the map ID - the network ID is guaranteed to
// remain the same as long as there is more than a single map.
if key[2..][..2] != self.context[2..][..2] {
// If the map ID is different, it's the end of iteration.
break;
}

// Increment the length and go to the next record.
len += 1;
iter.next();
} else {
break;
}

// Increment the length and go to the next record.
len += 1;
iter.next();
}

len
Expand Down Expand Up @@ -401,7 +405,7 @@ pub struct Iter<
K: 'a + Debug + PartialEq + Eq + Hash + Serialize + DeserializeOwned,
V: 'a + PartialEq + Eq + Serialize + DeserializeOwned,
> {
db_iter: rocksdb::DBIterator<'a>,
db_iter: rocksdb::DBRawIterator<'a>,
_phantom: PhantomData<(K, V)>,
}

Expand All @@ -412,7 +416,7 @@ impl<
> Iter<'a, K, V>
{
pub(super) fn new(db_iter: rocksdb::DBIterator<'a>) -> Self {
Self { db_iter, _phantom: PhantomData }
Self { db_iter: db_iter.into(), _phantom: PhantomData }
}
}

Expand All @@ -425,96 +429,92 @@ impl<
type Item = (Cow<'a, K>, Cow<'a, V>);

fn next(&mut self) -> Option<Self::Item> {
let (key, value) = self
.db_iter
.next()?
.map_err(|e| {
error!("RocksDB Iter iterator error: {e}");
})
.ok()?;
if !self.db_iter.valid() {
return None;
}

let (key, value) = self.db_iter.item()?;

// Deserialize the key and value.
let key = bincode::deserialize(&key[PREFIX_LEN..])
.map_err(|e| {
error!("RocksDB Iter deserialize(key) error: {e}");
})
.ok()?;
let value = bincode::deserialize(&value)
let value = bincode::deserialize(value)
.map_err(|e| {
error!("RocksDB Iter deserialize(value) error: {e}");
})
.ok()?;

self.db_iter.next();

Some((Cow::Owned(key), Cow::Owned(value)))
}
}

/// An iterator over the keys of a prefix.
pub struct Keys<'a, K: 'a + Debug + PartialEq + Eq + Hash + Serialize + DeserializeOwned> {
db_iter: rocksdb::DBIterator<'a>,
db_iter: rocksdb::DBRawIterator<'a>,
_phantom: PhantomData<K>,
}

impl<'a, K: 'a + Debug + PartialEq + Eq + Hash + Serialize + DeserializeOwned> Keys<'a, K> {
pub(crate) fn new(db_iter: rocksdb::DBIterator<'a>) -> Self {
Self { db_iter, _phantom: PhantomData }
Self { db_iter: db_iter.into(), _phantom: PhantomData }
}
}

impl<'a, K: 'a + Clone + Debug + PartialEq + Eq + Hash + Serialize + DeserializeOwned> Iterator for Keys<'a, K> {
type Item = Cow<'a, K>;

fn next(&mut self) -> Option<Self::Item> {
let (key, _) = self
.db_iter
.next()?
.map_err(|e| {
error!("RocksDB Keys iterator error: {e}");
})
.ok()?;
if !self.db_iter.valid() {
return None;
}

// Deserialize the key.
let key = bincode::deserialize(&key[PREFIX_LEN..])
let key = bincode::deserialize(&self.db_iter.key()?[PREFIX_LEN..])
.map_err(|e| {
error!("RocksDB Keys deserialize(key) error: {e}");
})
.ok()?;

self.db_iter.next();

Some(Cow::Owned(key))
}
}

/// An iterator over the values of a prefix.
pub struct Values<'a, V: 'a + PartialEq + Eq + Serialize + DeserializeOwned> {
db_iter: rocksdb::DBIterator<'a>,
db_iter: rocksdb::DBRawIterator<'a>,
_phantom: PhantomData<V>,
}

impl<'a, V: 'a + PartialEq + Eq + Serialize + DeserializeOwned> Values<'a, V> {
pub(crate) fn new(db_iter: rocksdb::DBIterator<'a>) -> Self {
Self { db_iter, _phantom: PhantomData }
Self { db_iter: db_iter.into(), _phantom: PhantomData }
}
}

impl<'a, V: 'a + Clone + PartialEq + Eq + Serialize + DeserializeOwned> Iterator for Values<'a, V> {
type Item = Cow<'a, V>;

fn next(&mut self) -> Option<Self::Item> {
let (_, value) = self
.db_iter
.next()?
.map_err(|e| {
error!("RocksDB Values iterator error: {e}");
})
.ok()?;
if !self.db_iter.valid() {
return None;
}

// Deserialize the value.
let value = bincode::deserialize(&value)
let value = bincode::deserialize(self.db_iter.value()?)
.map_err(|e| {
error!("RocksDB Values deserialize(value) error: {e}");
})
.ok()?;

self.db_iter.next();

Some(Cow::Owned(value))
}
}
Expand Down
86 changes: 45 additions & 41 deletions ledger/store/src/helpers/rocksdb/internal/nested_map.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -372,18 +372,22 @@ impl<

// Count the number of keys belonging to the nested map.
let mut len = 0usize;
while let Some(key) = iter.key() {
// Only compare the nested map - the network ID and the outer map
// ID are guaranteed to remain the same as long as there is more
// than a single map in the database.
if !key[PREFIX_LEN + 4..].starts_with(serialized_map) {
// If the nested map ID is different, it's the end of iteration.
while iter.valid() {
if let Some(key) = iter.key() {
// Only compare the nested map - the network ID and the outer map
// ID are guaranteed to remain the same as long as there is more
// than a single map in the database.
if !key[PREFIX_LEN + 4..].starts_with(serialized_map) {
// If the nested map ID is different, it's the end of iteration.
break;
}

// Increment the length and go to the next record.
len += 1;
iter.next();
} else {
break;
}

// Increment the length and go to the next record.
len += 1;
iter.next();
}

Ok(len)
Expand Down Expand Up @@ -595,7 +599,7 @@ pub struct NestedIter<
K: 'a + Debug + PartialEq + Eq + Serialize + DeserializeOwned,
V: 'a + PartialEq + Eq + Serialize + DeserializeOwned,
> {
db_iter: rocksdb::DBIterator<'a>,
db_iter: rocksdb::DBRawIterator<'a>,
_phantom: PhantomData<(M, K, V)>,
}

Expand All @@ -607,7 +611,7 @@ impl<
> NestedIter<'a, M, K, V>
{
pub(super) fn new(db_iter: rocksdb::DBIterator<'a>) -> Self {
Self { db_iter, _phantom: PhantomData }
Self { db_iter: db_iter.into(), _phantom: PhantomData }
}
}

Expand All @@ -621,16 +625,14 @@ impl<
type Item = (Cow<'a, M>, Cow<'a, K>, Cow<'a, V>);

fn next(&mut self) -> Option<Self::Item> {
let (map_key, value) = self
.db_iter
.next()?
.map_err(|e| {
error!("RocksDB NestedIter iterator error: {e}");
})
.ok()?;
if !self.db_iter.valid() {
return None;
}

let (map_key, value) = self.db_iter.item()?;

// Extract the bytes belonging to the map and the key.
let (entry_map, entry_key) = get_map_and_key(&map_key)
let (entry_map, entry_key) = get_map_and_key(map_key)
.map_err(|e| {
error!("RocksDB NestedIter get_map_and_key error: {e}");
})
Expand All @@ -648,12 +650,14 @@ impl<
})
.ok()?;
// Deserialize the value.
let value = bincode::deserialize(&value)
let value = bincode::deserialize(value)
.map_err(|e| {
error!("RocksDB NestedIter deserialize(value) error: {e}");
})
.ok()?;

self.db_iter.next();

Some((Cow::Owned(map), Cow::Owned(key), Cow::Owned(value)))
}
}
Expand All @@ -664,7 +668,7 @@ pub struct NestedKeys<
M: 'a + Clone + Debug + PartialEq + Eq + Hash + Serialize + DeserializeOwned,
K: 'a + Clone + Debug + PartialEq + Eq + Serialize + DeserializeOwned,
> {
db_iter: rocksdb::DBIterator<'a>,
db_iter: rocksdb::DBRawIterator<'a>,
_phantom: PhantomData<(M, K)>,
}

Expand All @@ -675,7 +679,7 @@ impl<
> NestedKeys<'a, M, K>
{
pub(crate) fn new(db_iter: rocksdb::DBIterator<'a>) -> Self {
Self { db_iter, _phantom: PhantomData }
Self { db_iter: db_iter.into(), _phantom: PhantomData }
}
}

Expand All @@ -688,16 +692,14 @@ impl<
type Item = (Cow<'a, M>, Cow<'a, K>);

fn next(&mut self) -> Option<Self::Item> {
let (map_key, _) = self
.db_iter
.next()?
.map_err(|e| {
error!("RocksDB NestedKeys iterator error: {e}");
})
.ok()?;
if !self.db_iter.valid() {
return None;
}

let map_key = self.db_iter.key()?;

// Extract the bytes belonging to the map and the key.
let (entry_map, entry_key) = get_map_and_key(&map_key)
let (entry_map, entry_key) = get_map_and_key(map_key)
.map_err(|e| {
error!("RocksDB NestedKeys get_map_and_key error: {e}");
})
Expand All @@ -715,41 +717,43 @@ impl<
})
.ok()?;

self.db_iter.next();

Some((Cow::Owned(map), Cow::Owned(key)))
}
}

/// An iterator over the values of a prefix.
pub struct NestedValues<'a, V: 'a + PartialEq + Eq + Serialize + DeserializeOwned> {
db_iter: rocksdb::DBIterator<'a>,
db_iter: rocksdb::DBRawIterator<'a>,
_phantom: PhantomData<V>,
}

impl<'a, V: 'a + PartialEq + Eq + Serialize + DeserializeOwned> NestedValues<'a, V> {
pub(crate) fn new(db_iter: rocksdb::DBIterator<'a>) -> Self {
Self { db_iter, _phantom: PhantomData }
Self { db_iter: db_iter.into(), _phantom: PhantomData }
}
}

impl<'a, V: 'a + Clone + PartialEq + Eq + Serialize + DeserializeOwned> Iterator for NestedValues<'a, V> {
type Item = Cow<'a, V>;

fn next(&mut self) -> Option<Self::Item> {
let (_, value) = self
.db_iter
.next()?
.map_err(|e| {
error!("RocksDB NestedValues iterator error: {e}");
})
.ok()?;
if !self.db_iter.valid() {
return None;
}

let value = self.db_iter.value()?;

// Deserialize the value.
let value = bincode::deserialize(&value)
let value = bincode::deserialize(value)
.map_err(|e| {
error!("RocksDB NestedValues deserialize(value) error: {e}");
})
.ok()?;

self.db_iter.next();

Some(Cow::Owned(value))
}
}
Expand Down

0 comments on commit 460f637

Please sign in to comment.