Enable force shutdown with 2nd Ctrl+C

quickwit/quickwit-cli/src/service.rs

@@ -18,9 +18,11 @@
 // along with this program. If not, see <http://www.gnu.org/licenses/>.
 
 use std::collections::HashSet;
+use std::pin::pin;
 use std::str::FromStr;
 
 use clap::{arg, ArgAction, ArgMatches, Command};
+use colored::Colorize;
 use futures::future::select;
 use itertools::Itertools;
 use quickwit_common::runtimes::RuntimesConfig;
@@ -33,6 +35,7 @@ use quickwit_telemetry::payload::{QuickwitFeature, QuickwitTelemetryInfo, Teleme
 use tokio::signal;
 use tracing::{debug, info};
 
+use crate::checklist::{BLUE_COLOR, RED_COLOR};
 use crate::{config_cli_arg, get_resolvers, load_node_config, start_actor_runtimes};
 
 pub fn build_run_command() -> Command {
@@ -53,6 +56,38 @@ pub struct RunCliCommand {
     pub services: Option<HashSet<QuickwitService>>,
 }
 
+async fn listen_interrupt() {
+    async fn ctrl_c() {
+        signal::ctrl_c()
+            .await
+            .expect("registering a signal handler for SIGINT should not fail");
+        // carriage return to hide the ^C echo from the terminal
+        print!("\r");
+    }
+    ctrl_c().await;
+    println!(
+        "{} Graceful shutdown initiated. Waiting for ingested data to be indexed. This may take a \
+         few minutes. Press Ctrl+C again to force shutdown.",
+        "❢".color(BLUE_COLOR)
+    );
+    tokio::spawn(async {
+        ctrl_c().await;
+        println!(
+            "{} Quickwit was forcefully shut down. Some data might not have been indexed.",
+            "✘".color(RED_COLOR)
+        );
+        std::process::exit(1);
+    });
+}
+
+async fn listen_sigterm() {
+    signal::unix::signal(signal::unix::SignalKind::terminate())
+        .expect("registering a signal handler for SIGTERM should not fail")
+        .recv()
+        .await;
+    info!("SIGTERM received");
+}
+
 impl RunCliCommand {
     pub fn parse_cli_args(mut matches: ArgMatches) -> anyhow::Result<Self> {
         let config_uri = matches
@@ -95,20 +130,7 @@ impl RunCliCommand {
         let runtimes_config = RuntimesConfig::default();
         start_actor_runtimes(runtimes_config, &node_config.enabled_services)?;
         let shutdown_signal = Box::pin(async {
-            select(
-                Box::pin(async {
-                    signal::ctrl_c()
-                        .await
-                        .expect("registering a signal handler for SIGINT should not fail");
-                }),
-                Box::pin(async {
-                    signal::unix::signal(signal::unix::SignalKind::terminate())
-                        .expect("registering a signal handler for SIGTERM should not fail")
-                        .recv()
-                        .await;
-                }),
-            )
-            .await;
+            select(pin!(listen_interrupt()), pin!(listen_sigterm())).await;
         });
         let serve_result = serve_quickwit(
             node_config,
@@ -129,6 +151,7 @@ impl RunCliCommand {
             telemetry_handle.terminate_telemetry().await;
         }
         serve_result?;
+        info!("quickwit successfully terminated");
         Ok(())
     }
 }