From fb28a61043e6d7adfc45ffb8fe8d37cd4d8325c1 Mon Sep 17 00:00:00 2001
From: Noah Citron <noah@jeff.org>
Date: Fri, 23 Sep 2022 22:09:03 -0400
Subject: [PATCH] feat: add forced shutdowns (#56)

---
 cli/src/main.rs      | 47 +++++++++++++++++++++++++++++++++++++-------
 client/src/client.rs |  3 ---
 2 files changed, 40 insertions(+), 10 deletions(-)

diff --git a/cli/src/main.rs b/cli/src/main.rs
index 453fd72..079c442 100644
--- a/cli/src/main.rs
+++ b/cli/src/main.rs
@@ -1,4 +1,9 @@
-use std::{fs, path::PathBuf, process::exit};
+use std::{
+    fs,
+    path::PathBuf,
+    process::exit,
+    sync::{Arc, Mutex},
+};
 
 use clap::Parser;
 use common::utils::hex_str_to_bytes;
@@ -6,9 +11,10 @@ use dirs::home_dir;
 use env_logger::Env;
 use eyre::Result;
 
-use client::Client;
+use client::{database::FileDB, Client};
 use config::{networks, Config};
 use futures::executor::block_on;
+use log::info;
 
 #[tokio::main]
 async fn main() -> Result<()> {
@@ -19,14 +25,41 @@ async fn main() -> Result<()> {
 
     client.start().await?;
 
-    ctrlc::set_handler(move || {
-        block_on(client.shutdown());
-        exit(0);
-    })?;
-
+    register_shutdown_handler(client);
     std::future::pending().await
 }
 
+fn register_shutdown_handler(client: Client<FileDB>) {
+    let client = Arc::new(client);
+    let shutdown_counter = Arc::new(Mutex::new(0));
+
+    ctrlc::set_handler(move || {
+        let mut counter = shutdown_counter.lock().unwrap();
+        *counter += 1;
+
+        let counter_value = *counter;
+
+        if counter_value == 3 {
+            info!("forced shutdown");
+            exit(0);
+        }
+
+        info!(
+            "shutting down... press ctrl-c {} more times to force quit",
+            3 - counter_value
+        );
+
+        if counter_value == 1 {
+            let client = client.clone();
+            std::thread::spawn(move || {
+                block_on(client.shutdown());
+                exit(0);
+            });
+        }
+    })
+    .expect("could not register shutdown handler");
+}
+
 fn get_config() -> Config {
     let cli = Cli::parse();
     let mut config = match cli.network.as_str() {
diff --git a/client/src/client.rs b/client/src/client.rs
index 96e5dc5..18ce440 100644
--- a/client/src/client.rs
+++ b/client/src/client.rs
@@ -67,9 +67,6 @@ impl<DB: Database> Client<DB> {
     }
 
     pub async fn shutdown(&self) {
-        println!();
-        info!("shutting down");
-
         let node = self.node.read().await;
         let checkpoint = if let Some(checkpoint) = node.get_last_checkpoint() {
             checkpoint