wip: graceful shutdown

Author: jonaro00

Cargo.lock

@@ -56,6 +56,7 @@ hyper-util = { version = "0.1.10", features = ["full"] }
 ignore = "0.4.20"
 indicatif = "0.17.3"
 indoc = "2.0.1"
+nix = "0.30"
 percent-encoding = "2.2"
 portpicker = "0.1.1"
 pretty_assertions = "1.3.0"

Cargo.toml

@@ -40,6 +40,7 @@ hyper-util = { workspace = true }
 ignore = { workspace = true }
 indicatif = { workspace = true }
 indoc = { workspace = true }
+nix = { workspace = true, features = ["signal"] }
 portpicker = { workspace = true }
 regex = { workspace = true }
 reqwest = { workspace = true }

cargo-shuttle/Cargo.toml

@@ -1440,6 +1440,8 @@ impl Shuttle {
         .kill_on_drop(true)
         .spawn()
         .context("spawning runtime process")?;
+        #[allow(unused)]
+        let pid = runtime.id().context("getting runtime's PID")?;
 
         // Start background tasks for reading runtime's stdout and stderr
         let raw = run_args.raw;
@@ -1492,7 +1494,7 @@ impl Shuttle {
         });
 
         #[cfg(target_family = "unix")]
-        let exit_result = {
+        let (exit_result, interrupted) = {
             let mut sigterm_notif =
                 tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
                     .expect("Can not get the SIGTERM signal receptor");
@@ -1501,20 +1503,20 @@ impl Shuttle {
                     .expect("Can not get the SIGINT signal receptor");
             tokio::select! {
                 exit_result = runtime.wait() => {
-                    Some(exit_result)
+                    (Some(exit_result), false)
                 }
                 _ = sigterm_notif.recv() => {
-                    eprintln!("Received SIGTERM. Killing the runtime...");
-                    None
+                    eprintln!("Received SIGTERM. Shutting down the runtime...");
+                    (None, false)
                 },
                 _ = sigint_notif.recv() => {
-                    eprintln!("Received SIGINT. Killing the runtime...");
-                    None
+                    eprintln!("Received SIGINT. Shutting down the runtime...");
+                    (None, true)
                 }
             }
         };
         #[cfg(target_family = "windows")]
-        let exit_result = {
+        let (exit_result, interrupted) = {
             let mut ctrl_break_notif = tokio::signal::windows::ctrl_break()
                 .expect("Can not get the CtrlBreak signal receptor");
             let mut ctrl_c_notif =
@@ -1527,27 +1529,27 @@ impl Shuttle {
                 .expect("Can not get the CtrlShutdown signal receptor");
             tokio::select! {
                 exit_result = runtime.wait() => {
-                    Some(exit_result)
+                    (Some(exit_result), false)
                 }
                 _ = ctrl_break_notif.recv() => {
                     eprintln!("Received ctrl-break.");
-                    None
+                    (None, false)
                 },
                 _ = ctrl_c_notif.recv() => {
                     eprintln!("Received ctrl-c.");
-                    None
+                    (None, true)
                 },
                 _ = ctrl_close_notif.recv() => {
                     eprintln!("Received ctrl-close.");
-                    None
+                    (None, false)
                 },
                 _ = ctrl_logoff_notif.recv() => {
                     eprintln!("Received ctrl-logoff.");
-                    None
+                    (None, false)
                 },
                 _ = ctrl_shutdown_notif.recv() => {
                     eprintln!("Received ctrl-shutdown.");
-                    None
+                    (None, false)
                 }
             }
         };
@@ -1562,6 +1564,29 @@ impl Shuttle {
                 bail!("Failed to wait for runtime process to exit: {e}");
             }
             None => {
+                #[cfg(target_family = "unix")]
+                {
+                    if interrupted {
+                        nix::sys::signal::kill(
+                            nix::unistd::Pid::from_raw(pid as i32),
+                            nix::sys::signal::SIGINT,
+                        )
+                        .context("Sending SIGINT to runtime process")?;
+                        match tokio::time::timeout(Duration::from_secs(30), runtime.wait()).await {
+                            Ok(exit_result) => {
+                                debug!("Runtime exited {:?}", exit_result)
+                            }
+                            Err(_) => {
+                                eprintln!("Runtime shutdown timed out. Sending SIGKILL");
+                                runtime.kill().await?;
+                            }
+                        };
+                    } else {
+                        runtime.kill().await?;
+                    }
+                }
+
+                #[cfg(target_family = "windows")]
                 runtime.kill().await?;
             }
         }

cargo-shuttle/src/lib.rs

@@ -3,6 +3,7 @@ use std::{
     iter::FromIterator,
     net::{IpAddr, Ipv4Addr, SocketAddr},
     process::exit,
+    time::Duration,
 };
 
 use anyhow::Context;
@@ -253,16 +254,28 @@ pub async fn start(
     //
     info!("Starting service");
 
-    let service_bind = service.bind(service_addr);
-
     #[cfg(target_family = "unix")]
-    let interrupted = {
+    async fn shutdown_signal() {
         let mut sigterm_notif =
             tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
                 .expect("Can not get the SIGTERM signal receptor");
         let mut sigint_notif =
             tokio::signal::unix::signal(tokio::signal::unix::SignalKind::interrupt())
                 .expect("Can not get the SIGINT signal receptor");
+        tokio::select! {
+            _ = sigterm_notif.recv() => {
+                tracing::warn!("Runtime received SIGTERM. Shutting down...");
+            },
+            _ = sigint_notif.recv() => {
+                tracing::warn!("Runtime received SIGINT. Shutting down...");
+            }
+        }
+    }
+
+    let service_bind = service.clone().bind(service_addr);
+
+    #[cfg(target_family = "unix")]
+    let interrupted = {
         tokio::select! {
             res = service_bind => {
                 if let Err(e) = res {
@@ -272,12 +285,7 @@ pub async fn start(
                 tracing::warn!("Service terminated on its own. Shutting down the runtime...");
                 false
             }
-            _ = sigterm_notif.recv() => {
-                tracing::warn!("Received SIGTERM. Shutting down the runtime...");
-                true
-            },
-            _ = sigint_notif.recv() => {
-                tracing::warn!("Received SIGINT. Shutting down the runtime...");
+            _ = shutdown_signal() => {
                 true
             }
         }
@@ -306,19 +314,19 @@ pub async fn start(
             _ = ctrl_break_notif.recv() => {
                 tracing::warn!("Received ctrl-break. Shutting down the runtime...");
                 true
-            },
+            }
             _ = ctrl_c_notif.recv() => {
                 tracing::warn!("Received ctrl-c. Shutting down the runtime...");
                 true
-            },
+            }
             _ = ctrl_close_notif.recv() => {
                 tracing::warn!("Received ctrl-close. Shutting down the runtime...");
                 true
-            },
+            }
             _ = ctrl_logoff_notif.recv() => {
                 tracing::warn!("Received ctrl-logoff. Shutting down the runtime...");
                 true
-            },
+            }
             _ = ctrl_shutdown_notif.recv() => {
                 tracing::warn!("Received ctrl-shutdown. Shutting down the runtime...");
                 true
@@ -327,7 +335,19 @@ pub async fn start(
     };
 
     if interrupted {
-        return 10;
+        match tokio::time::timeout(Duration::from_secs(60), service.shutdown()).await {
+            Err(_) => {
+                tracing::error!("Service graceful shutdown timed out internally");
+                return 11;
+            }
+            Ok(Err(shutdown_err)) => {
+                tracing::error!("Service shutdown error: {shutdown_err}");
+                return 12;
+            }
+            _ => {
+                return 10;
+            }
+        };
     }
 
     0

runtime/src/rt.rs

@@ -107,10 +107,14 @@ impl<R: Serialize + DeserializeOwned + Send> IntoResource<R> for R {
 /// An `Into<Service>` implementor is what is returned in the `shuttle_runtime::main` macro
 /// in order to run it on the Shuttle servers.
 #[async_trait]
-pub trait Service: Send {
+pub trait Service: Send + Clone {
     /// This function is run exactly once on startup of a deployment.
     ///
     /// The passed [`SocketAddr`] receives proxied HTTP traffic from your Shuttle subdomain (or custom domain).
     /// Binding to the address is only relevant if this service is an HTTP server.
-    async fn bind(mut self, addr: SocketAddr) -> Result<(), error::Error>;
+    async fn bind(self, addr: SocketAddr) -> Result<(), error::Error>;
+
+    async fn shutdown(self) -> Result<(), CustomError> {
+        Ok(())
+    }
 }