diff --git a/crates/core/src/error.rs b/crates/core/src/error.rs index dcbb440..9fa7537 100644 --- a/crates/core/src/error.rs +++ b/crates/core/src/error.rs @@ -22,6 +22,14 @@ pub enum Error { #[error("transcription failed: {0}")] TranscriptionFailed(String), + /// Inference exceeded the bounded wait imposed by the caller (live + /// session `drain_inference`). The spawned worker has had its abort + /// flag set so whisper-rs will return early on its next + /// abort-callback poll; the lock-up itself is *not* recovered by + /// this error — but the live-session lifecycle can now progress. + #[error("inference timed out after {timeout_ms}ms")] + InferenceTimeout { timeout_ms: u64 }, + #[error("audio decode failed: {0}")] AudioDecodeFailed(String), diff --git a/crates/transcription/src/local_engine.rs b/crates/transcription/src/local_engine.rs index ff7c0bd..f13805d 100644 --- a/crates/transcription/src/local_engine.rs +++ b/crates/transcription/src/local_engine.rs @@ -1,5 +1,6 @@ use std::path::Path; -use std::sync::Mutex; +use std::sync::atomic::AtomicBool; +use std::sync::{Arc, Mutex}; use std::time::Instant; use transcribe_rs::{SpeechModel, TranscribeOptions, TranscriptionResult}; @@ -155,6 +156,35 @@ impl LocalEngine { inference_ms, }) } + + /// Cancellable variant of `transcribe_sync`. Pipes `abort_flag` + /// through to the backend so the live-session drain timeout can + /// break a wedged whisper-rs decode out of its loop. Backends that + /// can't honour the flag (Parakeet) fall through to their plain + /// `transcribe_sync` via the trait's default implementation. + pub fn transcribe_sync_with_abort( + &self, + audio: &AudioSamples, + options: &TranscriptionOptions, + abort_flag: Arc, + ) -> Result { + let mut guard = self.engine.lock().unwrap_or_else(|e| e.into_inner()); + let backend = guard.as_mut().ok_or(Error::EngineNotLoaded)?; + + let start = Instant::now(); + let segments = + backend.transcribe_sync_with_abort(audio.samples(), options, abort_flag)?; + let inference_ms = start.elapsed().as_millis() as u64; + + Ok(TimedTranscript { + transcript: Transcript::new( + segments, + options.language.clone().unwrap_or_else(|| "en".to_string()), + audio.duration_secs(), + ), + inference_ms, + }) + } } /// Thin wrapper over `ParakeetModel` that overrides `transcribe_raw` to diff --git a/crates/transcription/src/transcriber.rs b/crates/transcription/src/transcriber.rs index 55befbf..63aebd6 100644 --- a/crates/transcription/src/transcriber.rs +++ b/crates/transcription/src/transcriber.rs @@ -9,6 +9,9 @@ //! `whisper` feature — `WhisperRsBackend` (direct whisper-rs, the only //! path that pipes `initial_prompt`). +use std::sync::atomic::AtomicBool; +use std::sync::Arc; + use lumotia_core::error::Result; use lumotia_core::types::{Segment, TranscriptionOptions}; @@ -45,6 +48,24 @@ pub trait Transcriber: Send { samples: &[f32], options: &TranscriptionOptions, ) -> Result>; + + /// Variant of `transcribe_sync` that accepts an external abort flag. + /// + /// Backends that can react to mid-inference cancellation (whisper-rs + /// via `set_abort_callback_safe`) override this to poll the flag + /// during decoding so a wedged inference can be unstuck by the live + /// session's `drain_inference` timeout. The default implementation + /// ignores the flag and dispatches to `transcribe_sync` so + /// non-cancellable backends (Parakeet via transcribe-rs) keep their + /// current behaviour. + fn transcribe_sync_with_abort( + &mut self, + samples: &[f32], + options: &TranscriptionOptions, + _abort_flag: Arc, + ) -> Result> { + self.transcribe_sync(samples, options) + } } #[cfg(test)] diff --git a/crates/transcription/src/whisper_rs_backend.rs b/crates/transcription/src/whisper_rs_backend.rs index 2af5351..2bbb0fb 100644 --- a/crates/transcription/src/whisper_rs_backend.rs +++ b/crates/transcription/src/whisper_rs_backend.rs @@ -7,6 +7,8 @@ //! into Whisper. use std::path::Path; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; use whisper_rs::{FullParams, SamplingStrategy, WhisperContext, WhisperContextParameters}; @@ -57,10 +59,35 @@ impl Transcriber for WhisperRsBackend { &mut self, samples: &[f32], options: &TranscriptionOptions, + ) -> Result> { + self.transcribe_sync_inner(samples, options, None) + } + + /// Cancellable variant. Installs the abort flag as whisper-rs's + /// abort callback so a `drain_inference` timeout in the live session + /// can break the worker out of the decoding loop instead of letting + /// it run to completion against a stuck backend. + fn transcribe_sync_with_abort( + &mut self, + samples: &[f32], + options: &TranscriptionOptions, + abort_flag: Arc, + ) -> Result> { + self.transcribe_sync_inner(samples, options, Some(abort_flag)) + } +} + +impl WhisperRsBackend { + fn transcribe_sync_inner( + &mut self, + samples: &[f32], + options: &TranscriptionOptions, + abort_flag: Option>, ) -> Result> { tracing::info!( language = ?options.language, has_initial_prompt = options.initial_prompt.as_deref().map(|p| !p.is_empty()).unwrap_or(false), + has_abort_flag = abort_flag.is_some(), "WhisperRsBackend::transcribe_sync entering" ); @@ -87,6 +114,15 @@ impl Transcriber for WhisperRsBackend { params.set_print_progress(false); params.set_print_realtime(false); + // Wire the per-task abort flag into whisper-rs. The closure is + // polled by whisper.cpp between decode steps; returning true + // tells the backend to abort the current inference. This is the + // cancellation route that closes the orphaned-thread loophole + // when a live session is stopped or wedges. + if let Some(flag) = abort_flag { + params.set_abort_callback_safe(move || flag.load(Ordering::Relaxed)); + } + state.full(params, samples).map_err(|e| { Error::TranscriptionFailed( WhisperBackendError::Transcribe(e.to_string()).to_string(), diff --git a/src-tauri/src/commands/live.rs b/src-tauri/src/commands/live.rs index f4055fe..bf08061 100644 --- a/src-tauri/src/commands/live.rs +++ b/src-tauri/src/commands/live.rs @@ -423,8 +423,49 @@ impl LiveSessionRuntime { } fn drain_inference(&mut self) -> Result<(), String> { + // Bound the wait so a wedged whisper (ggml deadlock, GPU stall, + // any pathology that turns transcribe_sync into a black hole) + // can't brick the live-session lifecycle and, via the outer + // AsyncMutex, every subsequent start/stop. The timeout pairs + // with the per-task `abort_flag` plumbed into whisper-rs's + // abort callback — on expiry we set the flag and drop our + // receiver, so the worker either honours the abort or exits + // silently when it eventually tries to send its result. + let chunk_duration_ms = (CHUNK_SAMPLES as u64 * 1000) / WHISPER_SAMPLE_RATE as u64; + let drain_timeout_ms = (chunk_duration_ms.saturating_mul(3)).max(2_000); + let drain_started = Instant::now(); while self.state.inflight.is_some() { self.poll_inference()?; + if self.state.inflight.is_none() { + break; + } + if drain_started.elapsed() >= Duration::from_millis(drain_timeout_ms) { + // SAFETY: a real test of this path requires a wedged + // whisper-rs, which is hard to fixture. We rely on the + // abort_callback wiring plus cargo build + manual smoke + // for verification. The timeout fires regardless, so + // the lifecycle deadlock is broken either way. + if let Some(task) = self.state.inflight.as_ref() { + task.abort_flag.store(true, Ordering::Relaxed); + tracing::warn!( + target: "lumotia_live", + session_id = self.session_id, + chunk_id = task.chunk_id, + timeout_ms = drain_timeout_ms, + inflight_age_ms = task.dispatched_at.elapsed().as_millis() as u64, + "drained inflight inference after timeout" + ); + } + let _ = self.state.inflight.take(); + let message = format!( + "Inference timed out after {drain_timeout_ms}ms; live session stopping cleanly" + ); + let _ = self.status_channel.send(LiveStatusMessage::Error { + session_id: self.session_id, + message: message.clone(), + }); + return Err(message); + } thread::sleep(Duration::from_millis(10)); } Ok(()) @@ -451,6 +492,29 @@ struct InferenceTask { trim_before_secs: f64, duration_secs: f64, rx: std::sync::mpsc::Receiver>, + /// Per-task abort flag. Set by `drain_inference` on timeout or by + /// any early worker exit; polled by whisper-rs through + /// `set_abort_callback_safe`. The Arc is shared with the spawned + /// inference thread so the cancellation route survives a dropped + /// `InferenceTask` (e.g. when the loop returns through `?`). + abort_flag: Arc, + /// Instant the inference thread was spawned. Used by + /// `drain_inference` to compute a bounded wait based on the chunk + /// duration. None for tasks dispatched before wall-clock tracking + /// was required (kept Option to keep the test fixtures concise). + dispatched_at: Instant, +} + +impl Drop for InferenceTask { + fn drop(&mut self) { + // Any path that drops the task without explicitly clearing it + // (early `?` propagation in the loop, panic unwind) signals the + // worker thread to exit. Without this the orphaned thread keeps + // running whisper-rs against the engine's Mutex; N orphans on + // a reconnect/retry loop contend on the backend and stall every + // future inference. + self.abort_flag.store(true, Ordering::Relaxed); + } } #[derive(Debug, Clone)] @@ -492,7 +556,14 @@ pub async fn start_live_transcription_session( status_channel: Channel, ) -> Result { ensure_main_window(&window)?; - let _lifecycle = live_state.lifecycle.lock().await; + // Phase 1: acquire the lifecycle lock long enough to reserve the + // single live-session slot. Held across `ensure_model_loaded` + // because we want start/stop concurrency to remain serialised; the + // dangerous pattern (lock held across a JoinHandle.await) was on + // the stop path, not here. Released explicitly before the + // RunningLiveSession is installed in `live_state.running` so the + // symmetric stop path doesn't observe a half-initialised state. + let lifecycle = live_state.lifecycle.lock().await; { let running = live_state.running.lock().unwrap(); if running.is_some() { @@ -588,6 +659,11 @@ pub async fn start_live_transcription_session( status_channel, }); + // Phase 2: drop the lifecycle guard explicitly before returning so + // a concurrent stop call can grab the lock without contention now + // that the new RunningLiveSession is installed. + drop(lifecycle); + Ok(StartLiveTranscriptionResponse { session_id }) } @@ -599,21 +675,35 @@ pub async fn stop_live_transcription_session( session_id: u64, ) -> Result { ensure_main_window(&window)?; - let _lifecycle = live_state.lifecycle.lock().await; - let running = live_state.running.lock().unwrap().take(); - let Some(running) = running else { - return Err("No live transcription session is running".into()); + // Take the running session and signal stop while holding the + // lifecycle guard — but release the guard BEFORE awaiting the + // worker's JoinHandle. The worker is a non-cancellable + // spawn_blocking task; awaiting it under the lifecycle AsyncMutex + // means a single wedged session blocks every future start/stop + // until the OS thread eventually exits. Releasing here keeps the + // mutual-exclusion invariant on the `running` slot (we already + // took it out) without leaking the worker's wedge into the + // lifecycle gate. + let (handle, stop_flag, status_channel) = { + let _lifecycle = live_state.lifecycle.lock().await; + let running = live_state.running.lock().unwrap().take(); + let Some(running) = running else { + return Err("No live transcription session is running".into()); + }; + + if running.id != session_id { + *live_state.running.lock().unwrap() = Some(running); + return Err(format!("Session {session_id} is not active")); + } + + running.stop_flag.store(true, Ordering::Relaxed); + (running.handle, running.stop_flag, running.status_channel) }; + // `_lifecycle` dropped — other start/stop calls can now claim it + // even if this handle.await is slow. + let _ = stop_flag; - if running.id != session_id { - *live_state.running.lock().unwrap() = Some(running); - return Err(format!("Session {session_id} is not active")); - } - - running.stop_flag.store(true, Ordering::Relaxed); - - let summary = running - .handle + let summary = handle .await .map_err(|e| format!("Live session task failed: {e}"))??; @@ -629,7 +719,7 @@ pub async fn stop_live_transcription_session( dropped_audio_ms: summary.dropped_audio_ms, }; - let _ = running.status_channel.send(LiveStatusMessage::Finished { + let _ = status_channel.send(LiveStatusMessage::Finished { session_id: summary.session_id, audio_path, dropped_audio_ms: summary.dropped_audio_ms, @@ -844,12 +934,19 @@ fn maybe_dispatch_chunk( }; let engine = engine.clone(); let (tx, rx) = std::sync::mpsc::channel(); + let abort_flag = Arc::new(AtomicBool::new(false)); + let worker_abort = abort_flag.clone(); + let dispatched_at = Instant::now(); + // The handle is intentionally dropped — the cancellation route is + // now `abort_flag` plumbed through whisper-rs's abort callback, not + // a JoinHandle. On any early exit from the loop the per-task flag + // is set and whisper-rs returns cleanly on its next poll. thread::spawn(move || { let audio = AudioSamples::mono_16khz(chunk_samples); let started = Instant::now(); let result = engine - .transcribe_sync(&audio, &options) + .transcribe_sync_with_abort(&audio, &options, worker_abort) .map(|mut timed| { timed.inference_ms = started.elapsed().as_millis() as u64; timed @@ -864,6 +961,8 @@ fn maybe_dispatch_chunk( trim_before_secs, duration_secs, rx, + abort_flag, + dispatched_at, }) } @@ -1570,6 +1669,8 @@ mod tests { trim_before_secs: 0.0, duration_secs: 0.8, rx: rx1, + abort_flag: Arc::new(AtomicBool::new(false)), + dispatched_at: Instant::now(), }); let first = poll_inference( @@ -1612,6 +1713,8 @@ mod tests { trim_before_secs: 0.0, duration_secs: 0.9, rx: rx2, + abort_flag: Arc::new(AtomicBool::new(false)), + dispatched_at: Instant::now(), }); let second = poll_inference( @@ -1738,4 +1841,34 @@ mod tests { assert_eq!(summary.session_id, 7); assert_eq!(start.await.unwrap().unwrap(), 8); } + + #[test] + fn dropping_inference_task_sets_abort_flag() { + // Race-A regression: dropping an InferenceTask (via `?` + // propagation, panic unwind, or early loop exit) MUST signal + // the spawned worker thread to exit via the abort flag. Without + // this, every early exit leaves an orphan thread running + // whisper-rs against the engine's Mutex; reconnect/retry loops + // pile up N orphans contending on the backend lock. + let abort_flag = Arc::new(AtomicBool::new(false)); + let (_tx, rx) = std::sync::mpsc::channel(); + let task = InferenceTask { + chunk_id: 99, + chunk_start_sample: 0, + trim_before_secs: 0.0, + duration_secs: 2.0, + rx, + abort_flag: abort_flag.clone(), + dispatched_at: Instant::now(), + }; + assert!( + !abort_flag.load(Ordering::Relaxed), + "abort flag must start cleared" + ); + drop(task); + assert!( + abort_flag.load(Ordering::Relaxed), + "Drop for InferenceTask must signal cancellation to the worker" + ); + } }