split startup & shtudown handlers

sparckles · Sep 29, 2024 · 77d5f16 · 77d5f16
1 parent c1660f8
commit 77d5f16
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 5 deletions.
diff --git a/src/executors/mod.rs b/src/executors/mod.rs
@@ -100,13 +100,35 @@ pub async fn execute_http_function(
     })
 }
 
-pub async fn execute_event_handler(
+pub async fn execute_startup_handler(
     event_handler: Option<Arc<FunctionInfo>>,
     task_locals: &TaskLocals,
 ) -> Result<()> {
     if let Some(function) = event_handler {
         if function.is_async {
             debug!("Startup event handler async");
+            Python::with_gil(|py| {
+                pyo3_asyncio::into_future_with_locals(
+                    task_locals,
+                    function.handler.as_ref(py).call0()?,
+                )
+            })?
+            .await?;
+        } else {
+            debug!("Startup event handler");
+            Python::with_gil(|py| function.handler.call0(py))?;
+        }
+    }
+    Ok(())
+}
+
+pub async fn execute_shutdown_handler(
+    event_handler: Option<Arc<FunctionInfo>>,
+    task_locals: &TaskLocals,
+) -> Result<()> {
+    if let Some(function) = event_handler {
+        if function.is_async {
+            debug!("Shutdown event handler async");
             Python::with_gil(|py| {
                 pyo3_asyncio::tokio::run_until_complete(
                     task_locals.event_loop(py),
@@ -117,7 +139,7 @@ pub async fn execute_event_handler(
                 )
             })?;
         } else {
-            debug!("Startup event handler");
+            debug!("Shutdown event handler");
             Python::with_gil(|py| function.handler.call0(py))?;
         }
     }

diff --git a/src/server.rs b/src/server.rs
@@ -1,4 +1,4 @@
-use crate::executors::{execute_event_handler, execute_http_function, execute_middleware_function};
+use crate::executors::{execute_http_function, execute_middleware_function, execute_shutdown_handler, execute_startup_handler};
 
 use crate::routers::const_router::ConstRouter;
 use crate::routers::Router;
@@ -122,7 +122,7 @@ impl Server {
         thread::spawn(move || {
             actix_web::rt::System::new().block_on(async move {
                 debug!("The number of workers is {}", workers);
-                execute_event_handler(startup_handler, &task_locals_copy)
+                execute_startup_handler(startup_handler, &task_locals_copy)
                     .await
                     .unwrap();
 
@@ -224,7 +224,7 @@ impl Server {
             debug!("Ctrl c handler");
             Python::with_gil(|py| {
                 pyo3_asyncio::tokio::run(py, async move {
-                    execute_event_handler(shutdown_handler, &task_locals.clone())
+                    execute_shutdown_handler(shutdown_handler, &task_locals.clone())
                         .await
                         .unwrap();
                     Ok(())