fix: handle signals earlier

Signed-off-by: ljedrz <[email protected]>

node/Cargo.toml

@@ -44,6 +44,9 @@ version = "2.0"
 [dependencies.num_cpus]
 version = "1"
 
+[dependencies.once_cell]
+version = "1"
+
 [dependencies.parking_lot]
 version = "0.12"
 

node/src/beacon/mod.rs

@@ -97,6 +97,9 @@ impl<N: Network, C: ConsensusStorage<N>> Beacon<N, C> {
     ) -> Result<Self> {
         let timer = timer!("Beacon::new");
 
+        // Initialize the signal handler.
+        let signal_node = Self::handle_signals();
+
         // Initialize the ledger.
         let ledger = Ledger::load(genesis, dev)?;
         lap!(timer, "Initialize the ledger");
@@ -155,8 +158,8 @@ impl<N: Network, C: ConsensusStorage<N>> Beacon<N, C> {
         node.initialize_routing().await;
         // Initialize the block production.
         node.initialize_block_production().await;
-        // Initialize the signal handler.
-        node.handle_signals();
+        // Pass the node to the signal handler.
+        let _ = signal_node.set(node.clone());
         lap!(timer, "Initialize the handlers");
 
         finish!(timer);

node/src/client/mod.rs

@@ -60,6 +60,9 @@ impl<N: Network, C: ConsensusStorage<N>> Client<N, C> {
         genesis: Block<N>,
         dev: Option<u16>,
     ) -> Result<Self> {
+        // Initialize the signal handler.
+        let signal_node = Self::handle_signals();
+
         // Initialize the node router.
         let router = Router::new(
             node_ip,
@@ -83,8 +86,8 @@ impl<N: Network, C: ConsensusStorage<N>> Client<N, C> {
         };
         // Initialize the routing.
         node.initialize_routing().await;
-        // Initialize the signal handler.
-        node.handle_signals();
+        // Pass the node to the signal handler.
+        let _ = signal_node.set(node.clone());
         // Return the node.
         Ok(node)
     }

node/src/prover/mod.rs

@@ -77,6 +77,9 @@ impl<N: Network, C: ConsensusStorage<N>> Prover<N, C> {
         genesis: Block<N>,
         dev: Option<u16>,
     ) -> Result<Self> {
+        // Initialize the signal handler.
+        let signal_node = Self::handle_signals();
+
         // Initialize the node router.
         let router = Router::new(
             node_ip,
@@ -108,8 +111,8 @@ impl<N: Network, C: ConsensusStorage<N>> Prover<N, C> {
         node.initialize_routing().await;
         // Initialize the coinbase puzzle.
         node.initialize_coinbase_puzzle().await;
-        // Initialize the signal handler.
-        node.handle_signals();
+        // Pass the node to the signal handler.
+        let _ = signal_node.set(node.clone());
         // Return the node.
         Ok(node)
     }

node/src/traits.rs

@@ -16,6 +16,9 @@ use snarkos_node_messages::NodeType;
 use snarkos_node_router::Routing;
 use snarkvm::prelude::{Address, Network, PrivateKey, ViewKey};
 
+use once_cell::sync::OnceCell;
+use std::sync::Arc;
+
 #[async_trait]
 pub trait NodeInterface<N: Network>: Routing<N> {
     /// Returns the node type.
@@ -45,17 +48,25 @@ pub trait NodeInterface<N: Network>: Routing<N> {
 
     /// Handles OS signals for the node to intercept and perform a clean shutdown.
     /// Note: Only Ctrl-C is supported; it should work on both Unix-family systems and Windows.
-    fn handle_signals(&self) {
-        let node = self.clone();
+    fn handle_signals() -> Arc<OnceCell<Self>> {
+        // In order for the signal handler to be started as early as possible, a reference to the node needs
+        // to be passed to it at a later time.
+        let node: Arc<OnceCell<Self>> = Default::default();
+
+        let node_clone = node.clone();
         tokio::task::spawn(async move {
             match tokio::signal::ctrl_c().await {
                 Ok(()) => {
-                    node.shut_down().await;
+                    if let Some(node) = node_clone.get() {
+                        node.shut_down().await;
+                    }
                     std::process::exit(0);
                 }
                 Err(error) => error!("tokio::signal::ctrl_c encountered an error: {}", error),
             }
         });
+
+        node
     }
 
     /// Shuts down the node.

node/src/validator/mod.rs

@@ -72,6 +72,9 @@ impl<N: Network, C: ConsensusStorage<N>> Validator<N, C> {
         cdn: Option<String>,
         dev: Option<u16>,
     ) -> Result<Self> {
+        // Initialize the signal handler.
+        let signal_node = Self::handle_signals();
+
         // Initialize the ledger.
         let ledger = Ledger::load(genesis, dev)?;
         // Initialize the CDN.
@@ -114,8 +117,8 @@ impl<N: Network, C: ConsensusStorage<N>> Validator<N, C> {
         node.initialize_sync()?;
         // Initialize the routing.
         node.initialize_routing().await;
-        // Initialize the signal handler.
-        node.handle_signals();
+        // Pass the node to the signal handler.
+        let _ = signal_node.set(node.clone());
         // Return the node.
         Ok(node)
     }