diff --git a/src-tauri/src/commands/models.rs b/src-tauri/src/commands/models.rs index 93cdf99..df77c1f 100644 --- a/src-tauri/src/commands/models.rs +++ b/src-tauri/src/commands/models.rs @@ -142,6 +142,58 @@ pub async fn ensure_model_loaded( } let engine = engine_for_name(state, engine_name)?; + + // Race-4/5 (conf 82): TOCTOU on the "is this model loaded?" check. + // Two concurrent invocations of `ensure_model_loaded` (a quick + // double-click on "Test" + "Transcribe" was the reproducer) both + // saw `loaded_model_id != requested`, both decided "not loaded — + // load it", both spawned `load_model_from_disk` + `engine.load(…)`, + // and the second silently overwrote the first. The first's + // multi-second GPU init was wasted; worse, on a tight VRAM budget + // it doubled peak residency. Holding `state.load_serialise` across + // the check-then-load section forces them to serialise: the second + // caller acquires the mutex only after the first releases it, sees + // the now-loaded state on the re-check, and short-circuits. + let serialise = Arc::clone(&state.load_serialise); + let llm_engine = Arc::clone(&state.llm_engine); + let _serialise_guard = serialise.lock().await; + + ensure_model_loaded_locked( + engine, + model_id, + concurrent, + || llm_engine.is_loaded(), + || llm_engine.unload().map_err(|e| e.to_string()), + |id| load_model_from_disk(&id), + ) + .await +} + +/// Body of `ensure_model_loaded` extracted from the Tauri-state +/// plumbing so it can be unit-tested. The caller is responsible for +/// holding `AppState::load_serialise` for the duration of this call. +/// +/// `llm_is_loaded` / `llm_unload` carry the sequential-GPU guard (brief +/// item A.1 #28) without forcing the test to construct a real +/// `LlmEngine`. `load_model_from_disk` is parameterised so tests can +/// inject a counting / sleepy fake. +async fn ensure_model_loaded_locked( + engine: Arc, + model_id: ModelId, + concurrent: Option, + llm_is_loaded: IsLoaded, + llm_unload: Unload, + load_model_from_disk: Loader, +) -> Result<(), String> +where + IsLoaded: Fn() -> bool, + Unload: FnOnce() -> Result<(), String>, + Loader: FnOnce(ModelId) -> Result, String> + Send + 'static, +{ + // Re-check loaded-id INSIDE the serialise lock: the previous + // holder of `load_serialise` may have just installed the same + // model, in which case we short-circuit instead of doing a + // redundant multi-second load. if engine .loaded_model_id() .as_ref() @@ -156,14 +208,14 @@ pub async fn ensure_model_loaded( // transcription engine on. None / Some(true) leaves the LLM // untouched (legacy parallel behaviour, safe on multi-GB VRAM // setups). Inverse guard lives in commands::llm::load_llm_model. - if concurrent == Some(false) && state.llm_engine.is_loaded() { - state.llm_engine.unload().map_err(|e| e.to_string())?; + if concurrent == Some(false) && llm_is_loaded() { + llm_unload()?; } let engine_clone = engine.clone(); let model_id_clone = model_id.clone(); tokio::task::spawn_blocking(move || { - let model = load_model_from_disk(&model_id_clone)?; + let model = load_model_from_disk(model_id_clone.clone())?; engine_clone.load(model, model_id_clone); Ok::<_, String>(()) }) @@ -699,4 +751,147 @@ mod tests { assert_eq!(full.first(), Some(&"cpu".to_string())); } } + + // ----------------------------------------------------------------- + // Race-4/5 — `ensure_model_loaded` TOCTOU + concurrent loads. + // ----------------------------------------------------------------- + + use lumotia_core::error::Result as CoreResult; + use lumotia_core::types::{EngineName, Segment}; + use lumotia_transcription::{LocalEngine, Transcriber, TranscriberCapabilities}; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + use std::time::Duration; + + /// Minimal `Transcriber` impl for tests. Carries no state; the + /// `transcribe_sync*` methods panic so the test fails loudly if + /// anyone accidentally tries to run inference on the fake. + struct FakeTranscriber; + + impl Transcriber for FakeTranscriber { + fn capabilities(&self) -> TranscriberCapabilities { + TranscriberCapabilities { + sample_rate: 16_000, + channels: 1, + supports_initial_prompt: false, + } + } + fn transcribe_sync( + &mut self, + _samples: &[f32], + _options: &TranscriptionOptions, + ) -> CoreResult> { + unreachable!("fake transcriber should never run inference"); + } + fn transcribe_sync_with_abort( + &mut self, + _samples: &[f32], + _options: &TranscriptionOptions, + _abort_flag: Arc, + ) -> CoreResult> { + unreachable!("fake transcriber should never run inference"); + } + } + + /// Race-4/5 regression. Two concurrent invocations that both + /// serialise on the same `tokio::sync::Mutex` must invoke the heavy + /// `load_model_from_disk` exactly once. The second caller acquires + /// the lock after the first releases it, observes the now-loaded + /// state inside the lock, and short-circuits. + /// + /// Pre-fix (no serialise lock, TOCTOU between `loaded_model_id` and + /// `engine.load`): both spawn_blocking tasks would run the loader, + /// `counter` would land at 2, and the second's result would + /// silently overwrite the first. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn concurrent_ensure_model_loaded_invokes_loader_once() { + let engine = Arc::new(LocalEngine::new(EngineName::new("whisper"))); + let serialise = Arc::new(tokio::sync::Mutex::new(())); + let counter = Arc::new(AtomicUsize::new(0)); + let model_id = ModelId::new("fake-model-for-test".to_string()); + + // Mimics the production call site: acquire serialise, then run + // the locked body. The fake loader sleeps briefly so the second + // caller really has a chance to race the first. + let run_once = |engine: Arc, + serialise: Arc>, + counter: Arc, + model_id: ModelId| async move { + let _g = serialise.lock().await; + ensure_model_loaded_locked( + engine, + model_id, + None, + || false, + || Ok(()), + move |_id| { + // Simulate a multi-millisecond disk + GPU init. + std::thread::sleep(Duration::from_millis(50)); + counter.fetch_add(1, Ordering::SeqCst); + Ok(Box::new(FakeTranscriber) as Box) + }, + ) + .await + }; + + let t1 = tokio::spawn(run_once( + Arc::clone(&engine), + Arc::clone(&serialise), + Arc::clone(&counter), + model_id.clone(), + )); + let t2 = tokio::spawn(run_once( + Arc::clone(&engine), + Arc::clone(&serialise), + Arc::clone(&counter), + model_id.clone(), + )); + + let (r1, r2) = tokio::join!(t1, t2); + r1.unwrap().expect("first ensure_model_loaded must succeed"); + r2.unwrap() + .expect("second ensure_model_loaded must succeed"); + + // The point: exactly ONE load happened, not two. + assert_eq!( + counter.load(Ordering::SeqCst), + 1, + "two concurrent ensure_model_loaded calls must coalesce into one load" + ); + assert_eq!(engine.loaded_model_id().as_ref(), Some(&model_id)); + } + + /// Even without serialise contention, the locked body must + /// short-circuit on the inside-lock re-check when the engine + /// already has the requested model resident. Guards against a + /// regression where the inner re-check is dropped and the lock + /// merely throttles parallel loads instead of preventing them. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn ensure_model_loaded_short_circuits_when_already_loaded() { + let engine = Arc::new(LocalEngine::new(EngineName::new("whisper"))); + let model_id = ModelId::new("fake-model-for-test".to_string()); + engine.load(Box::new(FakeTranscriber), model_id.clone()); + assert_eq!(engine.loaded_model_id().as_ref(), Some(&model_id)); + + let counter = Arc::new(AtomicUsize::new(0)); + let counter_for_loader = Arc::clone(&counter); + let result = ensure_model_loaded_locked( + Arc::clone(&engine), + model_id.clone(), + None, + || false, + || Ok(()), + move |_id| { + counter_for_loader.fetch_add(1, Ordering::SeqCst); + Ok(Box::new(FakeTranscriber) as Box) + }, + ) + .await; + + result.expect("already-loaded model should short-circuit"); + assert_eq!( + counter.load(Ordering::SeqCst), + 0, + "loader must not run when the requested model is already resident" + ); + } } diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 3bef58d..21f3b4a 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -55,6 +55,16 @@ pub struct AppState { pub parakeet_engine: Arc, pub db: SqlitePool, pub llm_engine: Arc, + /// Race-4/5 TOCTOU guard: two concurrent `ensure_model_loaded` + /// callers (e.g. a double-clicked "Test" + "Transcribe" button) + /// would both observe `loaded_model_id != requested`, both decide + /// "not loaded — load it", and both spawn a multi-second + /// `load_model_from_disk` whose result the second overwrites. + /// Holding this `tokio::sync::Mutex` across the check-then-load + /// critical section forces them to serialise; the second caller + /// acquires the lock after the first finishes, observes the + /// now-loaded state, and short-circuits. + pub load_serialise: Arc>, } /// Holds the preferences init script for injection into secondary windows. @@ -629,6 +639,7 @@ pub fn run() { parakeet_engine: Arc::new(LocalEngine::new(EngineName::new("parakeet"))), db, llm_engine: Arc::new(LlmEngine::new()), + load_serialise: Arc::new(tokio::sync::Mutex::new(())), }); // Runtime-warning banner: push CPU-feature + Vulkan-loader