Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Perf] Only use raw iterators with RocksDB and speed up ledger load #2518

Merged
merged 3 commits into from
Nov 4, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 3 additions & 10 deletions ledger/store/src/block/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1015,17 +1015,10 @@ impl<N: Network, B: BlockStorage<N>> BlockStore<N, B> {
// Compute the block tree.
let tree = {
// Prepare an iterator over the block heights.
let heights = storage.id_map().keys_confirmed();
let mut heights_hashes = storage.id_map().iter_confirmed().collect::<Vec<_>>();
heights_hashes.sort_unstable_by(|(h1, _), (h2, _)| h1.cmp(h2));
// Prepare the leaves of the block tree.
let hashes = match heights.max() {
Some(height) => cfg_into_iter!(0..=cow_to_copied!(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![],
};
let hashes = cfg_into_iter!(heights_hashes).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
Original file line number Diff line number Diff line change
Expand Up @@ -276,17 +276,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 @@ -400,7 +404,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 @@ -411,7 +415,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 @@ -424,96 +428,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
Original file line number Diff line number Diff line change
Expand Up @@ -371,18 +371,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 @@ -594,7 +598,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 @@ -606,7 +610,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 @@ -620,16 +624,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 @@ -647,12 +649,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 @@ -663,7 +667,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 @@ -674,7 +678,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 @@ -687,16 +691,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 @@ -714,41 +716,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