Merge pull request #3030 from ljedrz/perf/cache_block_tree

[Perf] Facilitate the caching of the block tree

Author: vicsn

Cargo.lock

@@ -54,6 +54,10 @@ version = "0.13"
 features = [ "digest", "ecdsa" ]
 default-features = false
 
+[dependencies.serde]
+workspace = true
+features = ["rc"]
+
 [dependencies.smallvec]
 workspace = true
 features = [ "const_generics", "const_new" ]
@@ -77,9 +81,6 @@ workspace = true
 [dev-dependencies.hex]
 workspace = true
 
-[dev-dependencies.serde]
-workspace = true
-
 [dev-dependencies.serde_json]
 workspace = true
 features = [ "preserve_order" ]

console/algorithms/Cargo.toml

@@ -19,6 +19,7 @@ use crate::Blake2Xs;
 use snarkvm_console_types::prelude::*;
 use snarkvm_utilities::BigInteger;
 
+use serde::{Deserialize, Serialize};
 use std::sync::Arc;
 
 /// The BHP chunk size (this implementation is for a 3-bit BHP).
@@ -27,7 +28,8 @@ pub(super) const BHP_LOOKUP_SIZE: usize = 1 << BHP_CHUNK_SIZE;
 
 /// BHP is a collision-resistant hash function that takes a variable-length input.
 /// The BHP hasher is used to process one internal iteration of the BHP hash function.
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Debug, PartialEq, Deserialize, Serialize)]
+#[serde(bound = "E: Serialize + DeserializeOwned")]
 pub struct BHPHasher<E: Environment, const NUM_WINDOWS: u8, const WINDOW_SIZE: u8> {
     /// The bases for the BHP hash.
     bases: Arc<Vec<Vec<Group<E>>>>,

console/algorithms/src/bhp/hasher/mod.rs

@@ -23,6 +23,7 @@ mod hash_uncompressed;
 
 use snarkvm_console_types::prelude::*;
 
+use serde::{Deserialize, Serialize};
 use std::sync::Arc;
 
 const BHP_CHUNK_SIZE: usize = 3;
@@ -50,7 +51,8 @@ pub type BHP1024<E> = BHP<E, 8, 54>; // Supports inputs up to 1044 bits (4 u8 +
 /// ```text
 /// DIGEST_N+1 = BHP([ DIGEST_N[0..DATA_BITS] || INPUT[(N+1)*BLOCK_SIZE..(N+2)*BLOCK_SIZE] ]);
 /// ```
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Debug, PartialEq, Deserialize, Serialize)]
+#[serde(bound = "E: Serialize + DeserializeOwned")]
 pub struct BHP<E: Environment, const NUM_WINDOWS: u8, const WINDOW_SIZE: u8> {
     /// The domain separator for the BHP hash function.
     domain: Vec<bool>,

console/algorithms/src/bhp/mod.rs

@@ -31,6 +31,9 @@ workspace = true
 [dependencies.rayon]
 workspace = true
 
+[dependencies.serde]
+workspace = true
+
 [dev-dependencies.snarkvm-console-network]
 path = "../network"
 

console/collections/Cargo.toml

@@ -26,12 +26,14 @@ use snarkvm_console_types::prelude::*;
 
 use aleo_std::prelude::*;
 
+use serde::{Deserialize, Serialize};
 use std::collections::BTreeMap;
 
 #[cfg(not(feature = "serial"))]
 use rayon::prelude::*;
 
-#[derive(Clone)]
+#[derive(Clone, Deserialize, Serialize)]
+#[serde(bound = "E: Serialize + DeserializeOwned, LH: Serialize + DeserializeOwned, PH: Serialize + DeserializeOwned")]
 pub struct MerkleTree<E: Environment, LH: LeafHash<Hash = PH::Hash>, PH: PathHash<Hash = Field<E>>, const DEPTH: u8> {
     /// The leaf hasher for the Merkle tree.
     leaf_hasher: LH,

console/collections/src/merkle_tree/mod.rs

@@ -114,14 +114,24 @@ pub enum RecordsFilter<N: Network> {
 /// which loads the ledger from storage,
 /// or initializes it with the genesis block if the storage is empty
 #[derive(Clone)]
-pub struct Ledger<N: Network, C: ConsensusStorage<N>> {
+pub struct Ledger<N: Network, C: ConsensusStorage<N>>(Arc<InnerLedger<N, C>>);
+
+impl<N: Network, C: ConsensusStorage<N>> Deref for Ledger<N, C> {
+    type Target = InnerLedger<N, C>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+#[doc(hidden)]
+pub struct InnerLedger<N: Network, C: ConsensusStorage<N>> {
     /// The VM state.
     vm: VM<N, C>,
     /// The genesis block.
     genesis_block: Block<N>,
     /// The current epoch hash.
-    current_epoch_hash: Arc<RwLock<Option<N::BlockHash>>>,
-
+    current_epoch_hash: RwLock<Option<N::BlockHash>>,
     /// The committee resulting from all the on-chain staking activity.
     ///
     /// This includes any bonding and unbonding transactions in the latest block.
@@ -137,23 +147,21 @@ pub struct Ledger<N: Network, C: ConsensusStorage<N>> {
     /// So the `Option` should always be `Some`,
     /// but there are cases in which it is `None`,
     /// probably only temporarily when loading/initializing the ledger,
-    current_committee: Arc<RwLock<Option<Committee<N>>>>,
+    current_committee: RwLock<Option<Committee<N>>>,
 
     /// The latest block that was added to the ledger.
     ///
     /// This lock is also used as a way to prevent concurrent updates to the ledger, and to ensure that
     /// the ledger does not advance while certain check happen.
-    current_block: Arc<RwLock<Block<N>>>,
-
+    current_block: RwLock<Block<N>>,
     /// The recent committees of interest paired with their applicable rounds.
     ///
     /// Each entry consisting of a round `R` and a committee `C`,
     /// says that `C` is the bonded committee at round `R`,
     /// i.e. resulting from all the bonding and unbonding transactions before `R`.
     /// If `L` is the lookback round distance, `C` is the active committee at round `R + L`
     /// (i.e. the committee in charge of running consensus at round `R + L`).
-    committee_cache: Arc<Mutex<LruCache<u64, Committee<N>>>>,
-
+    committee_cache: Mutex<LruCache<u64, Committee<N>>>,
     /// The cache that holds the provers and the number of solutions they have submitted for the current epoch.
     epoch_provers_cache: Arc<RwLock<IndexMap<Address<N>, u32>>>,
 }
@@ -211,47 +219,73 @@ impl<N: Network, C: ConsensusStorage<N>> Ledger<N, C> {
         let current_committee = vm.finalize_store().committee_store().current_committee().ok();
 
         // Create a committee cache.
-        let committee_cache = Arc::new(Mutex::new(LruCache::new(COMMITTEE_CACHE_SIZE.try_into().unwrap())));
+        let committee_cache = Mutex::new(LruCache::new(COMMITTEE_CACHE_SIZE.try_into().unwrap()));
 
         // Initialize the ledger.
-        let mut ledger = Self {
+        let ledger = Self(Arc::new(InnerLedger {
             vm,
             genesis_block: genesis_block.clone(),
             current_epoch_hash: Default::default(),
-            current_committee: Arc::new(RwLock::new(current_committee)),
-            current_block: Arc::new(RwLock::new(genesis_block.clone())),
+            current_committee: RwLock::new(current_committee),
+            current_block: RwLock::new(genesis_block.clone()),
             committee_cache,
             epoch_provers_cache: Default::default(),
-        };
+        }));
+
+        // Attempt to obtain the maximum height from the storage.
+        let max_stored_height = ledger.vm.block_store().max_height();
 
         // If the block store is empty, add the genesis block.
-        if ledger.vm.block_store().max_height().is_none() {
-            // Add the genesis block.
+        let latest_height = if let Some(max_height) = max_stored_height {
+            max_height
+        } else {
             ledger.advance_to_next_block(&genesis_block)?;
-        }
+            0
+        };
         lap!(timer, "Initialize genesis");
 
-        // Retrieve the latest height.
-        let latest_height =
-            ledger.vm.block_store().max_height().with_context(|| "Failed to load blocks from the ledger")?;
+        // Ensure that the greatest stored height matches that of the block tree.
+        ensure!(
+            latest_height == ledger.vm().block_store().current_block_height(),
+            "The stored height is different than the one in the block tree; \
+            please ensure that the cached block tree is valid or delete the \
+            'block_tree' file from the ledger folder"
+        );
+
+        // Verify that the root of the cached block tree matches the one in the storage.
+        let tree_root = <N::StateRoot>::from(ledger.vm().block_store().get_block_tree_root());
+        let state_root = ledger
+            .vm()
+            .block_store()
+            .get_state_root(latest_height)?
+            .ok_or_else(|| anyhow!("Missing state root in the storage"))?;
+        ensure!(
+            tree_root == state_root,
+            "The stored state root is different than the one in the block tree;
+            please ensure that the cached block tree is valid or delete the \
+            'block_tree' file from the ledger folder"
+        );
+
         // Fetch the latest block.
         let block = ledger
             .get_block(latest_height)
             .with_context(|| format!("Failed to load block {latest_height} from the ledger"))?;
 
         // Set the current block.
-        ledger.current_block = Arc::new(RwLock::new(block));
+        *ledger.current_block.write() = block;
         // Set the current committee (and ensures the latest committee exists).
-        ledger.current_committee = Arc::new(RwLock::new(Some(ledger.latest_committee()?)));
+        *ledger.current_committee.write() = Some(ledger.latest_committee()?);
         // Set the current epoch hash.
-        ledger.current_epoch_hash = Arc::new(RwLock::new(Some(ledger.get_epoch_hash(latest_height)?)));
+        *ledger.current_epoch_hash.write() = Some(ledger.get_epoch_hash(latest_height)?);
         // Set the epoch prover cache.
-        ledger.epoch_provers_cache = Arc::new(RwLock::new(ledger.load_epoch_provers()));
+        *ledger.epoch_provers_cache.write() = ledger.load_epoch_provers();
 
         finish!(timer, "Initialize ledger");
         Ok(ledger)
     }
+}
 
+impl<N: Network, C: ConsensusStorage<N>> Ledger<N, C> {
     /// Creates a rocksdb checkpoint in the specified directory, which needs to not exist at the
     /// moment of calling. The checkpoints are based on hard links, which means they can both be
     /// incremental (i.e. they aren't full physical copies), and used as full rollback points
@@ -261,6 +295,11 @@ impl<N: Network, C: ConsensusStorage<N>> Ledger<N, C> {
         self.vm.block_store().backup_database(path).map_err(|err| anyhow!(err))
     }
 
+    #[cfg(feature = "rocks")]
+    pub fn cache_block_tree(&self) -> Result<()> {
+        self.vm.block_store().cache_block_tree()
+    }
+
     /// Loads the provers and the number of solutions they have submitted for the current epoch.
     pub fn load_epoch_provers(&self) -> IndexMap<Address<N>, u32> {
         // Fetch the block heights that belong to the current epoch.
@@ -285,12 +324,12 @@ impl<N: Network, C: ConsensusStorage<N>> Ledger<N, C> {
     }
 
     /// Returns the VM.
-    pub const fn vm(&self) -> &VM<N, C> {
+    pub fn vm(&self) -> &VM<N, C> {
         &self.vm
     }
 
     /// Returns the puzzle.
-    pub const fn puzzle(&self) -> &Puzzle<N> {
+    pub fn puzzle(&self) -> &Puzzle<N> {
         self.vm.puzzle()
     }
 
@@ -481,6 +520,20 @@ impl<N: Network, C: ConsensusStorage<N>> Ledger<N, C> {
     }
 }
 
+#[cfg(feature = "rocks")]
+impl<N: Network, C: ConsensusStorage<N>> Drop for InnerLedger<N, C> {
+    fn drop(&mut self) {
+        // Cache the block tree in order to speed up the next startup; this operation
+        // is guaranteed to conclude as long as the destructors are allowed to run
+        // (a clean shutdown, panic = "unwind", an explicit call to `drop`, etc.).
+        // At the moment this code is executed, the Ledger is guaranteed to be owned
+        // exclusively by this method, so no other activity may interrupt it.
+        if let Err(e) = self.vm.block_store().cache_block_tree() {
+            error!("Couldn't cache the block tree: {e}");
+        }
+    }
+}
+
 pub mod prelude {
     pub use crate::{Ledger, authority, block, block::*, committee, helpers::*, narwhal, puzzle, query, store};
 }

ledger/src/lib.rs

@@ -48,13 +48,16 @@ use snarkvm_ledger_puzzle::{Solution, SolutionID};
 use snarkvm_synthesizer_program::{FinalizeOperation, Program};
 
 use aleo_std_storage::StorageMode;
+#[cfg(feature = "rocks")]
+use aleo_std_storage::aleo_ledger_dir;
 use anyhow::{Context, Result};
 #[cfg(feature = "locktick")]
 use locktick::{LockGuard, parking_lot::RwLock};
 #[cfg(not(feature = "locktick"))]
 use parking_lot::RwLock;
 use std::{borrow::Cow, sync::Arc};
-use tracing::debug;
+#[cfg(feature = "rocks")]
+use std::{fs, io::BufWriter};
 
 #[cfg(not(feature = "serial"))]
 use rayon::prelude::*;
@@ -991,6 +994,8 @@ pub trait BlockStorage<N: Network>: 'static + Clone + Send + Sync {
 
     #[cfg(feature = "rocks")]
     fn backup_database<P: AsRef<std::path::Path>>(&self, path: P) -> Result<(), String>;
+
+    fn create_block_tree(&self) -> Result<BlockTree<N>>;
 }
 
 /// The `BlockStore` is the user facing API that either uses `BlockMemory` or `BlockDB` as its storae backend.
@@ -1009,17 +1014,7 @@ impl<N: Network, B: BlockStorage<N>> BlockStore<N, B> {
     pub fn open<S: Into<StorageMode>>(storage: S) -> Result<Self> {
         let storage = B::open(storage)?;
 
-        // 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.
-        debug!("Found {} blocks on disk", hashes.len());
-        let tree = N::merkle_tree_bhp(&hashes)?;
+        let tree = storage.create_block_tree()?;
 
         let mut initial_cache = Vec::new();
         let cache_end_height = u32::try_from(tree.number_of_leaves())?;
@@ -1032,10 +1027,12 @@ impl<N: Network, B: BlockStorage<N>> BlockStore<N, B> {
             }
 
             // Get the hash for the next block to add to the cache.
-            let hash = storage
-                .id_map()
-                .get_confirmed(&height)?
-                .with_context(|| format!("Block {height} exists in tree but not in storage"))?;
+            let hash = storage.id_map().get_confirmed(&height)?.with_context(|| {
+                format!(
+                    "Block {height} exists in tree but not in storage;\
+                    perhaps you used a wrong block tree cache file?"
+                )
+            })?;
 
             initial_cache.push(
                 storage.get_block(&hash)?.with_context(|| format!("Block {hash} exists in tree but not in storage"))?,
@@ -1200,10 +1197,40 @@ impl<N: Network, B: BlockStorage<N>> BlockStore<N, B> {
         self.storage.unpause_atomic_writes::<DISCARD_BATCH>()
     }
 
+    /// Stores a database backup at the given location.
     #[cfg(feature = "rocks")]
     pub fn backup_database<P: AsRef<std::path::Path>>(&self, path: P) -> Result<(), String> {
         self.storage.backup_database(path)
     }
+
+    /// Serializes and persists the current block tree.
+    #[cfg(feature = "rocks")]
+    pub fn cache_block_tree(&self) -> Result<()> {
+        // Prepare the path for the target file.
+        let mut path = aleo_ledger_dir(N::ID, self.storage.storage_mode());
+        path.push("block_tree");
+
+        // Create the target file.
+        let file = fs::File::create(path)?;
+        // The block tree can become quite large, so use a BufWriter in order to
+        // not have to keep the entire serialized tree in memory, and to reduce
+        // the number of syscalls involved with disk writes. 1MiB should provide
+        // a good balance between the CPU cache and maximum disk throughput.
+        let mut writer = BufWriter::with_capacity(1024 * 1024, file);
+        bincode::serialize_into(&mut writer, &&*self.tree.read())?;
+        writer.flush()?;
+        // TODO(ljedrz): this operation can already take ~2.5s, so we may want
+        // to perform chunking and parallel serialization. This may be useful
+        // for other applications, so it should be implemented as a common
+        // utility.
+
+        Ok(())
+    }
+
+    /// Returns the root of the block tree.
+    pub fn get_block_tree_root(&self) -> Field<N> {
+        *self.tree.read().root()
+    }
 }
 
 impl<N: Network, B: BlockStorage<N>> BlockStore<N, B> {