[refactor] replace the certificate_store with a wrapper struct (#825)

* first iteration

* fix after rebase

* update mysten-infra to use the new additions in the typed-store reverse iterator

* fix workerspace-hack issue

* fix: remove some double iterations from the read_all, delete_all operations

* fix: improve the efficiency of `after_round`

* fix: improve consensus comment

Co-authored-by: François Garillot <francois@garillot.net>

Author: akichidis

Cargo.lock

@@ -1,5 +1,5 @@
 [workspace]
-members = ["config", "consensus", "crypto", "dag", "examples", "executor", "network", "node", "primary", "test_utils", "types", "worker", "workspace-hack"]
+members = ["config", "consensus", "crypto", "dag", "examples", "executor", "network", "node", "primary", "storage", "test_utils", "types", "worker", "workspace-hack"]
 
 [profile.release]
 codegen-units = 1

Cargo.toml

@@ -13,14 +13,15 @@ bytes = "1.2.1"
 match_opt = "0.1.2"
 rand = { version = "0.8.5", optional = true }
 serde = { version = "1.0.144", features = ["derive"] }
-store = { git = "https://github.com/mystenlabs/mysten-infra.git", package = "typed-store", rev = "8d090689be14078f2ca41c356e7bbc0af21f73ab" }
+store = { git = "https://github.com/mystenlabs/mysten-infra.git", package = "typed-store", rev = "696fa25982acdc87ab64ec6a33cb111edc63a331" }
 thiserror = "1.0.32"
 tokio = { version = "1.20.1", features = ["sync"] }
 tracing = "0.1.36"
 
 config = { path = "../config" }
 fastcrypto = { git = "https://github.com/MystenLabs/fastcrypto", rev = "c022a2ae23ca7cc2778293fd3b1db42e8cd02d3b" }
 crypto = { path = "../crypto" }
+storage = { path = "../storage" }
 dag = { path = "../dag" }
 prometheus = "0.13.1"
 types = { path = "../types" }

consensus/Cargo.toml

@@ -10,7 +10,7 @@ use std::{
     collections::HashMap,
     sync::Arc,
 };
-use store::Store;
+use storage::CertificateStore;
 use tokio::{sync::watch, task::JoinHandle};
 use tracing::{info, instrument};
 use types::{
@@ -61,7 +61,7 @@ impl ConsensusState {
         genesis: Vec<Certificate>,
         metrics: Arc<ConsensusMetrics>,
         recover_last_committed: HashMap<PublicKey, Round>,
-        cert_store: Store<CertificateDigest, Certificate>,
+        cert_store: CertificateStore,
         gc_depth: Round,
     ) -> Self {
         let last_committed_round = *recover_last_committed
@@ -88,7 +88,7 @@ impl ConsensusState {
 
     #[instrument(level = "info", skip_all)]
     pub async fn construct_dag_from_cert_store(
-        cert_store: Store<CertificateDigest, Certificate>,
+        cert_store: CertificateStore,
         last_committed_round: Round,
         gc_depth: Round,
     ) -> Dag {
@@ -99,14 +99,11 @@ impl ConsensusState {
         );
 
         let min_round = last_committed_round.saturating_sub(gc_depth);
-        let cert_map = cert_store
-            .iter(Some(Box::new(move |(_dig, cert)| {
-                cert.header.round > min_round
-            })))
-            .await;
+        // get all certificates at a round > min_round
+        let cert_map = cert_store.after_round(min_round + 1).unwrap();
 
         let num_certs = cert_map.len();
-        for (digest, cert) in cert_map {
+        for (digest, cert) in cert_map.into_iter().map(|c| (c.digest(), c)) {
             let inner = dag.get_mut(&cert.header.round);
             match inner {
                 Some(m) => {
@@ -205,7 +202,7 @@ where
     pub fn spawn(
         committee: Committee,
         store: Arc<ConsensusStore>,
-        cert_store: Store<CertificateDigest, Certificate>,
+        cert_store: CertificateStore,
         rx_reconfigure: watch::Receiver<ReconfigureNotification>,
         rx_primary: metered_channel::Receiver<Certificate>,
         tx_primary: metered_channel::Sender<Certificate>,
@@ -248,7 +245,7 @@ where
     async fn run(
         &mut self,
         recover_last_committed: HashMap<PublicKey, Round>,
-        cert_store: Store<CertificateDigest, Certificate>,
+        cert_store: CertificateStore,
         gc_depth: Round,
     ) -> StoreResult<()> {
         // The consensus state (everything else is immutable).

consensus/src/consensus.rs

@@ -10,6 +10,7 @@ use fastcrypto::traits::KeyPair;
 use prometheus::Registry;
 #[cfg(test)]
 use std::collections::{BTreeSet, VecDeque};
+use storage::CertificateStore;
 use store::{reopen, rocks, rocks::DBMap};
 use test_utils::mock_committee;
 #[allow(unused_imports)]
@@ -32,19 +33,23 @@ pub fn make_consensus_store(store_path: &std::path::Path) -> Arc<ConsensusStore>
     Arc::new(ConsensusStore::new(last_committed_map, sequence_map))
 }
 
-pub fn make_certificate_store(
-    store_path: &std::path::Path,
-) -> store::Store<CertificateDigest, Certificate> {
+pub fn make_certificate_store(store_path: &std::path::Path) -> CertificateStore {
     const CERTIFICATES_CF: &str = "certificates";
-
-    let rocksdb =
-        rocks::open_cf(store_path, None, &[CERTIFICATES_CF]).expect("Failed creating database");
-
-    let certificate_map = reopen!(&rocksdb,
-        CERTIFICATES_CF;<CertificateDigest, Certificate>
+    const CERTIFICATE_ID_BY_ROUND_CF: &str = "certificate_id_by_round";
+
+    let rocksdb = rocks::open_cf(
+        store_path,
+        None,
+        &[CERTIFICATES_CF, CERTIFICATE_ID_BY_ROUND_CF],
+    )
+    .expect("Failed creating database");
+
+    let (certificate_map, certificate_id_by_round_map) = reopen!(&rocksdb,
+        CERTIFICATES_CF;<CertificateDigest, Certificate>,
+        CERTIFICATE_ID_BY_ROUND_CF;<(Round, CertificateDigest), u8>
     );
 
-    store::Store::new(certificate_map)
+    CertificateStore::new(certificate_map, certificate_id_by_round_map)
 }
 
 // Run for 4 dag rounds in ideal conditions (all nodes reference all other nodes). We should commit

consensus/src/tests/bullshark_tests.rs

@@ -10,6 +10,7 @@ use fastcrypto::traits::KeyPair;
 use prometheus::Registry;
 #[cfg(test)]
 use std::collections::{BTreeSet, VecDeque};
+use storage::CertificateStore;
 use store::{reopen, rocks, rocks::DBMap};
 use test_utils::mock_committee;
 #[allow(unused_imports)]
@@ -32,19 +33,23 @@ pub fn make_consensus_store(store_path: &std::path::Path) -> Arc<ConsensusStore>
     Arc::new(ConsensusStore::new(last_committed_map, sequence_map))
 }
 
-pub fn make_certificate_store(
-    store_path: &std::path::Path,
-) -> store::Store<CertificateDigest, Certificate> {
+pub fn make_certificate_store(store_path: &std::path::Path) -> CertificateStore {
     const CERTIFICATES_CF: &str = "certificates";
-
-    let rocksdb =
-        rocks::open_cf(store_path, None, &[CERTIFICATES_CF]).expect("Failed creating database");
-
-    let certificate_map = reopen!(&rocksdb,
-        CERTIFICATES_CF;<CertificateDigest, Certificate>
+    const CERTIFICATE_ID_BY_ROUND_CF: &str = "certificate_id_by_round";
+
+    let rocksdb = rocks::open_cf(
+        store_path,
+        None,
+        &[CERTIFICATES_CF, CERTIFICATE_ID_BY_ROUND_CF],
+    )
+    .expect("Failed creating database");
+
+    let (certificate_map, certificate_id_by_round_map) = reopen!(&rocksdb,
+        CERTIFICATES_CF;<CertificateDigest, Certificate>,
+        CERTIFICATE_ID_BY_ROUND_CF;<(Round, CertificateDigest), u8>
     );
 
-    store::Store::new(certificate_map)
+    CertificateStore::new(certificate_map, certificate_id_by_round_map)
 }
 
 // Run for 4 dag rounds in ideal conditions (all nodes reference all other nodes). We should commit

consensus/src/tests/tusk_tests.rs

@@ -26,9 +26,9 @@ prometheus = "0.13.1"
 backoff = { version = "0.4.0", features = ["tokio"] }
 
 types = { path = "../types" }
-mysten-network = { git = "https://github.com/mystenlabs/mysten-infra.git", rev = "8d090689be14078f2ca41c356e7bbc0af21f73ab" }
+mysten-network = { git = "https://github.com/mystenlabs/mysten-infra.git", rev = "696fa25982acdc87ab64ec6a33cb111edc63a331" }
 
-store = { git = "https://github.com/mystenlabs/mysten-infra.git", package = "typed-store", rev = "8d090689be14078f2ca41c356e7bbc0af21f73ab" }
+store = { git = "https://github.com/mystenlabs/mysten-infra.git", package = "typed-store", rev = "696fa25982acdc87ab64ec6a33cb111edc63a331" }
 workspace-hack = { version = "0.1", path = "../workspace-hack" }
 match_opt = "0.1.2"
 
@@ -39,4 +39,4 @@ tempfile = "3.3.0"
 primary = { path = "../primary" }
 test_utils = { path = "../test_utils" }
 types = { path = "../types" }
-telemetry-subscribers = { git = "https://github.com/mystenlabs/mysten-infra.git", rev = "8d090689be14078f2ca41c356e7bbc0af21f73ab" }
+telemetry-subscribers = { git = "https://github.com/mystenlabs/mysten-infra.git", rev = "696fa25982acdc87ab64ec6a33cb111edc63a331" }

executor/Cargo.toml

@@ -22,8 +22,8 @@ tonic = { version = "0.7.2", features = ["tls"] }
 tracing = "0.1.36"
 types = { path = "../types" }
 
-mysten-network = { git = "https://github.com/mystenlabs/mysten-infra.git", rev = "8d090689be14078f2ca41c356e7bbc0af21f73ab" }
 serde = "1.0.144"
+mysten-network = { git = "https://github.com/mystenlabs/mysten-infra.git", rev = "696fa25982acdc87ab64ec6a33cb111edc63a331" }
 workspace-hack = { version = "0.1", path = "../workspace-hack" }
 eyre = "0.6.8"
 

network/Cargo.toml

@@ -15,10 +15,10 @@ clap = "2.34"
 dhat = { version = "0.3.0", optional = true }
 futures = "0.3.23"
 multiaddr = "0.14.0"
-mysten-network = { git = "https://github.com/mystenlabs/mysten-infra.git", rev = "8d090689be14078f2ca41c356e7bbc0af21f73ab" }
+mysten-network = { git = "https://github.com/mystenlabs/mysten-infra.git", rev = "696fa25982acdc87ab64ec6a33cb111edc63a331" }
 rand = "0.8.5"
-store = { git = "https://github.com/mystenlabs/mysten-infra.git", package = "typed-store", rev = "8d090689be14078f2ca41c356e7bbc0af21f73ab" }
-telemetry-subscribers = { git = "https://github.com/mystenlabs/mysten-infra.git", rev = "8d090689be14078f2ca41c356e7bbc0af21f73ab" }
+store = { git = "https://github.com/mystenlabs/mysten-infra.git", package = "typed-store", rev = "696fa25982acdc87ab64ec6a33cb111edc63a331" }
+telemetry-subscribers = { git = "https://github.com/mystenlabs/mysten-infra.git", rev = "696fa25982acdc87ab64ec6a33cb111edc63a331" }
 thiserror = "1.0.32"
 tokio = { version = "1.20.1", features = ["full"] }
 tokio-stream = "0.1.9"
@@ -37,6 +37,7 @@ executor = { path = "../executor" }
 network = { path = "../network" }
 primary = { path = "../primary" }
 prometheus = "0.13.1"
+storage = { path = "../storage" }
 types = { path = "../types" }
 worker = { path = "../worker" }
 workspace-hack = { version = "0.1", path = "../workspace-hack" }

node/Cargo.toml

@@ -14,6 +14,7 @@ use fastcrypto::traits::{KeyPair as _, VerifyingKey};
 use primary::{BlockCommand, NetworkModel, PayloadToken, Primary, PrimaryChannelMetrics};
 use prometheus::{IntGauge, Registry};
 use std::{fmt::Debug, sync::Arc};
+use storage::{CertificateStore, CertificateToken};
 use store::{
     reopen,
     rocks::{open_cf, DBMap},
@@ -34,7 +35,7 @@ pub mod restarter;
 /// All the data stores of the node.
 pub struct NodeStorage {
     pub header_store: Store<HeaderDigest, Header>,
-    pub certificate_store: Store<CertificateDigest, Certificate>,
+    pub certificate_store: CertificateStore,
     pub payload_store: Store<(BatchDigest, WorkerId), PayloadToken>,
     pub batch_store: Store<BatchDigest, SerializedBatchMessage>,
     pub consensus_store: Arc<ConsensusStore>,
@@ -45,6 +46,7 @@ impl NodeStorage {
     /// The datastore column family names.
     const HEADERS_CF: &'static str = "headers";
     const CERTIFICATES_CF: &'static str = "certificates";
+    const CERTIFICATE_ID_BY_ROUND_CF: &'static str = "certificate_id_by_round";
     const PAYLOAD_CF: &'static str = "payload";
     const BATCHES_CF: &'static str = "batches";
     const LAST_COMMITTED_CF: &'static str = "last_committed";
@@ -59,6 +61,7 @@ impl NodeStorage {
             &[
                 Self::HEADERS_CF,
                 Self::CERTIFICATES_CF,
+                Self::CERTIFICATE_ID_BY_ROUND_CF,
                 Self::PAYLOAD_CF,
                 Self::BATCHES_CF,
                 Self::LAST_COMMITTED_CF,
@@ -71,6 +74,7 @@ impl NodeStorage {
         let (
             header_map,
             certificate_map,
+            certificate_id_by_round_map,
             payload_map,
             batch_map,
             last_committed_map,
@@ -79,6 +83,7 @@ impl NodeStorage {
         ) = reopen!(&rocksdb,
             Self::HEADERS_CF;<HeaderDigest, Header>,
             Self::CERTIFICATES_CF;<CertificateDigest, Certificate>,
+            Self::CERTIFICATE_ID_BY_ROUND_CF;<(Round, CertificateDigest), CertificateToken>,
             Self::PAYLOAD_CF;<(BatchDigest, WorkerId), PayloadToken>,
             Self::BATCHES_CF;<BatchDigest, SerializedBatchMessage>,
             Self::LAST_COMMITTED_CF;<PublicKey, Round>,
@@ -87,7 +92,7 @@ impl NodeStorage {
         );
 
         let header_store = Store::new(header_map);
-        let certificate_store = Store::new(certificate_map);
+        let certificate_store = CertificateStore::new(certificate_map, certificate_id_by_round_map);
         let payload_store = Store::new(payload_map);
         let batch_store = Store::new(batch_map);
         let consensus_store = Arc::new(ConsensusStore::new(last_committed_map, sequence_map));