Fixed graceful shutdown (#414)

Author: taco-paco

test-bins/src/rpc.rs

@@ -1,3 +1,5 @@
+mod shutdown;
+
 use log::*;
 use magicblock_api::{
     ledger,
@@ -8,6 +10,8 @@ use magicblock_config::{EphemeralConfig, GeyserGrpcConfig};
 use solana_sdk::signature::{Keypair, Signer};
 use test_tools::init_logger;
 
+use crate::shutdown::Shutdown;
+
 // mAGicPQYBMvcYveUZA5F5UNNwyHvfYh5xkLS2Fr1mev
 const TEST_KEYPAIR_BYTES: [u8; 64] = [
     7, 83, 184, 55, 200, 223, 238, 137, 166, 244, 107, 126, 189, 16, 194, 36,
@@ -98,7 +102,6 @@ async fn main() {
         ledger::lock_ledger(api.ledger().ledger_path(), &mut ledger_lock);
 
     api.start().await.expect("Failed to start validator");
-
     info!("");
     info!("🧙 Magicblock Validator is running!");
     info!(
@@ -113,11 +116,9 @@ async fn main() {
     info!("Ready for connections!");
     info!("");
 
-    // validator is supposed to run forever, so we wait for
-    // termination signal to initiate a graceful shutdown
-    let _ = tokio::signal::ctrl_c().await;
-
-    info!("SIGTERM has been received, initiating graceful shutdown");
+    if let Err(err) = Shutdown::wait().await {
+        error!("Failed to gracefully shutdown: {}", err);
+    }
     // weird panic behavior in json rpc http server, which panics when stopped from
     // within async context, so we just move it to a different thread for shutdown
     //

test-bins/src/shutdown.rs

@@ -0,0 +1,35 @@
+use std::io;
+
+use log::info;
+use tokio::{signal, signal::unix::SignalKind};
+
+pub struct Shutdown;
+impl Shutdown {
+    pub async fn wait() -> Result<(), io::Error> {
+        #[cfg(unix)]
+        return Self::wait_unix().await;
+        #[cfg(not(unix))]
+        return Self::wait_other().await;
+    }
+
+    #[cfg(unix)]
+    async fn wait_unix() -> Result<(), io::Error> {
+        let mut terminate_signal =
+            signal::unix::signal(SignalKind::terminate())?;
+        tokio::select! {
+            _ = terminate_signal.recv() => {
+                info!("SIGTERM has been received, initiating graceful shutdown");
+            },
+            _ = tokio::signal::ctrl_c() => {
+                info!("SIGINT signal received, initiating graceful shutdown");
+            },
+        }
+
+        Ok(())
+    }
+
+    #[cfg(not(unix))]
+    async fn wait_other() -> Result<(), io::Error> {
+        tokio::signal::ctrl_c().await
+    }
+}