75 gracefully shutdown server

generic-server/Cargo.toml

@@ -25,8 +25,10 @@ server-plugin = { path = "../server-plugin" }
 chrono = { version = "0.4.26", optional = true }
 did-endpoint = { path = "../did-endpoint", optional = true }
 oob-messages = { path = "../oob-messages", optional = true }
+once_cell = "1.19.0"
 
 [dev-dependencies]
+nix = { version = "0.29.0", features = ["feature"] }
 tower = { version = "0.4.13", features = ["util"] }
 
 [features]

generic-server/src/lib.rs

@@ -1,18 +1,21 @@
 pub mod plugin;
 pub mod util;
-
 use plugin::container::PluginContainer;
 
 use axum::Router;
 use tower_http::catch_panic::CatchPanicLayer;
 use tower_http::trace::TraceLayer;
 
-pub fn app() -> Router {
+pub fn app() -> (PluginContainer<'static>, Router) {
     let mut container = PluginContainer::default();
     let _ = container.load();
-
-    Router::new() //
-        .merge(container.routes().unwrap_or_default())
-        .layer(TraceLayer::new_for_http())
-        .layer(CatchPanicLayer::new())
+
+    let router = Router::new() 
+            .merge(container.routes().unwrap_or_default())
+            .layer(TraceLayer::new_for_http())
+            .layer(CatchPanicLayer::new());
+    (container, router)
 }
+
+
+

generic-server/src/main.rs

@@ -1,28 +1,43 @@
-use generic_server::app;
-
 use axum::Server;
+use generic_server::app;
 use std::net::SocketAddr;
 
 #[tokio::main]
 async fn main() {
+    // Start server
+    let port = std::env::var("SERVER_LOCAL_PORT").unwrap_or("3000".to_owned());
+    let addr: SocketAddr = format!("127.0.0.1:{port}").parse().unwrap();
+    tracing::info!("listening on {addr}");
+    generic_server_with_gracefull_shutdown(addr).await
+}
+
+async fn generic_server_with_gracefull_shutdown(addr: SocketAddr) {
     // Load dotenv-flow variables
     dotenv_flow::dotenv_flow().ok();
 
     // Enable logging
     config_tracing();
 
-    // Start server
-    let port = std::env::var("SERVER_LOCAL_PORT").unwrap_or("3000".to_owned());
-    let addr: SocketAddr = format!("127.0.0.1:{port}").parse().unwrap();
-    tracing::info!("listening on {addr}");
-    Server::bind(&addr)
-        .serve(app().into_make_service())
-        .await
-        .unwrap();
-}
+    // Load plugins
+    let (mut plugin_container, router) = app();
 
+    // spawn task for server
+    tokio::spawn(async move {
+        Server::bind(&addr)
+            .serve(router.into_make_service())
+            .await
+            .unwrap();
+    });
+
+    tokio::select! {
+        _ = tokio::signal::ctrl_c() => {
+           tracing::info!("\nshuting down gracefully");
+            let _ = plugin_container.unload();}
+    };
+}
 fn config_tracing() {
     use tracing::Level;
+
     use tracing_subscriber::{filter, layer::SubscriberExt, util::SubscriberInitExt};
 
     let tracing_layer = tracing_subscriber::fmt::layer();
@@ -36,3 +51,20 @@ fn config_tracing() {
         .with(filter)
         .init();
 }
+
+#[cfg(test)]
+mod tests {
+    use super::generic_server_with_gracefull_shutdown;
+    use std::net::SocketAddr;
+
+    #[tokio::test]
+    async fn test_server_shutdown() {
+        let port = std::env::var("SERVER_LOCAL_PORT").unwrap_or("3000".to_owned());
+        let addr: SocketAddr = format!("127.0.0.1:{port}").parse().unwrap();
+        tracing::info!("listening on {addr}");
+
+        // run server in background
+        tokio::spawn(generic_server_with_gracefull_shutdown(addr));
+        // send a shutdown signal to main thread
+    }
+}

generic-server/src/plugin/container.rs

@@ -1,7 +1,6 @@
-use std::collections::{HashMap, HashSet};
-
 use axum::Router;
 use server_plugin::{Plugin, PluginError};
+use std::collections::{HashMap, HashSet};
 
 use super::PLUGINS;
 
@@ -89,6 +88,37 @@ impl<'a> PluginContainer<'a> {
         }
     }
 
+    /// unload container plugins
+    pub fn unload(&mut self) -> Result<(), PluginContainerError> {
+        // Unmount plugins and clearing the vector of routes
+        let errors: HashMap<_, _> = self
+            .plugins
+            .iter()
+            .filter_map(|plugin| match plugin.unmount() {
+                Ok(_) => {
+                    tracing::info!("Unmounted plugin {}", plugin.name());
+                    None
+                }
+                Err(err) => {
+                    tracing::error!("error unmounting plugin {}", plugin.name());
+                    Some((plugin.name().to_owned(), err))
+                }
+            })
+            .collect();
+
+        // Flag as unloaded
+        self.loaded = false;
+
+        // Return state of completion
+        if errors.is_empty() {
+            self.collected_routes.clear();
+            tracing::debug!("plugin container unloaded");
+            Ok(())
+        } else {
+            Err(PluginContainerError::PluginErrorMap(errors))
+        }
+    }
+
     /// Merge collected routes from all plugins successfully initialized.
     pub fn routes(&self) -> Result<Router, PluginContainerError> {
         if self.loaded {
@@ -268,4 +298,20 @@ mod tests {
             PluginContainerError::Unloaded
         );
     }
+    #[test]
+    fn test_unloading() {
+        // loading container from test_loading
+        let mut container = PluginContainer {
+            loaded: false,
+            collected_routes: vec![],
+            plugins: &vec![Box::new(FirstPlugin {}), Box::new(SecondPlugin {})],
+        };
+        assert!(container.load().is_ok());
+        assert!(container.routes().is_ok());
+
+        // unloading container and clearing routes
+        assert!(container.unload().is_ok());
+
+        assert_eq!(container.collected_routes.len(), 0);
+    }
 }