diff --git a/Cargo.toml b/Cargo.toml index 38e071b12..e1fd9d7cc 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -33,7 +33,7 @@ tokio-test = "0.4" criterion = { version = "0.5", features = ["html_reports"] } [profile.release] -panic = "abort" +panic = "unwind" codegen-units = 1 lto = true opt-level = 3 # Use speed optimization, not size - critical for SurrealDB performance diff --git a/packages/desktop-app/src-tauri/Cargo.toml b/packages/desktop-app/src-tauri/Cargo.toml index 5a78a2cac..5afc73829 100644 --- a/packages/desktop-app/src-tauri/Cargo.toml +++ b/packages/desktop-app/src-tauri/Cargo.toml @@ -39,6 +39,7 @@ tracing-subscriber = { version = "0.3", features = ["env-filter"] } futures = "0.3" axum = { version = "0.7", features = ["default"] } tokio-stream = "0.1" +tokio-util = "0.7" async-stream = "0.3" tower-http = { version = "0.6", features = ["cors"] } surrealdb = { workspace = true } diff --git a/packages/desktop-app/src-tauri/src/commands/db.rs b/packages/desktop-app/src-tauri/src/commands/db.rs index b9fb6c0f7..381a2e5c2 100644 --- a/packages/desktop-app/src-tauri/src/commands/db.rs +++ b/packages/desktop-app/src-tauri/src/commands/db.rs @@ -168,9 +168,12 @@ async fn init_services(app: &AppHandle, db_path: PathBuf) -> Result<(), String> eprintln!("✅ [init_services] All services registered with Tauri"); tracing::info!("✅ [init_services] All services registered with Tauri"); + // Retrieve the shutdown token from Tauri state for background task coordination + let shutdown_token: tauri::State = app.state(); + // Initialize MCP server now that NodeService is available // MCP will use the same NodeService as Tauri commands - if let Err(e) = crate::initialize_mcp_server(app.clone()) { + if let Err(e) = crate::initialize_mcp_server(app.clone(), shutdown_token.child_token()) { tracing::error!("❌ Failed to initialize MCP server: {}", e); // Don't fail database init if MCP fails - MCP is optional } @@ -178,9 +181,12 @@ async fn init_services(app: &AppHandle, db_path: PathBuf) -> Result<(), String> // Initialize domain event forwarding with client filtering (#665) // Events that originated from this Tauri client are filtered out to prevent feedback loops let client_id = "tauri-main".to_string(); - if let Err(e) = - crate::initialize_domain_event_forwarder(app.clone(), node_service_arc.clone(), client_id) - { + if let Err(e) = crate::initialize_domain_event_forwarder( + app.clone(), + node_service_arc.clone(), + client_id, + shutdown_token.child_token(), + ) { tracing::error!("❌ Failed to initialize domain event forwarder: {}", e); // Don't fail database init if event forwarding fails - it's not critical } diff --git a/packages/desktop-app/src-tauri/src/lib.rs b/packages/desktop-app/src-tauri/src/lib.rs index af652491b..49036d439 100644 --- a/packages/desktop-app/src-tauri/src/lib.rs +++ b/packages/desktop-app/src-tauri/src/lib.rs @@ -36,10 +36,14 @@ mod tests; /// achieving real-time sync through event-driven architecture. /// /// Events that originated from this Tauri client are filtered out (prevents feedback loop). +/// +/// The `cancel_token` is used for graceful shutdown - when cancelled, the forwarder +/// will stop its event loop and exit cleanly before the Tokio runtime drops. pub fn initialize_domain_event_forwarder( app: tauri::AppHandle, node_service: std::sync::Arc, client_id: String, + cancel_token: tokio_util::sync::CancellationToken, ) -> anyhow::Result<()> { use crate::services::DomainEventForwarder; use futures::FutureExt; @@ -53,7 +57,7 @@ pub fn initialize_domain_event_forwarder( tauri::async_runtime::spawn(async move { let result = std::panic::AssertUnwindSafe(async { let forwarder = DomainEventForwarder::new(node_service, app, client_id); - forwarder.run().await + forwarder.run(cancel_token).await }) .catch_unwind() .await; @@ -84,8 +88,14 @@ pub fn initialize_domain_event_forwarder( /// and NodeEmbeddingService and spawns the MCP server task with them, /// ensuring MCP and Tauri commands operate on the same database. /// +/// The `cancel_token` is used for graceful shutdown - when cancelled, the MCP +/// server task will be aborted before the Tokio runtime drops. +/// /// As of Issue #715, uses McpServerService from nodespace-core for managed lifecycle. -pub fn initialize_mcp_server(app: tauri::AppHandle) -> anyhow::Result<()> { +pub fn initialize_mcp_server( + app: tauri::AppHandle, + cancel_token: tokio_util::sync::CancellationToken, +) -> anyhow::Result<()> { use crate::commands::embeddings::EmbeddingState; use futures::FutureExt; use nodespace_core::NodeService; @@ -119,13 +129,22 @@ pub fn initialize_mcp_server(app: tauri::AppHandle) -> anyhow::Result<()> { // Spawn MCP server task with Tauri event emissions // Uses panic protection to prevent silent background task failures + // Monitors cancel_token for graceful shutdown before runtime drops tauri::async_runtime::spawn(async move { // Use FutureExt::catch_unwind for proper async panic catching // This avoids the deadlock risk of using block_on inside an async task // AssertUnwindSafe is needed because services contain non-UnwindSafe types - let result = std::panic::AssertUnwindSafe(mcp_service.start_with_callback(callback)) - .catch_unwind() - .await; + let result = std::panic::AssertUnwindSafe(async { + tokio::select! { + res = mcp_service.start_with_callback(callback) => res, + _ = cancel_token.cancelled() => { + tracing::info!("MCP server received shutdown signal"); + Ok(()) + } + } + }) + .catch_unwind() + .await; match result { Ok(Ok(_)) => { @@ -133,11 +152,9 @@ pub fn initialize_mcp_server(app: tauri::AppHandle) -> anyhow::Result<()> { } Ok(Err(e)) => { tracing::error!("❌ MCP server error: {}", e); - // TODO: Consider emitting Tauri event to notify UI of MCP failure } Err(panic_info) => { tracing::error!("💥 MCP server panicked: {:?}", panic_info); - // TODO: Consider attempting automatic restart or notifying user } } }); @@ -145,14 +162,45 @@ pub fn initialize_mcp_server(app: tauri::AppHandle) -> anyhow::Result<()> { Ok(()) } +/// Shared shutdown token for graceful background task termination. +/// +/// Managed as Tauri state so it can be accessed from both the setup phase +/// (where background tasks are spawned) and the run event handler (where +/// shutdown is triggered). When cancelled, all background tasks (MCP server, +/// domain event forwarder) exit their loops before the Tokio runtime drops. +#[derive(Clone)] +pub struct ShutdownToken(tokio_util::sync::CancellationToken); + +impl ShutdownToken { + fn new() -> Self { + Self(tokio_util::sync::CancellationToken::new()) + } + + /// Create a child token for a background task. + /// Cancelling the parent automatically cancels all children. + pub fn child_token(&self) -> tokio_util::sync::CancellationToken { + self.0.child_token() + } + + /// Signal all background tasks to shut down. + /// Idempotent - safe to call multiple times. + pub fn cancel(&self) { + self.0.cancel(); + } +} + #[cfg_attr(mobile, tauri::mobile_entry_point)] pub fn run() { use tauri::{menu::*, Emitter, Manager, RunEvent}; + // Create shutdown token for coordinating graceful background task termination + let shutdown_token = ShutdownToken::new(); + let shutdown_token_for_setup = shutdown_token.clone(); + let app = tauri::Builder::default() .plugin(tauri_plugin_opener::init()) .plugin(tauri_plugin_dialog::init()) - .setup(|app| { + .setup(move |app| { // Create menu items let toggle_sidebar = MenuItemBuilder::new("Toggle Sidebar") .id("toggle_sidebar") @@ -209,6 +257,10 @@ pub fn run() { // Set the menu app.set_menu(menu)?; + // Register shutdown token as managed state so commands/db.rs can access it + // when spawning background tasks (MCP server, domain event forwarder) + app.manage(shutdown_token_for_setup); + // Note: MCP server initialization is deferred until database is initialized // See commands/db.rs::init_services() which calls initialize_mcp_server() // after NodeService is available in Tauri state @@ -308,26 +360,34 @@ pub fn run() { .expect("error while building tauri application"); // Run with event handler for graceful shutdown - // This allows proper cleanup of Metal/GPU resources before exit + // This allows proper cleanup of Metal/GPU resources and background tasks before exit // // Note: On macOS, RunEvent::ExitRequested may not fire reliably (Tauri issue #9198) - // We handle cleanup in multiple places to ensure GPU resources are released: + // We handle cleanup in multiple places to ensure resources are released: // 1. WindowEvent::CloseRequested - when user clicks X or Cmd+Q // 2. RunEvent::ExitRequested - when app is about to exit // 3. RunEvent::Exit - final cleanup before process termination - app.run(|app_handle, event| { + // + // The shutdown_token is cancelled to signal background tasks (MCP server, + // domain event forwarder) to exit their loops before the Tokio runtime drops. + let shutdown_token_for_events = shutdown_token.clone(); + app.run(move |app_handle, event| { match event { RunEvent::WindowEvent { label, event: tauri::WindowEvent::CloseRequested { .. }, .. } => { - // Window close requested - release GPU resources before the window closes + // Window close requested - signal background tasks and release GPU resources // This is the most reliable place to do cleanup on macOS tracing::info!( - "Window '{}' close requested, releasing GPU context...", + "Window '{}' close requested, signaling shutdown and releasing GPU context...", label ); + shutdown_token_for_events.cancel(); + // Brief pause to let background tasks exit their loops before + // releasing GPU resources they may still reference + std::thread::sleep(std::time::Duration::from_millis(50)); release_gpu_resources(app_handle); } RunEvent::ExitRequested { code, .. } => { @@ -336,23 +396,32 @@ pub fn run() { "App exit requested (code: {:?}), performing cleanup...", code ); + shutdown_token_for_events.cancel(); release_gpu_resources(app_handle); tracing::info!("Cleanup complete, exiting..."); } RunEvent::Exit => { - // Final exit - last chance to cleanup (may be too late for GPU resources) - tracing::info!("App exiting normally"); + // Final exit - ensure shutdown signal is sent (idempotent) + tracing::info!("App exiting, ensuring shutdown signal sent..."); + shutdown_token_for_events.cancel(); } _ => {} } }); } -/// Release GPU resources (Metal context) to prevent SIGABRT crash on exit. +/// Release GPU resources (Metal context and backend) to prevent SIGABRT crash on exit. /// /// This must be called before the app exits to properly clean up Metal/GPU /// resources. The crash occurs when ggml_metal_rsets_free is called during /// static destruction (__cxa_finalize_ranges) while resources are still in use. +/// +/// Cleanup order is critical: +/// 1. Release LlamaState (context + model) - frees Metal residency sets +/// 2. Release global LLAMA_BACKEND - frees the Metal backend itself +/// +/// This prevents __cxa_finalize_ranges from encountering Metal resources during +/// static destruction, which would cause SIGABRT in production builds. fn release_gpu_resources(app_handle: &tauri::AppHandle) { use tauri::Manager; @@ -360,9 +429,12 @@ fn release_gpu_resources(app_handle: &tauri::AppHandle) { app_handle.try_state::() { tracing::info!("Releasing GPU context to prevent Metal crash..."); - // Release the LlamaContext which holds Metal residency sets - // This must happen before the global LLAMA_BACKEND static is destroyed + // Step 1: Release the LlamaState (context + model) which holds Metal residency sets embedding_state.service.nlp_engine().release_gpu_context(); tracing::info!("GPU context released successfully"); } + + // Step 2: Release the global llama backend itself + // Must happen AFTER all models/contexts are dropped (step 1) + nodespace_nlp_engine::release_llama_backend(); } diff --git a/packages/desktop-app/src-tauri/src/services/domain_event_forwarder.rs b/packages/desktop-app/src-tauri/src/services/domain_event_forwarder.rs index 05e4e6a3e..897714f04 100644 --- a/packages/desktop-app/src-tauri/src/services/domain_event_forwarder.rs +++ b/packages/desktop-app/src-tauri/src/services/domain_event_forwarder.rs @@ -60,7 +60,10 @@ impl DomainEventForwarder { /// /// Filters out events where source_client_id matches this client's ID /// to prevent feedback loops from the Tauri frontend receiving its own changes. - pub async fn run(self) -> Result<()> { + /// + /// The `cancel_token` allows graceful shutdown - when cancelled, the loop exits + /// cleanly instead of being killed when the Tokio runtime drops. + pub async fn run(self, cancel_token: tokio_util::sync::CancellationToken) -> Result<()> { info!( "🔧 Starting domain event forwarding service (client_id: {})", self.client_id @@ -74,22 +77,29 @@ impl DomainEventForwarder { info!("✅ Event subscription established successfully"); - // Listen for events and forward to frontend + // Listen for events and forward to frontend, with graceful shutdown support loop { - match rx.recv().await { - Ok(event) => { - debug!("Received domain event: {:?}", event); - self.forward_event(&event); - } - Err(broadcast::error::RecvError::Lagged(_)) => { - // Event queue was too long, skip missed events - debug!("Event queue lagged, some events may have been skipped"); + tokio::select! { + result = rx.recv() => { + match result { + Ok(event) => { + debug!("Received domain event: {:?}", event); + self.forward_event(&event); + } + Err(broadcast::error::RecvError::Lagged(_)) => { + debug!("Event queue lagged, some events may have been skipped"); + } + Err(broadcast::error::RecvError::Closed) => { + error!("Event broadcast channel closed, stopping synchronization service"); + self.emit_status("disconnected", Some("channel-closed")); + return Err(anyhow::anyhow!("Event broadcast channel closed")); + } + } } - Err(broadcast::error::RecvError::Closed) => { - // Broadcast channel closed, service should shut down - error!("Event broadcast channel closed, stopping synchronization service"); - self.emit_status("disconnected", Some("channel-closed")); - return Err(anyhow::anyhow!("Event broadcast channel closed")); + _ = cancel_token.cancelled() => { + info!("Domain event forwarder received shutdown signal, exiting gracefully"); + self.emit_status("disconnected", Some("shutdown")); + return Ok(()); } } } diff --git a/packages/nlp-engine/src/embedding.rs b/packages/nlp-engine/src/embedding.rs index e5dba9b7a..6c2367b32 100644 --- a/packages/nlp-engine/src/embedding.rs +++ b/packages/nlp-engine/src/embedding.rs @@ -13,7 +13,7 @@ use crate::config::EmbeddingConfig; use crate::error::{EmbeddingError, Result}; use lru::LruCache; use std::num::NonZeroUsize; -use std::sync::{Arc, Mutex, OnceLock}; +use std::sync::{Arc, Mutex}; /// Embedding vector dimension for nomic-embed-vision-v1.5 pub const EMBEDDING_DIMENSION: usize = 768; @@ -36,76 +36,57 @@ use llama_cpp_2::model::params::LlamaModelParams; use llama_cpp_2::model::{AddBos, LlamaModel}; /// Global llama backend singleton. -/// The llama.cpp backend can only be initialized once per process, so we use a OnceLock -/// to ensure thread-safe initialization and allow multiple EmbeddingService instances -/// to share the same backend (important for tests running in parallel). +/// +/// Uses `Mutex>` instead of `OnceLock` so the backend can be +/// explicitly cleared during shutdown. This prevents SIGABRT crashes from +/// `__cxa_finalize_ranges` trying to destroy Metal/GPU resources during static +/// destruction when the C runtime state is already partially torn down. +/// +/// The llama.cpp backend can only be initialized once per process. +/// The Mutex ensures thread-safe initialization and allows multiple EmbeddingService +/// instances to share the same backend (important for tests running in parallel). #[cfg(feature = "embedding-service")] -static LLAMA_BACKEND: OnceLock = OnceLock::new(); +static LLAMA_BACKEND: Mutex> = Mutex::new(None); /// Initialize or get the global llama backend. -/// Returns a reference to the singleton backend instance. +/// Returns a guard holding a reference to the singleton backend instance. /// /// This handles the case where multiple threads try to initialize the backend /// concurrently. The llama.cpp backend is a global singleton that can only be /// initialized once per process. +/// +/// The returned `BackendGuard` must be held for the duration of backend usage +/// (e.g., model loading, context creation). #[cfg(feature = "embedding-service")] -fn get_or_init_backend() -> Result<&'static LlamaBackend> { +fn get_or_init_backend() -> Result { use llama_cpp_2::LlamaCppError; - // Try to get existing backend first - if let Some(backend) = LLAMA_BACKEND.get() { - return Ok(backend); + let mut guard = LLAMA_BACKEND.lock().unwrap_or_else(|p| p.into_inner()); + + // Already initialized - return the guard + if guard.is_some() { + return Ok(BackendGuard(guard)); } - // Try to initialize - only one thread will succeed + // Try to initialize - only one thread will succeed (Mutex ensures exclusivity) match LlamaBackend::init() { Ok(backend) => { - // Try to store it - might fail if another thread beat us - match LLAMA_BACKEND.set(backend) { - Ok(()) => {} - Err(_already_set) => { - // Another thread beat us - that's fine, we'll use theirs - // The backend we created will be dropped, but that's safe - // because llama_backend_free is idempotent - } - } - // Return whatever is stored (either ours or the other thread's) - Ok(LLAMA_BACKEND.get().expect("Backend must be initialized")) + *guard = Some(backend); + Ok(BackendGuard(guard)) } Err(LlamaCppError::BackendAlreadyInitialized) => { - // The C library backend is initialized, but our OnceLock might not have - // the reference yet. This can happen when: - // 1. Another thread is about to store the backend (race condition) - // 2. A previous test run left stale C global state - // - // For case 1, wait with exponential backoff for the OnceLock to be set. - // For case 2, we need to create a new wrapper that references the existing - // C backend (llama.cpp backend init is idempotent - it reuses existing state). - for i in 0..20 { - if let Some(backend) = LLAMA_BACKEND.get() { - return Ok(backend); - } - // Exponential backoff: 1ms, 2ms, 4ms... up to ~1s total - std::thread::sleep(std::time::Duration::from_millis(1 << i.min(8))); - } - - // Last resort: try to create a new LlamaBackend anyway. - // llama.cpp's backend_init is idempotent and will return OK if already initialized. - // This handles the case where the C backend is initialized but our Rust wrapper isn't. + // The C library backend is initialized but our Option is None. + // This can happen if the backend was cleared during a previous shutdown + // but the C global state persists (e.g., in tests). + // Try again - llama.cpp backend init is idempotent. match LlamaBackend::init() { Ok(backend) => { - let _ = LLAMA_BACKEND.set(backend); - Ok(LLAMA_BACKEND.get().expect("Backend must be set")) - } - Err(_) => { - // If still failing, check one more time if another thread set it - if let Some(backend) = LLAMA_BACKEND.get() { - return Ok(backend); - } - Err(EmbeddingError::ModelLoadError( - "Backend initialization failed after multiple attempts".to_string(), - )) + *guard = Some(backend); + Ok(BackendGuard(guard)) } + Err(_) => Err(EmbeddingError::ModelLoadError( + "Backend initialization failed after multiple attempts".to_string(), + )), } } Err(e) => Err(EmbeddingError::ModelLoadError(format!( @@ -115,6 +96,46 @@ fn get_or_init_backend() -> Result<&'static LlamaBackend> { } } +/// RAII guard that holds a lock on the global LLAMA_BACKEND. +/// +/// Dereferences to `&LlamaBackend` for convenient use with llama.cpp APIs. +/// The lock is released when the guard is dropped. +#[cfg(feature = "embedding-service")] +struct BackendGuard(std::sync::MutexGuard<'static, Option>); + +#[cfg(feature = "embedding-service")] +impl std::ops::Deref for BackendGuard { + type Target = LlamaBackend; + + fn deref(&self) -> &Self::Target { + self.0 + .as_ref() + .expect("BackendGuard created with None backend") + } +} + +/// Explicitly release the global llama backend to prevent SIGABRT on process exit. +/// +/// This drops the `LlamaBackend` while we still have full control of the process, +/// rather than letting `__cxa_finalize_ranges` destroy it during static destruction +/// when Metal/GPU state may already be partially torn down. +/// +/// Must be called AFTER all `LlamaState` instances (models, contexts) are dropped. +/// Safe to call multiple times (idempotent). +#[cfg(feature = "embedding-service")] +pub fn release_llama_backend() { + let mut guard = LLAMA_BACKEND.lock().unwrap_or_else(|p| p.into_inner()); + if guard.take().is_some() { + tracing::info!("Global LLAMA_BACKEND released, Metal resources freed"); + } +} + +/// Explicitly release the global llama backend (no-op when embedding-service feature is disabled). +#[cfg(not(feature = "embedding-service"))] +pub fn release_llama_backend() { + // No-op: no backend to release when embedding-service feature is disabled +} + /// Wrapper to hold model and context together with proper lifetimes. /// /// ## Safety (Issue #776) @@ -190,9 +211,9 @@ impl LlamaState { .with_n_threads_batch(self.n_threads) .with_embeddings(true); - // Get the global backend reference + // Get the global backend (holds Mutex lock for duration of context creation) let backend = get_or_init_backend()?; - let ctx = self.model.new_context(backend, ctx_params).map_err(|e| { + let ctx = self.model.new_context(&backend, ctx_params).map_err(|e| { EmbeddingError::InferenceError(format!("Context creation failed: {}", e)) })?; @@ -289,13 +310,14 @@ impl EmbeddingService { tracing::info!("Model path resolved to: {:?}", model_path); // Initialize llama.cpp backend (uses global singleton) + // BackendGuard holds the Mutex lock for the duration of model loading let backend = get_or_init_backend()?; // Load model with GPU offloading let model_params = LlamaModelParams::default().with_n_gpu_layers(self.config.n_gpu_layers); - let model = LlamaModel::load_from_file(backend, &model_path, &model_params) + let model = LlamaModel::load_from_file(&backend, &model_path, &model_params) .map_err(|e| EmbeddingError::ModelLoadError(format!("Model load failed: {}", e)))?; // Get actual embedding dimension from model diff --git a/packages/nlp-engine/src/lib.rs b/packages/nlp-engine/src/lib.rs index 9ed5a10f7..87be80f17 100644 --- a/packages/nlp-engine/src/lib.rs +++ b/packages/nlp-engine/src/lib.rs @@ -39,5 +39,5 @@ pub mod error; // Re-export main types pub use config::EmbeddingConfig; -pub use embedding::{EmbeddingService, EMBEDDING_DIMENSION}; +pub use embedding::{release_llama_backend, EmbeddingService, EMBEDDING_DIMENSION}; pub use error::{EmbeddingError, Result};