Skip to content

Simplify Storage::rounds map. #3602

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.

Sign up for GitHub

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

Merged
kaimast merged 1 commit into from
May 20, 2025
Merged

Simplify Storage::rounds map. #3602

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
58 changes: 23 additions & 35 deletions node/bft/src/helpers/storage.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,13 +86,13 @@ pub struct StorageInner<N: Network> {
/// The maximum number of rounds to keep in storage.
max_gc_rounds: u64,
/* Once per batch */
/// The map of `round` to a list of `(certificate ID, batch ID, author)` entries.
rounds: RwLock<IndexMap<u64, IndexSet<(Field<N>, Field<N>, Address<N>)>>>,
/// The map of `round` to a list of `(certificate ID, author)` entries.
rounds: RwLock<IndexMap<u64, IndexSet<(Field<N>, Address<N>)>>>,
/// A cache of `certificate ID` to unprocessed `certificate`.
unprocessed_certificates: RwLock<LruCache<Field<N>, BatchCertificate<N>>>,
/// The map of `certificate ID` to `certificate`.
certificates: RwLock<IndexMap<Field<N>, BatchCertificate<N>>>,
/// The map of `batch ID` to `round`.
/// The map of `certificate ID` to `round`.
batch_ids: RwLock<IndexMap<Field<N>, u64>>,
/// The map of `transmission ID` to `(transmission, certificate IDs)` entries.
transmissions: Arc<dyn StorageService<N>>,
Expand Down Expand Up @@ -259,7 +259,7 @@ impl<N: Network> Storage<N> {

/// Returns `true` if the storage contains a certificate from the specified `author` in the given `round`.
pub fn contains_certificate_in_round_from(&self, round: u64, author: Address<N>) -> bool {
self.rounds.read().get(&round).map_or(false, |set| set.iter().any(|(_, _, a)| a == &author))
self.rounds.read().get(&round).map_or(false, |set| set.iter().any(|(_, a)| a == &author))
}

/// Returns `true` if the storage contains the specified `certificate ID`.
Expand Down Expand Up @@ -328,7 +328,7 @@ impl<N: Network> Storage<N> {
if let Some(entries) = self.rounds.read().get(&round) {
let certificates = self.certificates.read();
entries.iter().find_map(
|(certificate_id, _, a)| if a == &author { certificates.get(certificate_id).cloned() } else { None },
|(certificate_id, a)| if a == &author { certificates.get(certificate_id).cloned() } else { None },
)
} else {
Default::default()
Expand All @@ -345,7 +345,7 @@ impl<N: Network> Storage<N> {
// Retrieve the certificates.
if let Some(entries) = self.rounds.read().get(&round) {
let certificates = self.certificates.read();
entries.iter().flat_map(|(certificate_id, _, _)| certificates.get(certificate_id).cloned()).collect()
entries.iter().flat_map(|(certificate_id, _)| certificates.get(certificate_id).cloned()).collect()
} else {
Default::default()
}
Expand All @@ -360,7 +360,7 @@ impl<N: Network> Storage<N> {
}
// Retrieve the certificates.
if let Some(entries) = self.rounds.read().get(&round) {
entries.iter().map(|(certificate_id, _, _)| *certificate_id).collect()
entries.iter().map(|(certificate_id, _)| *certificate_id).collect()
} else {
Default::default()
}
Expand All @@ -375,7 +375,7 @@ impl<N: Network> Storage<N> {
}
// Retrieve the certificates.
if let Some(entries) = self.rounds.read().get(&round) {
entries.iter().map(|(_, _, author)| *author).collect()
entries.iter().map(|(_, author)| *author).collect()
} else {
Default::default()
}
Expand All @@ -392,7 +392,7 @@ impl<N: Network> Storage<N> {
cfg_sorted_by!(rounds.clone(), |a, _, b, _| a.cmp(b))
.flat_map(|(_, certificates_for_round)| {
// Iterate over the certificates for the round.
cfg_into_iter!(certificates_for_round).filter_map(|(certificate_id, _, _)| {
cfg_into_iter!(certificates_for_round).filter_map(|(certificate_id, _)| {
// Skip the certificate if it already exists in the ledger.
if self.ledger.contains_certificate(&certificate_id).unwrap_or(false) {
None
Expand Down Expand Up @@ -658,21 +658,19 @@ impl<N: Network> Storage<N> {
let round = certificate.round();
// Retrieve the certificate ID.
let certificate_id = certificate.id();
// Retrieve the batch ID.
let batch_id = certificate.batch_id();
// Retrieve the author of the batch.
let author = certificate.author();

// Insert the round to certificate ID entry.
self.rounds.write().entry(round).or_default().insert((certificate_id, batch_id, author));
self.rounds.write().entry(round).or_default().insert((certificate_id, author));
// Obtain the certificate's transmission ids.
let transmission_ids = certificate.transmission_ids().clone();
// Insert the certificate.
self.certificates.write().insert(certificate_id, certificate);
// Remove the unprocessed certificate.
self.unprocessed_certificates.write().pop(&certificate_id);
// Insert the batch ID.
self.batch_ids.write().insert(batch_id, round);
self.batch_ids.write().insert(certificate_id, round);
// Insert the certificate ID for each of the transmissions into storage.
self.transmissions.insert_transmissions(
certificate_id,
Expand Down Expand Up @@ -708,8 +706,6 @@ impl<N: Network> Storage<N> {
};
// Retrieve the round.
let round = certificate.round();
// Retrieve the batch ID.
let batch_id = certificate.batch_id();
// Compute the author of the batch.
let author = certificate.author();

Expand All @@ -721,7 +717,7 @@ impl<N: Network> Storage<N> {
match self.rounds.write().entry(round) {
Entry::Occupied(mut entry) => {
// Remove the round to certificate ID entry.
entry.get_mut().swap_remove(&(certificate_id, batch_id, author));
entry.get_mut().swap_remove(&(certificate_id, author));
// If the round is empty, remove it.
if entry.get().is_empty() {
entry.swap_remove();
Expand All @@ -734,7 +730,7 @@ impl<N: Network> Storage<N> {
// Remove the unprocessed certificate.
self.unprocessed_certificates.write().pop(&certificate_id);
// Remove the batch ID.
self.batch_ids.write().swap_remove(&batch_id);
self.batch_ids.write().swap_remove(&certificate_id);
// Remove the transmission entries in the certificate from storage.
self.transmissions.remove_transmissions(&certificate_id, certificate.transmission_ids());
// Return successfully.
Expand Down Expand Up @@ -876,7 +872,7 @@ impl<N: Network> Storage<N> {
}

/// Returns an iterator over the `(round, (certificate ID, batch ID, author))` entries.
pub fn rounds_iter(&self) -> impl Iterator<Item = (u64, IndexSet<(Field<N>, Field<N>, Address<N>)>)> {
pub fn rounds_iter(&self) -> impl Iterator<Item = (u64, IndexSet<(Field<N>, Address<N>)>)> {
self.rounds.read().clone().into_iter()
}

Expand Down Expand Up @@ -907,19 +903,17 @@ impl<N: Network> Storage<N> {
let round = certificate.round();
// Retrieve the certificate ID.
let certificate_id = certificate.id();
// Retrieve the batch ID.
let batch_id = certificate.batch_id();
// Retrieve the author of the batch.
let author = certificate.author();

// Insert the round to certificate ID entry.
self.rounds.write().entry(round).or_default().insert((certificate_id, batch_id, author));
self.rounds.write().entry(round).or_default().insert((certificate_id, author));
// Obtain the certificate's transmission ids.
let transmission_ids = certificate.transmission_ids().clone();
// Insert the certificate.
self.certificates.write().insert(certificate_id, certificate);
// Insert the batch ID.
self.batch_ids.write().insert(batch_id, round);
self.batch_ids.write().insert(certificate_id, round);

// Construct the dummy missing transmissions (for testing purposes).
let missing_transmissions = transmission_ids
Expand Down Expand Up @@ -954,7 +948,7 @@ pub(crate) mod tests {
/// Asserts that the storage matches the expected layout.
pub fn assert_storage<N: Network>(
storage: &Storage<N>,
rounds: &[(u64, IndexSet<(Field<N>, Field<N>, Address<N>)>)],
rounds: &[(u64, IndexSet<(Field<N>, Address<N>)>)],
certificates: &[(Field<N>, BatchCertificate<N>)],
batch_ids: &[(Field<N>, u64)],
transmissions: &HashMap<TransmissionID<N>, (Transmission<N>, IndexSet<Field<N>>)>,
Expand Down Expand Up @@ -1032,8 +1026,6 @@ pub(crate) mod tests {
let certificate_id = certificate.id();
// Retrieve the round.
let round = certificate.round();
// Retrieve the batch ID.
let batch_id = certificate.batch_id();
// Retrieve the author of the batch.
let author = certificate.author();

Expand All @@ -1052,11 +1044,11 @@ pub(crate) mod tests {
// Check that the underlying storage representation is correct.
{
// Construct the expected layout for 'rounds'.
let rounds = [(round, indexset! { (certificate_id, batch_id, author) })];
let rounds = [(round, indexset! { (certificate_id, author) })];
// Construct the expected layout for 'certificates'.
let certificates = [(certificate_id, certificate.clone())];
// Construct the expected layout for 'batch_ids'.
let batch_ids = [(batch_id, round)];
let batch_ids = [(certificate_id, round)];
// Assert the storage is well-formed.
assert_storage(&storage, &rounds, &certificates, &batch_ids, &transmissions);
}
Expand Down Expand Up @@ -1098,17 +1090,15 @@ pub(crate) mod tests {
let certificate_id = certificate.id();
// Retrieve the round.
let round = certificate.round();
// Retrieve the batch ID.
let batch_id = certificate.batch_id();
// Retrieve the author of the batch.
let author = certificate.author();

// Construct the expected layout for 'rounds'.
let rounds = [(round, indexset! { (certificate_id, batch_id, author) })];
let rounds = [(round, indexset! { (certificate_id, author) })];
// Construct the expected layout for 'certificates'.
let certificates = [(certificate_id, certificate.clone())];
// Construct the expected layout for 'batch_ids'.
let batch_ids = [(batch_id, round)];
let batch_ids = [(certificate_id, round)];
// Construct the sample 'transmissions'.
let (missing_transmissions, transmissions) = sample_transmissions(&certificate, rng);

Expand Down Expand Up @@ -1350,17 +1340,15 @@ pub mod prop_tests {

// Retrieve the round.
let round = certificate.round();
// Retrieve the batch ID.
let batch_id = certificate.batch_id();
// Retrieve the author of the batch.
let author = certificate.author();

// Construct the expected layout for 'rounds'.
let rounds = [(round, indexset! { (certificate_id, batch_id, author) })];
let rounds = [(round, indexset! { (certificate_id, author) })];
// Construct the expected layout for 'certificates'.
let certificates = [(certificate_id, certificate.clone())];
// Construct the expected layout for 'batch_ids'.
let batch_ids = [(batch_id, round)];
let batch_ids = [(certificate_id, round)];

// Insert the certificate.
let missing_transmissions: HashMap<TransmissionID<CurrentNetwork>, Transmission<CurrentNetwork>> =
Expand Down