Implement graceful shutdown with tokio signal handling

Cargo.toml

@@ -7,7 +7,7 @@ edition = "2021"
 axum = "0.7.3"
 serde = { version = "1.0.195", features = ["serde_derive"] }
 serde_repr = "0.1.18"
-tokio = "1.35.1"
+tokio = { version = "1.35.1", features = ["full"] }
 tower-http = { version = "0.5.0", features = ["trace"] }
 tracing = "0.1.40"
 tracing-subscriber = "0.3.18"

src/main.rs

@@ -10,6 +10,7 @@ use oracle::*;
 use serde::{Deserialize, Serialize};
 use stark_vrf::{generate_public_key, BaseField, ScalarValue, StarkVRF};
 use std::str::FromStr;
+use tokio::signal;
 use tower_http::trace::TraceLayer;
 use tracing::debug;
 
@@ -86,14 +87,48 @@ async fn main() {
         .with_max_level(tracing::Level::DEBUG)
         .init();
 
+    async fn index() -> &'static str {
+        "OK"
+    }
+
     let app = Router::new()
+        .route("/", get(index))
         .route("/info", get(vrf_info))
         .route("/stark_vrf", post(stark_vrf))
         .layer(TraceLayer::new_for_http());
 
     let listener = tokio::net::TcpListener::bind("0.0.0.0:3000")
         .await
         .expect("Failed to bind to port 3000, port already in use by another process. Change the port or terminate the other process.");
+
     debug!("Server started on http://0.0.0.0:3000");
-    axum::serve(listener, app).await.unwrap();
+
+    axum::serve(listener, app)
+        .with_graceful_shutdown(shutdown_signal())
+        .await
+        .unwrap();
+}
+
+async fn shutdown_signal() {
+    let ctrl_c = async {
+        signal::ctrl_c()
+            .await
+            .expect("failed to install Ctrl+C handler");
+    };
+
+    #[cfg(unix)]
+    let terminate = async {
+        signal::unix::signal(signal::unix::SignalKind::terminate())
+            .expect("failed to install signal handler")
+            .recv()
+            .await;
+    };
+
+    #[cfg(not(unix))]
+    let terminate = std::future::pending::<()>();
+
+    tokio::select! {
+        _ = ctrl_c => {},
+        _ = terminate => {},
+    }
 }