feat: also gracefully shutdown on SIGTERM

Signed-off-by: Bugen Zhao <[email protected]>

src/utils/runtime/src/lib.rs

@@ -35,6 +35,7 @@ mod panic_hook;
 pub use panic_hook::*;
 mod prof;
 use prof::*;
+use tokio::signal::unix::SignalKind;
 
 /// Start RisingWave components with configs from environment variable.
 ///
@@ -96,24 +97,26 @@ where
         spawn_prof_thread(profile_path);
     }
 
+    let mut sigint = tokio::signal::unix::signal(SignalKind::interrupt()).unwrap();
+    let mut sigterm = tokio::signal::unix::signal(SignalKind::terminate()).unwrap();
+
     let future_with_shutdown = async move {
         let shutdown = CancellationToken::new();
         let mut fut = pin!(f(shutdown.clone()));
 
         tokio::select! {
             biased;
-            result = tokio::signal::ctrl_c() => {
-                result.expect("failed to receive ctrl-c signal");
-                tracing::info!("received ctrl-c, shutting down... (press ctrl-c again to force shutdown)");
 
-                // Send shutdown signal.
+            // Watch SIGINT, typically originating from user pressing ctrl-c.
+            // Attempt to shutdown gracefully and force shutdown on the next signal.
+            _ = sigint.recv() => {
+                tracing::info!("received ctrl-c, shutting down... (press ctrl-c again to force shutdown)");
                 shutdown.cancel();
 
                 // While waiting for the future to finish, listen for the second ctrl-c signal.
                 tokio::select! {
                     biased;
-                    result = tokio::signal::ctrl_c() => {
-                        result.expect("failed to receive ctrl-c signal");
+                    _ = sigint.recv() => {
                         tracing::warn!("forced shutdown");
 
                         // Directly exit the process **here** instead of returning from the future, since
@@ -123,6 +126,16 @@ where
                     _ = &mut fut => {},
                 }
             }
+
+            // Watch SIGTERM, typically originating from Kubernetes.
+            // Attempt to shutdown gracefully. No need to force shutdown since it will send SIGKILL after a timeout.
+            _ = sigterm.recv() => {
+                tracing::info!("received SIGTERM, shutting down...");
+                shutdown.cancel();
+                fut.await;
+            }
+
+            // Proceed with the future.
             _ = &mut fut => {},
         }
     };