add on_shutdown hook

src/app.rs

@@ -16,7 +16,7 @@ use axum::Router as AxumRouter;
 use crate::controller::channels::AppChannels;
 use crate::{
     bgworker::{self, Queue},
-    boot::{BootResult, ServeParams, StartMode},
+    boot::{shutdown_signal, BootResult, StartMode},
     cache::{self},
     config::{self, Config},
     controller::{
@@ -63,7 +63,7 @@ pub struct AppContext {
 /// the application's routing, worker connections, task registration, and
 /// database actions according to their specific requirements and use cases.
 #[async_trait]
-pub trait Hooks {
+pub trait Hooks: Send {
     /// Defines the composite app version
     #[must_use]
     fn app_version() -> String {
@@ -111,17 +111,23 @@ pub trait Hooks {
     ///
     /// # Returns
     /// A Result indicating success () or an error if the server fails to start.
-    async fn serve(app: AxumRouter, server_config: ServeParams) -> Result<()> {
+    async fn serve(app: AxumRouter, ctx: &AppContext) -> Result<()> {
         let listener = tokio::net::TcpListener::bind(&format!(
             "{}:{}",
-            server_config.binding, server_config.port
+            ctx.config.server.binding, ctx.config.server.port
         ))
         .await?;
 
+        let cloned_ctx = ctx.clone();
         axum::serve(
             listener,
             app.into_make_service_with_connect_info::<SocketAddr>(),
         )
+        .with_graceful_shutdown(async move {
+            shutdown_signal().await;
+            tracing::info!("shutting down...");
+            Self::on_shutdown(&cloned_ctx).await;
+        })
         .await?;
 
         Ok(())
@@ -208,6 +214,11 @@ pub trait Hooks {
     /// Seeds the database with initial data.
     #[cfg(feature = "with-db")]
     async fn seed(db: &DatabaseConnection, path: &Path) -> Result<()>;
+
+    /// Called when the application is shutting down.
+    /// This function allows users to perform any necessary cleanup or final
+    /// actions before the application stops completely.
+    async fn on_shutdown(_ctx: &AppContext) {}
 }
 
 /// An initializer.

src/boot.rs

@@ -6,6 +6,7 @@ use std::path::PathBuf;
 use axum::Router;
 #[cfg(feature = "with-db")]
 use sea_orm_migration::MigratorTrait;
+use tokio::signal;
 use tracing::{debug, error, info, warn};
 
 #[cfg(feature = "with-db")]
@@ -77,27 +78,26 @@ pub async fn start<H: Hooks>(boot: BootResult, server_config: ServeParams) -> Re
 
     match (router, run_worker) {
         (Some(router), false) => {
-            H::serve(router, server_config).await?;
+            H::serve(router, &app_context).await?;
         }
         (Some(router), true) => {
-            tokio::spawn(async move {
-                debug!("note: worker is run in-process (tokio spawn)");
-                if let Some(queue) = &app_context.queue_provider {
-                    let res = queue.run().await;
+            debug!("note: worker is run in-process (tokio spawn)");
+            if let Some(queue) = &app_context.queue_provider {
+                let cloned_queue = queue.clone();
+                tokio::spawn(async move {
+                    let res = cloned_queue.run().await;
                     if res.is_err() {
                         error!(
                             err = res.unwrap_err().to_string(),
                             "error while running worker"
                         );
                     }
-                } else {
-                    error!(
-                        err = Error::QueueProviderMissing.to_string(),
-                        "cannot start worker"
-                    );
-                }
-            });
-            H::serve(router, server_config).await?;
+                });
+            } else {
+                return Err(Error::QueueProviderMissing);
+            }
+
+            H::serve(router, &app_context).await?;
         }
         (None, true) => {
             if let Some(queue) = &app_context.queue_provider {
@@ -383,6 +383,36 @@ pub fn list_endpoints<H: Hooks>(ctx: &AppContext) -> Vec<ListRoutes> {
     H::routes(ctx).collect()
 }
 
+/// Waits for a shutdown signal, either via Ctrl+C or termination signal.
+///
+/// # Panics
+///
+/// This function will panic if it fails to install the signal handlers for
+/// Ctrl+C or the terminate signal on Unix-based systems.
+pub 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 => {},
+    }
+}
+
 pub struct MiddlewareInfo {
     pub id: String,
     pub enabled: bool,

src/tests_cfg/config.rs

@@ -19,7 +19,7 @@ pub fn test_config() -> Config {
         },
         server: config::Server {
             binding: "localhost".to_string(),
-            port: 3000,
+            port: 5555,
             host: "localhost".to_string(),
             ident: None,
             middlewares: middleware::Config::default(),