Merge pull request #1 from hotdata-dev/feat/graceful-shutdown

feat(server): add graceful shutdown handling

Author: zfarrell

src/bin/server.rs

@@ -40,6 +40,7 @@ async fn main() -> Result<()> {
 
     // Create router
     let app = AppServer::new(engine);
+    let engine = app.engine.clone();
 
     // Create server address
     let addr = format!("{}:{}", config.server.host, config.server.port);
@@ -49,7 +50,43 @@ async fn main() -> Result<()> {
     tracing::info!("Server listening on {}", addr);
 
     // Start server
-    axum::serve(listener, app.router).await?;
+    let server = axum::serve(listener, app.router)
+        .with_graceful_shutdown(shutdown());
+
+    server.await?;
+
+    // Explicitly shutdown engine to close catalog connection
+    if let Err(e) = engine.shutdown() {
+        tracing::error!("Error during engine shutdown: {}", e);
+    }
+
+    tracing::info!("Server shutdown complete");
 
     Ok(())
 }
+
+async fn shutdown() {
+    let ctrl_c = async {
+        tokio::signal::ctrl_c()
+            .await
+            .expect("failed to install Ctrl+C handler");
+    };
+
+    #[cfg(unix)]
+    let terminate = async {
+        tokio::signal::unix::signal(tokio::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 => {},
+    }
+
+    tracing::info!("Shutdown signal received, stopping server...");
+}

src/http/app_server.rs

@@ -9,6 +9,7 @@ use std::sync::Arc;
 
 pub struct AppServer {
     pub router: Router,
+    pub engine: Arc<HotDataEngine>,
 }
 
 pub const PATH_QUERY: &str = "/query";
@@ -18,6 +19,7 @@ pub const PATH_CONNECTIONS: &str = "/connections";
 
 impl AppServer {
     pub fn new(engine: HotDataEngine) -> Self {
+        let engine = Arc::new(engine);
         AppServer {
             router: Router::new()
                 .route(PATH_QUERY, post(query_handler))
@@ -27,7 +29,8 @@ impl AppServer {
                     PATH_CONNECTIONS,
                     post(create_connection_handler).get(list_connections_handler),
                 )
-                .with_state(Arc::new(engine)),
+                .with_state(engine.clone()),
+            engine,
         }
     }
 }