Merge pull request #4234 from ProvableHQ/fix/validator_block_tree_dump

[Fix] Make the validators dump their block tree correctly

Author: kaimast

cli/src/commands/start.rs

@@ -64,7 +64,7 @@ use std::{
     sync::{Arc, atomic::AtomicBool},
 };
 use tokio::{
-    runtime::{self, Runtime},
+    runtime::{self, Handle, Runtime},
     sync::mpsc,
     task,
 };
@@ -323,7 +323,9 @@ impl Start {
         }
 
         // Initialize the runtime.
-        Self::runtime().block_on(async move {
+        let runtime = Self::runtime();
+        let handle = runtime.handle().clone();
+        runtime.block_on(async move {
             // Error messages.
             let node_parse_error = || "Failed to start node";
 
@@ -332,9 +334,15 @@ impl Start {
 
             // Parse the node arguments, start it, and block until shutdown.
             match self_.network {
-                MainnetV0::ID => self_.parse_node::<MainnetV0>(log_receiver).await.with_context(node_parse_error)?,
-                TestnetV0::ID => self_.parse_node::<TestnetV0>(log_receiver).await.with_context(node_parse_error)?,
-                CanaryV0::ID => self_.parse_node::<CanaryV0>(log_receiver).await.with_context(node_parse_error)?,
+                MainnetV0::ID => {
+                    self_.parse_node::<MainnetV0>(handle, log_receiver).await.with_context(node_parse_error)?
+                }
+                TestnetV0::ID => {
+                    self_.parse_node::<TestnetV0>(handle, log_receiver).await.with_context(node_parse_error)?
+                }
+                CanaryV0::ID => {
+                    self_.parse_node::<CanaryV0>(handle, log_receiver).await.with_context(node_parse_error)?
+                }
                 _ => panic!("Invalid network ID specified"),
             };
 
@@ -665,7 +673,7 @@ impl Start {
 
     /// Start the node and blocks until it terminates.
     #[rustfmt::skip]
-    async fn parse_node<N: Network>(&mut self, log_receiver: mpsc::Receiver<Vec<u8>>) -> Result<()> {
+    async fn parse_node<N: Network>(&mut self, handle: Handle, log_receiver: mpsc::Receiver<Vec<u8>>) -> Result<()> {
         if !self.nobanner {
             // Print the welcome banner.
             println!("{}", crate::helpers::welcome_message());
@@ -842,7 +850,7 @@ impl Start {
         }
 
         // Register the signal handler.
-        let signal_handler = SignalHandler::new();
+        let signal_handler = SignalHandler::new(Some(handle));
 
         // Initialize the node.
         let node = match node_type {

node/bft/ledger-service/src/ledger.rs

@@ -474,4 +474,8 @@ impl<N: Network, C: ConsensusStorage<N>> LedgerService<N> for CoreLedgerService<
             }
         }
     }
+
+    fn is_stopped(&self) -> bool {
+        self.stoppable.is_stopped()
+    }
 }

node/bft/ledger-service/src/mock.rs

@@ -235,4 +235,9 @@ impl<N: Network> LedgerService<N> for MockLedgerService<N> {
         let height = N::CONSENSUS_HEIGHT(consensus_version).unwrap();
         Ok(consensus_config_value!(N, TRANSACTION_SPEND_LIMIT, height).unwrap())
     }
+
+    fn is_stopped(&self) -> bool {
+        // No ledger to check against.
+        false
+    }
 }

node/bft/ledger-service/src/prover.rs

@@ -190,4 +190,9 @@ impl<N: Network> LedgerService<N> for ProverLedgerService<N> {
     ) -> Result<u64> {
         bail!("Cannot compute spent_cost in prover")
     }
+
+    fn is_stopped(&self) -> bool {
+        // No ledger to check against.
+        false
+    }
 }

node/bft/ledger-service/src/traits.rs

@@ -160,4 +160,6 @@ pub trait LedgerService<N: Network>: std::fmt::Debug + Send + Sync {
         transaction: &Transaction<N>,
         consensus_version: ConsensusVersion,
     ) -> Result<u64>;
+
+    fn is_stopped(&self) -> bool;
 }

node/bft/ledger-service/src/translucent.rs

@@ -200,4 +200,8 @@ impl<N: Network, C: ConsensusStorage<N>> LedgerService<N> for TranslucentLedgerS
     ) -> Result<u64> {
         self.inner.transaction_spend_in_microcredits(transaction, consensus_version)
     }
+
+    fn is_stopped(&self) -> bool {
+        self.inner.is_stopped()
+    }
 }

node/bft/src/primary.rs

@@ -2055,6 +2055,8 @@ impl<N: Network> Primary<N> {
         info!("Shutting down the primary...");
         // Remove the callback.
         self.primary_callback.clear();
+        // Shut down the sync service.
+        self.sync.shut_down().await;
         // Shut down the workers.
         self.workers().iter().for_each(|worker| worker.shut_down());
         // Abort the tasks.

node/bft/src/worker.rs

@@ -676,6 +676,7 @@ mod tests {
             fn check_block_subdag(&self, _block: Block<N>, _prefix: &[PendingBlock<N>]) -> Result<PendingBlock<N>, CheckBlockError<N>>;
             fn begin_ledger_update<'a>(&'a self) -> Result<Box<dyn LedgerUpdateService<N> + 'a>, BeginLedgerUpdateError>;
             fn transaction_spend_in_microcredits(&self, transaction: &Transaction<N>, consensus_version: ConsensusVersion) -> Result<u64>;
+            fn is_stopped(&self) -> bool;
         }
     }
 

node/rest/src/lib.rs

@@ -139,6 +139,11 @@ impl<N: Network, C: ConsensusStorage<N>, R: Routing<N>> Rest<N, C, R> {
     pub const fn handles(&self) -> &Arc<Mutex<Vec<JoinHandle<()>>>> {
         &self.handles
     }
+
+    /// Shuts down the REST instance.
+    pub fn shut_down(&self) {
+        self.handles.lock().iter().for_each(|handle| handle.abort());
+    }
 }
 
 impl<N: Network, C: ConsensusStorage<N>, R: Routing<N>> Rest<N, C, R> {

node/router/src/lib.rs

@@ -238,6 +238,11 @@ impl<N: Network> Router<N> {
         &self.cache
     }
 
+    /// Returns a reference to the ledger.
+    pub fn ledger(&self) -> &Arc<dyn LedgerService<N>> {
+        &self.ledger
+    }
+
     /// Returns `true` if the node is only engaging with trusted peers.
     pub fn trusted_peers_only(&self) -> bool {
         self.trusted_peers_only