chore: change ipc server stop at last

GitDataAI · Aug 5, 2024 · bbbe488 · bbbe488
1 parent ea10e63
commit bbbe488
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 19 deletions.
diff --git a/crates/compute_unit_runner/src/bin/compute_unit_runner.rs b/crates/compute_unit_runner/src/bin/compute_unit_runner.rs
@@ -107,9 +107,17 @@ async fn main() -> Result<()> {
         //listen unix socket
         let unix_socket_addr = args.unix_socket_addr.clone();
         let program = program_safe.clone();
+        let token = token.clone();
         join_set.spawn(async move {
             info!("start ipc server {}", &unix_socket_addr);
-            ipc::start_ipc_server(unix_socket_addr, program).await
+            let handler = ipc::start_ipc_server(&unix_socket_addr, program)?;
+            select! {
+                _ = token.cancelled() => {
+                   handler.stop(true).await;
+                   info!("ipc server stopped");
+                   return Ok(());
+                }
+            };
         });
     }
     {

diff --git a/crates/compute_unit_runner/src/ipc.rs b/crates/compute_unit_runner/src/ipc.rs
@@ -1,19 +1,12 @@
 use crate::media_data_tracker::MediaDataTracker;
 use actix_web::{
-    error,
-    http::StatusCode,
-    middleware,
-    web,
-    web::Data,
-    App,
-    HttpRequest,
-    HttpResponse,
-    HttpServer,
+    dev::ServerHandle, error, http::StatusCode, middleware, web::{self, Data}, App, HttpRequest, HttpResponse, HttpServer
 };
 use anyhow::{
     anyhow,
     Result,
 };
+use std::os::unix::net::UnixListener;
 use core::str;
 use http_body_util::Collected;
 use jz_action::{
@@ -32,11 +25,10 @@ use std::{
     time::Duration,
 };
 use tokio::{
-    sync::{
+     sync::{
         oneshot,
         Mutex,
-    },
-    time::sleep,
+    }, time::sleep
 };
 use tracing::info;
 
@@ -225,10 +217,10 @@ where
     }
 }
 
-pub async fn start_ipc_server<R>(
-    unix_socket_addr: String,
+pub fn start_ipc_server<R>(
+    unix_socket_addr: &str,
     program: Arc<Mutex<MediaDataTracker<R>>>,
-) -> Result<()>
+) -> Result<ServerHandle>
 where
     R: DbRepo + Clone + Send + Sync + 'static,
 {
@@ -262,9 +254,11 @@ where
             )
             .app_data(web::JsonConfig::default().error_handler(json_error_handler))
     })
-    .bind_uds(unix_socket_addr)?;
-    server.run().await.unwrap();
-    Ok(())
+    .disable_signals()
+    .bind_uds(unix_socket_addr)?
+    .run();
+
+    Ok( server.handle())
 }
 
 pub trait IPCClient {