Better graceful shutdown

Dockerfile

@@ -1,4 +1,4 @@
-FROM rust:1.67 as builder
+FROM rust:1.67-bookworm as builder
 
 WORKDIR /usr/src/app
 COPY . .
@@ -33,4 +33,4 @@ ENV VSS_PORT=8080
 EXPOSE $VSS_PORT
 
 # Run the app
-CMD ./vss-rs
+CMD ["./vss-rs"]

src/main.rs

@@ -8,7 +8,10 @@ use axum::{http, Extension, Router, TypedHeader};
 use diesel::r2d2::{ConnectionManager, Pool};
 use diesel::PgConnection;
 use diesel_migrations::MigrationHarness;
+use log::{error, info};
 use secp256k1::{All, PublicKey, Secp256k1};
+use tokio::signal::unix::{signal, SignalKind};
+use tokio::sync::oneshot;
 use tower_http::cors::{AllowOrigin, CorsLayer};
 
 mod auth;
@@ -133,21 +136,47 @@ async fn main() -> anyhow::Result<()> {
         .layer(DefaultBodyLimit::max(100_000_000)) // max 100mb body size
         .layer(Extension(state));
 
+    // Set up a oneshot channel to handle shutdown signal
+    let (tx, rx) = oneshot::channel();
+
+    // Spawn a task to listen for shutdown signals
+    tokio::spawn(async move {
+        let mut term_signal = signal(SignalKind::terminate())
+            .map_err(|e| error!("failed to install TERM signal handler: {e}"))
+            .unwrap();
+        let mut int_signal = signal(SignalKind::interrupt())
+            .map_err(|e| {
+                error!("failed to install INT signal handler: {e}");
+            })
+            .unwrap();
+
+        tokio::select! {
+            _ = term_signal.recv() => {
+                info!("Received SIGTERM");
+            },
+            _ = int_signal.recv() => {
+                info!("Received SIGINT");
+            },
+        }
+
+        let _ = tx.send(());
+    });
+
     let server = axum::Server::bind(&addr).serve(server_router.into_make_service());
 
-    println!("Webserver running on http://{addr}");
+    info!("Webserver running on http://{addr}");
 
     let graceful = server.with_graceful_shutdown(async {
-        tokio::signal::ctrl_c()
-            .await
-            .expect("failed to create Ctrl+C shutdown signal");
+        let _ = rx.await;
     });
 
     // Await the server to receive the shutdown signal
     if let Err(e) = graceful.await {
-        eprintln!("shutdown error: {e}");
+        error!("shutdown error: {e}");
     }
 
+    info!("Graceful shutdown complete");
+
     Ok(())
 }