From ad06b2a903c42b324533917d45fc3bac81b4d01a Mon Sep 17 00:00:00 2001 From: Akarshan Date: Mon, 23 Jun 2025 12:15:04 +0530 Subject: [PATCH] Move llama-server cleanup code to a separate file --- .../inference_llamacpp_extension/cleanup.rs | 58 ++++++++++++ .../inference_llamacpp_extension/mod.rs | 1 + src-tauri/src/lib.rs | 88 +------------------ 3 files changed, 61 insertions(+), 86 deletions(-) create mode 100644 src-tauri/src/core/utils/extensions/inference_llamacpp_extension/cleanup.rs diff --git a/src-tauri/src/core/utils/extensions/inference_llamacpp_extension/cleanup.rs b/src-tauri/src/core/utils/extensions/inference_llamacpp_extension/cleanup.rs new file mode 100644 index 000000000..853969dd5 --- /dev/null +++ b/src-tauri/src/core/utils/extensions/inference_llamacpp_extension/cleanup.rs @@ -0,0 +1,58 @@ +use tauri::State; +use crate::core::state::AppState; + +pub async fn cleanup_processes(state: State<'_, AppState>) { + let mut map = state.llama_server_process.lock().await; + let pids: Vec = map.keys().cloned().collect(); + for pid in pids { + if let Some(mut child) = map.remove(&pid) { + #[cfg(unix)] + { + use nix::sys::signal::{kill, Signal}; + use nix::unistd::Pid; + use tokio::time::{timeout, Duration}; + + if let Some(raw_pid) = child.id() { + let raw_pid = raw_pid as i32; + log::info!("Sending SIGTERM to PID {} during shutdown", raw_pid); + let _ = kill(Pid::from_raw(raw_pid), Signal::SIGTERM); + + match timeout(Duration::from_secs(2), child.wait()).await { + Ok(Ok(status)) => log::info!("Process {} exited gracefully: {}", raw_pid, status), + Ok(Err(e)) => log::error!("Error waiting after SIGTERM for {}: {}", raw_pid, e), + Err(_) => { + log::warn!("SIGTERM timed out for PID {}; sending SIGKILL", raw_pid); + let _ = kill(Pid::from_raw(raw_pid), Signal::SIGKILL); + let _ = child.wait().await; + } + } + } + } + + #[cfg(windows)] + { + use windows_sys::Win32::System::Console::{GenerateConsoleCtrlEvent, CTRL_C_EVENT}; + use windows_sys::Win32::Foundation::BOOL; + use tokio::time::{timeout, Duration}; + + if let Some(raw_pid) = child.id() { + log::info!("Sending Ctrl-C to PID {} during shutdown", raw_pid); + let ok: BOOL = unsafe { GenerateConsoleCtrlEvent(CTRL_C_EVENT, raw_pid) }; + if ok == 0 { + log::error!("Failed to send Ctrl-C to PID {}", raw_pid); + } + + match timeout(Duration::from_secs(2), child.wait()).await { + Ok(Ok(status)) => log::info!("Process {} exited after Ctrl-C: {}", raw_pid, status), + Ok(Err(e)) => log::error!("Error waiting after Ctrl-C for {}: {}", raw_pid, e), + Err(_) => { + log::warn!("Timed out for PID {}; force-killing", raw_pid); + let _ = child.kill().await; + let _ = child.wait().await; + } + } + } + } + } + } +} diff --git a/src-tauri/src/core/utils/extensions/inference_llamacpp_extension/mod.rs b/src-tauri/src/core/utils/extensions/inference_llamacpp_extension/mod.rs index 74f47ad34..35a24a4f9 100644 --- a/src-tauri/src/core/utils/extensions/inference_llamacpp_extension/mod.rs +++ b/src-tauri/src/core/utils/extensions/inference_llamacpp_extension/mod.rs @@ -1 +1,2 @@ pub mod server; +pub mod cleanup; diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 493efbdbe..fa8968a0d 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -8,6 +8,7 @@ use core::{ use reqwest::Client; use std::{collections::HashMap, sync::Arc}; use tauri::{Emitter, Manager}; +use core::utils::extensions::inference_llamacpp_extension::cleanup::cleanup_processes; use tokio::sync::Mutex; @@ -135,92 +136,7 @@ pub fn run() { let state = window.app_handle().state::(); tauri::async_runtime::block_on(async { - let mut map = state.llama_server_process.lock().await; - let pids: Vec = map.keys().cloned().collect(); - for pid in pids { - if let Some(mut child) = map.remove(&pid) { - #[cfg(unix)] - { - use nix::sys::signal::{kill, Signal}; - use nix::unistd::Pid; - use tokio::time::{timeout, Duration}; - - if let Some(raw_pid) = child.id() { - let raw_pid = raw_pid as i32; - log::info!( - "Sending SIGTERM to PID {} during shutdown", - raw_pid - ); - let _ = kill(Pid::from_raw(raw_pid), Signal::SIGTERM); - - match timeout(Duration::from_secs(2), child.wait()).await { - Ok(Ok(status)) => log::info!( - "Process {} exited gracefully: {}", - raw_pid, - status - ), - Ok(Err(e)) => log::error!( - "Error waiting after SIGTERM for {}: {}", - raw_pid, - e - ), - Err(_) => { - log::warn!( - "SIGTERM timed out for PID {}; sending SIGKILL", - raw_pid - ); - let _ = - kill(Pid::from_raw(raw_pid), Signal::SIGKILL); - let _ = child.wait().await; - } - } - } - } - - #[cfg(windows)] - { - use tokio::time::{timeout, Duration}; - use windows_sys::Win32::Foundation::BOOL; - use windows_sys::Win32::System::Console::{ - GenerateConsoleCtrlEvent, CTRL_C_EVENT, - }; - - if let Some(raw_pid) = child.id() { - log::info!( - "Sending Ctrl-C to PID {} during shutdown", - raw_pid - ); - let ok: BOOL = unsafe { - GenerateConsoleCtrlEvent(CTRL_C_EVENT, raw_pid) - }; - if ok == 0 { - log::error!("Failed to send Ctrl-C to PID {}", raw_pid); - } - - match timeout(Duration::from_secs(2), child.wait()).await { - Ok(Ok(status)) => log::info!( - "Process {} exited after Ctrl-C: {}", - raw_pid, - status - ), - Ok(Err(e)) => log::error!( - "Error waiting after Ctrl-C for {}: {}", - raw_pid, - e - ), - Err(_) => { - log::warn!( - "Timed out for PID {}; force-killing", - raw_pid - ); - let _ = child.kill().await; - let _ = child.wait().await; - } - } - } - } - } - } + cleanup_processes(state).await; }); } let client = Client::new();