From 5fe5fdffda179b400a738026afae299156ec6218 Mon Sep 17 00:00:00 2001
From: ljedrz <ljedrz@gmail.com>
Date: Tue, 20 Jun 2023 13:20:57 +0200
Subject: [PATCH] fix: handle signals earlier

Signed-off-by: ljedrz <ljedrz@gmail.com>
---
 node/Cargo.toml           |  3 +++
 node/src/beacon/mod.rs    |  7 +++++--
 node/src/client/mod.rs    |  7 +++++--
 node/src/prover/mod.rs    |  7 +++++--
 node/src/traits.rs        | 17 ++++++++++++++---
 node/src/validator/mod.rs |  7 +++++--
 6 files changed, 37 insertions(+), 11 deletions(-)

diff --git a/node/Cargo.toml b/node/Cargo.toml
index 920ecf1ce0..d46a51a207 100644
--- a/node/Cargo.toml
+++ b/node/Cargo.toml
@@ -44,6 +44,9 @@ version = "1"
 [dependencies.num_cpus]
 version = "1"
 
+[dependencies.once_cell]
+version = "1"
+
 [dependencies.parking_lot]
 version = "0.12"
 
diff --git a/node/src/beacon/mod.rs b/node/src/beacon/mod.rs
index 4752c30fb6..24a64ff462 100644
--- a/node/src/beacon/mod.rs
+++ b/node/src/beacon/mod.rs
@@ -98,6 +98,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");
@@ -156,8 +159,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);
diff --git a/node/src/client/mod.rs b/node/src/client/mod.rs
index 01623ad4b6..a8efb950b6 100644
--- a/node/src/client/mod.rs
+++ b/node/src/client/mod.rs
@@ -55,6 +55,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,
@@ -78,8 +81,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)
     }
diff --git a/node/src/prover/mod.rs b/node/src/prover/mod.rs
index e78e20b504..a44f14e5a4 100644
--- a/node/src/prover/mod.rs
+++ b/node/src/prover/mod.rs
@@ -72,6 +72,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,
@@ -103,8 +106,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)
     }
diff --git a/node/src/traits.rs b/node/src/traits.rs
index 173bafffe4..7675f45295 100644
--- a/node/src/traits.rs
+++ b/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.
diff --git a/node/src/validator/mod.rs b/node/src/validator/mod.rs
index 7a21078a4c..6e65accaa5 100644
--- a/node/src/validator/mod.rs
+++ b/node/src/validator/mod.rs
@@ -66,6 +66,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.
@@ -108,8 +111,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)
     }