feat: Add graceful shutdown signal handling

Cargo.toml

@@ -56,7 +56,7 @@ emulator-ui = { path = "crates/emulator-ui" }
 hybrid-axum-tonic = { path = "crates/hybrid-axum-tonic" }
 tikv-jemallocator = "0.5.4"
 time = { workspace = true, optional = true }
-tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
+tokio = { workspace = true, features = ["macros", "rt-multi-thread", "signal"] }
 tracing = { workspace = true }
 tracing-subscriber = { version = "0.3.18", optional = true }
 

src/lib.rs

@@ -1,22 +1,34 @@
-use std::net::SocketAddr;
+use std::{future::Future, net::SocketAddr};
 
 use axum::Router;
 use hybrid_axum_tonic::{NestTonic, RestGrpcService};
-use tracing::info;
+use tracing::{enabled, info, Level};
 
 #[tokio::main]
-pub async fn run(host_port: SocketAddr) -> color_eyre::Result<()> {
+pub async fn run(
+    host_port: SocketAddr,
+    shutdown: impl Future<Output = ()>,
+) -> color_eyre::Result<()> {
     let rest_router = emulator_ui::router();
     let grpc_router = Router::new().nest_tonic(emulator_grpc::service());
     let combined = RestGrpcService::new(rest_router, grpc_router).into_make_service();
-    let server = axum::Server::bind(&host_port).serve(combined);
+    let server = axum::Server::bind(&host_port)
+        .serve(combined)
+        .with_graceful_shutdown(shutdown);
 
-    info!("Firestore listening on {}", host_port);
-
-    #[cfg(not(feature = "tracing"))]
-    eprintln!("Firestore listening on {}", host_port);
+    if enabled!(Level::INFO) {
+        info!("Firestore emulator listening on {}", host_port);
+    } else {
+        eprintln!("Firestore emulator listening on {}", host_port);
+    }
 
     server.await?;
 
+    if enabled!(Level::INFO) {
+        info!("Firestore emulator stopped");
+    } else {
+        eprintln!("Firestore emulator stopped");
+    }
+
     Ok(())
 }

src/main.rs

@@ -3,6 +3,7 @@ use std::net::SocketAddr;
 use clap::Parser;
 use firestore_emulator::run;
 use tikv_jemallocator::Jemalloc;
+use tokio::signal::ctrl_c;
 
 #[global_allocator]
 static GLOBAL_ALLOC: Jemalloc = Jemalloc;
@@ -44,5 +45,7 @@ fn main() -> color_eyre::Result<()> {
         registry.init();
     }
 
-    run(Args::parse().host_port)
+    let ctrl_c_listener = async { ctrl_c().await.expect("failed to listen for ctrl-c event") };
+
+    run(Args::parse().host_port, ctrl_c_listener)
 }