[Perf] Cache the block tree on shutdown

Cargo.toml

@@ -45,9 +45,8 @@ version = "1.0.1"
 default-features = false
 
 [workspace.dependencies.snarkvm]
-#path = "../snarkVM"
-git = "https://github.com/ProvableHQ/snarkVM.git"
-rev = "fbd99608f"
+git = "https://github.com/ljedrz/snarkVM.git"
+branch = "perf/cache_block_tree"
 #version = "=4.3.0"
 default-features = false
 

node/src/client/mod.rs

@@ -23,7 +23,11 @@ use snarkos_node_cdn::CdnBlockSync;
 use snarkos_node_network::NodeType;
 use snarkos_node_rest::Rest;
 use snarkos_node_router::{
-    Heartbeat, Inbound, Outbound, Router, Routing,
+    Heartbeat,
+    Inbound,
+    Outbound,
+    Router,
+    Routing,
     messages::{Message, UnconfirmedSolution, UnconfirmedTransaction},
 };
 use snarkos_node_sync::{BLOCK_REQUEST_BATCH_DELAY, BlockSync, Ping, PrepareSyncRequest, locators::BlockLocators};
@@ -58,7 +62,8 @@ use std::{
     sync::{
         Arc,
         atomic::{
-            AtomicBool, AtomicUsize,
+            AtomicBool,
+            AtomicUsize,
             Ordering::{Acquire, Relaxed},
         },
     },
@@ -194,6 +199,9 @@ impl<N: Network, C: ConsensusStorage<N>> Client<N, C> {
             shutdown: shutdown.clone(),
         };
 
+        // Pass the node to the signal handler.
+        let _ = signal_node.set(node.clone());
+
         // Perform sync with CDN (if enabled).
         let cdn_sync = cdn.map(|base_url| {
             trace!("CDN sync is enabled");
@@ -229,8 +237,6 @@ impl<N: Network, C: ConsensusStorage<N>> Client<N, C> {
         node.initialize_execute_verification();
         // Initialize the notification message loop.
         node.handles.lock().push(crate::start_notification_message_loop());
-        // Pass the node to the signal handler.
-        let _ = signal_node.set(node.clone());
         // Return the node.
         Ok(node)
     }
@@ -617,6 +623,13 @@ impl<N: Network, C: ConsensusStorage<N>> NodeInterface<N> for Client<N, C> {
         // Shut down the router.
         self.router.shut_down().await;
 
+        // Cache the block tree.
+        let ledger = self.ledger.clone();
+        let _ = spawn_blocking!(ledger.cache_block_tree().map_err(|e| {
+            error!("Couldn't cache the block tree: {e}");
+            e
+        }));
+
         info!("Node has shut down.");
     }
 }

node/src/validator/mod.rs

@@ -24,7 +24,11 @@ use snarkos_node_consensus::Consensus;
 use snarkos_node_network::{NodeType, PeerPoolHandling};
 use snarkos_node_rest::Rest;
 use snarkos_node_router::{
-    Heartbeat, Inbound, Outbound, Router, Routing,
+    Heartbeat,
+    Inbound,
+    Outbound,
+    Router,
+    Routing,
     messages::{PuzzleResponse, UnconfirmedSolution, UnconfirmedTransaction},
 };
 use snarkos_node_sync::{BlockSync, Ping};
@@ -33,7 +37,8 @@ use snarkos_node_tcp::{
     protocols::{Disconnect, Handshake, OnConnect, Reading},
 };
 use snarkvm::prelude::{
-    Ledger, Network,
+    Ledger,
+    Network,
     block::{Block, Header},
     puzzle::Solution,
     store::ConsensusStorage,
@@ -152,6 +157,9 @@ impl<N: Network, C: ConsensusStorage<N>> Validator<N, C> {
             shutdown: shutdown.clone(),
         };
 
+        // Pass the node to the signal handler.
+        let _ = signal_node.set(node.clone());
+
         // Perform sync with CDN (if enabled).
         let cdn_sync = cdn.map(|base_url| Arc::new(CdnBlockSync::new(base_url, ledger.clone(), shutdown)));
 
@@ -187,8 +195,6 @@ impl<N: Network, C: ConsensusStorage<N>> Validator<N, C> {
         node.initialize_routing().await;
         // Initialize the notification message loop.
         node.handles.lock().push(crate::start_notification_message_loop());
-        // Pass the node to the signal handler.
-        let _ = signal_node.set(node.clone());
         // Return the node.
         Ok(node)
     }
@@ -475,6 +481,13 @@ impl<N: Network, C: ConsensusStorage<N>> NodeInterface<N> for Validator<N, C> {
         trace!("Shutting down consensus...");
         self.consensus.shut_down().await;
 
+        // Cache the block tree.
+        let ledger = self.ledger.clone();
+        let _ = spawn_blocking!(ledger.cache_block_tree().map_err(|e| {
+            error!("Couldn't cache the block tree: {e}");
+            e
+        }));
+
         info!("Node has shut down.");
     }
 }
@@ -483,7 +496,8 @@ impl<N: Network, C: ConsensusStorage<N>> NodeInterface<N> for Validator<N, C> {
 mod tests {
     use super::*;
     use snarkvm::prelude::{
-        MainnetV0, VM,
+        MainnetV0,
+        VM,
         store::{ConsensusStore, helpers::memory::ConsensusMemory},
     };