From 3633f25d0c4332873cad277f9090ae45052b4148 Mon Sep 17 00:00:00 2001
From: discord9 <55937128+discord9@users.noreply.github.com>
Date: Tue, 19 Nov 2024 23:20:33 +0800
Subject: [PATCH] feat: also shutdown gracefully on sigterm on unix (#5023)

* feat: support SIGTERM on unix

* chore: log

* fix: Result type
---
 src/cmd/src/lib.rs | 31 ++++++++++++++++++++++++++++---
 1 file changed, 28 insertions(+), 3 deletions(-)

diff --git a/src/cmd/src/lib.rs b/src/cmd/src/lib.rs
index 731f527daf7c..3a719b9589a7 100644
--- a/src/cmd/src/lib.rs
+++ b/src/cmd/src/lib.rs
@@ -43,6 +43,31 @@ lazy_static::lazy_static! {
         prometheus::register_int_gauge_vec!("greptime_app_version", "app version", &["version", "short_version", "app"]).unwrap();
 }
 
+/// wait for the close signal, for unix platform it's SIGINT or SIGTERM
+#[cfg(unix)]
+async fn start_wait_for_close_signal() -> std::io::Result<()> {
+    use tokio::signal::unix::{signal, SignalKind};
+    let mut sigint = signal(SignalKind::interrupt())?;
+    let mut sigterm = signal(SignalKind::terminate())?;
+
+    tokio::select! {
+        _ = sigint.recv() => {
+            info!("Received SIGINT, shutting down");
+        }
+        _ = sigterm.recv() => {
+            info!("Received SIGTERM, shutting down");
+        }
+    }
+
+    Ok(())
+}
+
+/// wait for the close signal, for non-unix platform it's ctrl-c
+#[cfg(not(unix))]
+async fn start_wait_for_close_signal() -> std::io::Result<()> {
+    tokio::signal::ctrl_c().await
+}
+
 #[async_trait]
 pub trait App: Send {
     fn name(&self) -> &str;
@@ -69,9 +94,9 @@ pub trait App: Send {
         self.start().await?;
 
         if self.wait_signal() {
-            if let Err(e) = tokio::signal::ctrl_c().await {
-                error!(e; "Failed to listen for ctrl-c signal");
-                // It's unusual to fail to listen for ctrl-c signal, maybe there's something unexpected in
+            if let Err(e) = start_wait_for_close_signal().await {
+                error!(e; "Failed to listen for close signal");
+                // It's unusual to fail to listen for close signal, maybe there's something unexpected in
                 // the underlying system. So we stop the app instead of running nonetheless to let people
                 // investigate the issue.
             }