install signal handlers early & handle SIGTERM

1 files changed, 52 insertions, 17 deletions

diff --git a/src/main.rs b/src/main.rs
index 2b90bef..d9d24f3 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -91,6 +91,10 @@ async fn main() -> Result<()> {
 
     let args = Args::parse();
 
+    // First thing: until a handler is installed, SIGINT and SIGTERM kill the process on
+    // the spot, and startup (spawning Mycelium, bootstrap) can take minutes.
+    let stop = watch_for_shutdown().context("Installing signal handlers")?;
+
     let mycelium_setup = {
         let mut spawn_config = MyceliumSpawnConfig {
             rest_url_opt: args.mycelium_rest_url,
@@ -102,7 +106,10 @@ async fn main() -> Result<()> {
         if let Some(net_name) = args.mycelium_network_name {
             spawn_config.network_name = net_name;
         }
-        MyceliumSetup::bootstrap(spawn_config).await?
+        tokio::select! {
+            setup = MyceliumSetup::bootstrap(spawn_config) => setup?,
+            _ = stopped(stop.clone()) => return Ok(()),
+        }
     };
 
     let mycelium_rest_url = mycelium_setup.rest_url.clone();
@@ -192,9 +199,13 @@ async fn main() -> Result<()> {
     tokio::spawn(remake::run_worker(state.clone()));
     trace!("Remake worker task spawned");
 
-    bootstrap::bootstrap(&state, &seeds)
-        .await
-        .context("Joining cluster")?;
+    tokio::select! {
+        joined = bootstrap::bootstrap(&state, &seeds) => joined.context("Joining cluster")?,
+        _ = stopped(stop.clone()) => {
+            info!("Stopped while joining the cluster");
+            return Ok(());
+        }
+    }
 
     // Start the change publisher after bootstrapping: streams only exist from then on,
     // and it avoids diffing the initial sync.
@@ -226,9 +237,6 @@ async fn main() -> Result<()> {
     let server_fut = server_builder.into_future();
     tokio::pin!(server_fut);
 
-    let mut terminate = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
-        .context("Installing the SIGTERM handler")?;
-
     tokio::select! {
         exit_status = state.transport.wait_for_exit() => {
             error!("Mycelium runtime process exited. Status: {:?}", exit_status);
@@ -245,16 +253,7 @@ async fn main() -> Result<()> {
             std::mem::drop(state);
         }
 
-        _ = tokio::signal::ctrl_c() => {
-            info!("Received manual interrupt signal. Shutting down...");
-            let _ = shutdown_tx.send(());
-            let _ = server_fut.as_mut().await;
-        }
-
-        // What service managers send; handled like Ctrl-C so the supervised Mycelium is
-        // stopped on the way out instead of relying on the kernel's parent-death signal.
-        _ = terminate.recv() => {
-            info!("Received termination signal. Shutting down...");
+        _ = stopped(stop.clone()) => {
             let _ = shutdown_tx.send(());
             let _ = server_fut.as_mut().await;
         }
@@ -264,5 +263,41 @@ async fn main() -> Result<()> {
     Ok(())
 }
 
+/// Turns SIGINT and SIGTERM into a shutdown request. A second signal exits at once (the
+/// supervised Mycelium still gets its parent-death signal then).
+fn watch_for_shutdown() -> Result<watch::Receiver<bool>> {
+    use tokio::signal::unix::{signal, SignalKind};
+
+    let mut interrupt = signal(SignalKind::interrupt())?;
+    let mut terminate = signal(SignalKind::terminate())?;
+    let (stop_tx, stop_rx) = watch::channel(false);
+
+    tokio::spawn(async move {
+        let name = tokio::select! {
+            _ = interrupt.recv() => "SIGINT",
+            _ = terminate.recv() => "SIGTERM",
+        };
+        info!(
+            signal = name,
+            "Shutting down (send it again to exit immediately)"
+        );
+        let _ = stop_tx.send(true);
+
+        tokio::select! {
+            _ = interrupt.recv() => {}
+            _ = terminate.recv() => {}
+        }
+        error!("Second signal; exiting immediately");
+        std::process::exit(130);
+    });
+
+    Ok(stop_rx)
+}
+
+/// Resolves once a shutdown was requested.
+async fn stopped(mut stop: watch::Receiver<bool>) {
+    let _ = stop.wait_for(|stop| *stop).await;
+}
+
 // === Footnotes ===============================================================
 // [concurrent]: these comments are in place so we remember what to do if we go from async to concurrent.