diff --git a/crates/llm/src/lib.rs b/crates/llm/src/lib.rs index df6456d..26016f4 100644 --- a/crates/llm/src/lib.rs +++ b/crates/llm/src/lib.rs @@ -1,5 +1,6 @@ use std::num::NonZeroU32; use std::path::Path; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex}; use encoding_rs::UTF_8; @@ -31,6 +32,8 @@ pub enum EngineError { NotLoaded, #[error("LLM load failed: {0}")] LoadFailed(String), + #[error("Another LLM load is already in flight; refusing to start a parallel load.")] + AlreadyLoading, #[error( "prompt too long: {prompt_tokens} prompt tokens exceed the {available_prompt_tokens}-token prompt budget for an {context_window}-token context with {max_tokens} reserved response tokens" )] @@ -83,6 +86,29 @@ struct LlmState { #[derive(Clone, Default)] pub struct LlmEngine { inner: Arc>, + /// Flag held for the duration of a model load. The std::sync::Mutex + /// covers cheap state mutations (~microseconds); the multi-second + /// `LlamaModel::load_from_file` call now runs *outside* the mutex, + /// so polls like `is_loaded()` / `loaded_model_id()` (called from + /// sync Tauri handlers without `spawn_blocking`) don't park tokio + /// worker threads on a slow C++ FFI call. This Atomic also doubles + /// as a TOCTOU guard: two concurrent `load_model` invocations on + /// the same engine will not both reach the heavy load — the second + /// returns `EngineError::AlreadyLoading`. + loading: Arc, +} + +/// RAII guard that clears the `loading` flag on drop, including on +/// panic / early-return. Prevents the engine getting stuck in a +/// "permanently loading" state if a load fails midway. +struct LoadingGuard { + flag: Arc, +} + +impl Drop for LoadingGuard { + fn drop(&mut self) { + self.flag.store(false, Ordering::Release); + } } impl LlmEngine { @@ -100,18 +126,91 @@ impl LlmEngine { model_path: &Path, use_gpu: bool, ) -> Result<(), EngineError> { - let mut guard = self.inner.lock().unwrap(); + self.load_model_with(model_id, model_path, use_gpu, |backend, path, params| { + LlamaModel::load_from_file(backend, path, params) + .map_err(|e| EngineError::LoadFailed(format!("model load: {e}"))) + }) + } - if let Some(loaded) = &guard.loaded { - if loaded.model_id == model_id.as_str() - && loaded.model_path == model_path.display().to_string() - && loaded.use_gpu == use_gpu - { - return Ok(()); + /// Core load implementation with a swappable file-loader closure. + /// Production callers use `load_model`, which delegates here with + /// the real `LlamaModel::load_from_file`. Tests inject a sleepy / + /// counting closure to exercise the locking discipline without + /// pulling a real GGUF off disk. + /// + /// Locking discipline (the whole point of this function): + /// 1. Take the mutex briefly to compare against the currently + /// loaded triple — if it matches, return early. No-op fast path. + /// 2. CAS the `loading` flag from false → true. If another load is + /// already in flight, refuse with `AlreadyLoading` rather than + /// starting a parallel one. A `LoadingGuard` ensures the flag + /// is cleared on every exit path including panic. + /// 3. Take the mutex briefly to drop the OLD model Arc (frees its + /// VRAM via `llama_free_model`) before the new load begins. + /// The backend Arc is preserved — `LlamaBackend::init()` is a + /// one-shot per process (an `AtomicBool` in llama-cpp-2 enforces + /// `BackendAlreadyInitialized` on a second call), so we must + /// never drop the backend while the process keeps running. + /// Note: `is_loaded()` reports false during the swap window — + /// that is the correct semantics. Callers wanting "model X is + /// loaded" must check `loaded_model_id()` against their target. + /// 4. Initialise the backend if absent (first-ever load only) and + /// run the slow `load_from_file` call — both OUTSIDE the mutex. + /// 5. Take the mutex briefly to install the new backend (if just + /// initialised) and the new model Arc. + fn load_model_with( + &self, + model_id: LlmModelId, + model_path: &Path, + use_gpu: bool, + loader: F, + ) -> Result<(), EngineError> + where + F: FnOnce(&LlamaBackend, &Path, &LlamaModelParams) -> Result, + { + // Step 1: short crit section — already-loaded fast path. + { + let guard = self.inner.lock().unwrap(); + if let Some(loaded) = &guard.loaded { + if loaded.model_id == model_id.as_str() + && loaded.model_path == model_path.display().to_string() + && loaded.use_gpu == use_gpu + { + return Ok(()); + } } } - let backend = match guard.backend.clone() { + // Step 2: claim the loading slot. Refuse if a parallel load is + // already mid-flight rather than starting a second slow load + // and silently overwriting the first. + if self + .loading + .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) + .is_err() + { + return Err(EngineError::AlreadyLoading); + } + let _loading_guard = LoadingGuard { + flag: Arc::clone(&self.loading), + }; + + // Step 3: short crit section — drop the OLD model so its VRAM is + // released BEFORE we allocate the new one. Without this, a swap + // briefly holds two models resident (Lifecycle-1: an + // ~17 GB Q4 27B swap on a 24 GB card OOMs even though either + // model fits alone). Keep the backend Arc — see locking notes. + let existing_backend = { + let mut guard = self.inner.lock().unwrap(); + guard.model = None; + guard.loaded = None; + guard.backend.clone() + }; + + // Step 4: heavy work OUTSIDE the mutex. `is_loaded()` and + // `loaded_model_id()` can be polled freely here without parking + // tokio worker threads. + let backend = match existing_backend { Some(existing) => existing, None => Arc::new( LlamaBackend::init() @@ -121,27 +220,72 @@ impl LlmEngine { let gpu_layers = if use_gpu { u32::MAX } else { 0 }; let params = LlamaModelParams::default().with_n_gpu_layers(gpu_layers); - let model = LlamaModel::load_from_file(&backend, model_path, ¶ms) - .map_err(|e| EngineError::LoadFailed(format!("model load: {e}")))?; + let model = loader(&backend, model_path, ¶ms)?; - guard.backend = Some(backend); - guard.model = Some(Arc::new(model)); - guard.loaded = Some(LoadedModelState { - model_id: model_id.as_str().to_string(), - model_path: model_path.display().to_string(), - use_gpu, - }); + // Step 5: short crit section — install the new state. + { + let mut guard = self.inner.lock().unwrap(); + guard.backend = Some(backend); + guard.model = Some(Arc::new(model)); + guard.loaded = Some(LoadedModelState { + model_id: model_id.as_str().to_string(), + model_path: model_path.display().to_string(), + use_gpu, + }); + } + // `_loading_guard` drops here and clears the flag. Ok(()) } pub fn unload(&self) -> Result<(), EngineError> { let mut guard = self.inner.lock().unwrap(); guard.model = None; - guard.backend = None; + // Backend is process-singleton (llama-cpp-2 enforces this via + // `LLAMA_BACKEND_INITIALIZED`). Dropping the Arc here would call + // `llama_backend_free` and a subsequent `init` would succeed, but + // we keep it resident to avoid the init/free churn on every + // load/unload cycle. guard.loaded = None; Ok(()) } + /// True iff a model load is currently in flight. Exposed for tests + /// + frontends that want to render a "loading…" state without + /// polling `is_loaded()` (which returns false during a swap). + pub fn is_loading(&self) -> bool { + self.loading.load(Ordering::Acquire) + } + + /// Test-only harness: runs `op` while holding the same locking + /// discipline as `load_model_with` (loading flag claimed, model + /// state cleared, slow op runs OUTSIDE the inner mutex, new state + /// installed at the end). Used by the regression test to verify + /// that `is_loaded()` / `loaded_model_id()` don't block on the + /// slow section. Not part of the public API. + #[cfg(test)] + pub(crate) fn __test_run_with_lock_discipline(&self, op: F) -> Result<(), EngineError> + where + F: FnOnce(), + { + if self + .loading + .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) + .is_err() + { + return Err(EngineError::AlreadyLoading); + } + let _loading_guard = LoadingGuard { + flag: Arc::clone(&self.loading), + }; + { + let mut guard = self.inner.lock().unwrap(); + guard.model = None; + guard.loaded = None; + } + op(); + Ok(()) + } + pub fn is_loaded(&self) -> bool { self.inner.lock().unwrap().model.is_some() } @@ -666,4 +810,125 @@ mod tests { let n_ctx = preflight_context_window(7_104, 1_024).unwrap(); assert_eq!(n_ctx, MAX_CONTEXT_TOKENS); } + + /// Race-3 regression. The inner `std::sync::Mutex` MUST NOT be held + /// across the slow `LlamaModel::load_from_file` call: sync Tauri + /// command handlers like `get_llm_status`, `check_llm_model`, + /// `delete_llm_model`, and `test_llm_model` call `is_loaded()` / + /// `loaded_model_id()` from tokio worker threads without + /// `spawn_blocking`. If the lock is held for the duration of a + /// 5-15 s load, parallel status polls from the frontend park the + /// tokio executor and the whole UI deadlocks. + /// + /// Pre-fix this test FAILS — both probes time out because the load + /// holds the mutex. Post-fix it PASSES — probes return in ≤50 ms. + #[test] + fn is_loaded_does_not_block_on_slow_load() { + use std::sync::mpsc; + use std::thread; + use std::time::{Duration, Instant}; + + let engine = LlmEngine::new(); + let load_started = Arc::new(std::sync::Barrier::new(2)); + let release_load = Arc::new(std::sync::Barrier::new(2)); + + let engine_for_loader = engine.clone(); + let load_started_for_loader = Arc::clone(&load_started); + let release_load_for_loader = Arc::clone(&release_load); + + let loader_handle = thread::spawn(move || { + engine_for_loader + .__test_run_with_lock_discipline(|| { + // Signal the probe thread that the load is mid-flight + // (loading flag claimed, inner mutex released). + load_started_for_loader.wait(); + // Wait until the probe thread says it's done so the + // load's "duration" is bounded by the probes. + release_load_for_loader.wait(); + }) + .unwrap(); + }); + + // Wait until the loader is inside its slow section. + load_started.wait(); + + // Now probe `is_loaded()` and `loaded_model_id()` from this + // thread. They MUST return without contending on the lock. + let probe_deadline = Duration::from_millis(50); + let (tx, rx) = mpsc::channel(); + let engine_for_probe = engine.clone(); + let probe_handle = thread::spawn(move || { + let start = Instant::now(); + let loaded = engine_for_probe.is_loaded(); + let id = engine_for_probe.loaded_model_id(); + let loading = engine_for_probe.is_loading(); + let elapsed = start.elapsed(); + tx.send((loaded, id, loading, elapsed)).unwrap(); + }); + + let result = rx + .recv_timeout(probe_deadline) + .expect("is_loaded / loaded_model_id probe must return within 50 ms"); + let (loaded, id, loading, elapsed) = result; + + assert!( + !loaded, + "is_loaded() should report false while a load is in flight" + ); + assert_eq!(id, None, "loaded_model_id() should be None mid-load"); + assert!(loading, "is_loading() should report true mid-load"); + assert!( + elapsed < probe_deadline, + "probe took {elapsed:?}, expected < {probe_deadline:?}" + ); + + probe_handle.join().unwrap(); + release_load.wait(); + loader_handle.join().unwrap(); + + // After the load completes the flag clears. + assert!(!engine.is_loading()); + } + + /// Race-3 / Race-4 — concurrent load attempts must not both reach + /// the heavy work. The second caller should be told `AlreadyLoading` + /// rather than starting a parallel load that silently overwrites + /// the first. + #[test] + fn second_concurrent_load_is_refused() { + use std::thread; + + let engine = LlmEngine::new(); + let load_started = Arc::new(std::sync::Barrier::new(2)); + let release_load = Arc::new(std::sync::Barrier::new(2)); + + let engine_for_loader = engine.clone(); + let load_started_for_loader = Arc::clone(&load_started); + let release_load_for_loader = Arc::clone(&release_load); + + let loader_handle = thread::spawn(move || { + engine_for_loader + .__test_run_with_lock_discipline(|| { + load_started_for_loader.wait(); + release_load_for_loader.wait(); + }) + .unwrap(); + }); + + load_started.wait(); + + // Second concurrent attempt MUST be refused, not parallel-load. + let second = engine.__test_run_with_lock_discipline(|| { + panic!("second concurrent load should never reach its op"); + }); + assert!(matches!(second, Err(EngineError::AlreadyLoading))); + + release_load.wait(); + loader_handle.join().unwrap(); + + // After the first load completes, a fresh attempt is allowed. + assert!(engine + .__test_run_with_lock_discipline(|| {}) + .is_ok()); + } }