rvsttd: 100ms trailing audio drain before final flush, cpal stays alive during drain

rvsttd/src/main.rs

@@ -11,7 +11,7 @@ use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::mpsc;
 use std::sync::{Arc, Mutex};
 use std::thread;
-use std::time::{SystemTime, UNIX_EPOCH};
+use std::time::{Duration, SystemTime, UNIX_EPOCH};
 
 include!(concat!(env!("OUT_DIR"), "/moonshine_bindings.rs"));
 
@@ -198,6 +198,7 @@ impl Shared {
 
 struct Session {
     shared: Arc<Shared>,
+    stop_signal: Arc<AtomicBool>,
     aborted: Arc<AtomicBool>,
     transcriber: thread::JoinHandle<()>,
     cpal_stream: Option<cpal::Stream>,
@@ -206,17 +207,20 @@ struct Session {
 
 impl Session {
     fn stop(mut self) {
-        // Drop cpal stream first — joins the callback thread, which drops
-        // the Sender, which unblocks the transcriber's recv().
-        self.cpal_stream.take();
+        // Signal transcriber to exit main loop, then wait for it to drain
+        // trailing audio + final flush. cpal stream stays alive during drain.
+        self.stop_signal.store(true, Ordering::SeqCst);
         self.transcriber.join().ok();
+        // Now safe to kill ALSA — transcriber is done
+        self.cpal_stream.take();
         unsafe { moonshine_free_stream(self.shared.transcriber_handle, self.stream_handle) };
     }
 
     fn abort(mut self) {
         self.aborted.store(true, Ordering::SeqCst);
-        self.cpal_stream.take();
+        self.stop_signal.store(true, Ordering::SeqCst);
         self.transcriber.join().ok();
+        self.cpal_stream.take();
         unsafe { moonshine_free_stream(self.shared.transcriber_handle, self.stream_handle) };
     }
 }
@@ -235,6 +239,7 @@ fn start_session(shared: Arc<Shared>) -> Option<Session> {
     }
 
     let (tx, rx) = mpsc::channel::<Vec<f32>>();
+    let stop_signal = Arc::new(AtomicBool::new(false));
     let aborted = Arc::new(AtomicBool::new(false));
 
     let cpal_stream = match start_cpal(tx) {
@@ -247,14 +252,16 @@ fn start_session(shared: Arc<Shared>) -> Option<Session> {
     };
 
     let shared_clone = shared.clone();
+    let stop_signal_clone = stop_signal.clone();
     let aborted_clone = aborted.clone();
 
     let transcriber = thread::spawn(move || {
-        transcriber_loop(shared_clone, rx, aborted_clone, stream_handle);
+        transcriber_loop(shared_clone, rx, stop_signal_clone, aborted_clone, stream_handle);
     });
 
     Some(Session {
         shared,
+        stop_signal,
         aborted,
         transcriber,
         cpal_stream,
@@ -265,6 +272,7 @@ fn start_session(shared: Arc<Shared>) -> Option<Session> {
 fn transcriber_loop(
     shared: Arc<Shared>,
     rx: mpsc::Receiver<Vec<f32>>,
+    stop_signal: Arc<AtomicBool>,
     aborted: Arc<AtomicBool>,
     stream_handle: i32,
 ) {
@@ -280,35 +288,61 @@ fn transcriber_loop(
         r.log_event(&format!("session {} started", shared.session_id));
     }
 
-    // Drain audio from the channel. When cpal stream is dropped, the Sender
-    // is dropped, recv() returns Err, and we exit the loop deterministically.
-    while let Ok(chunk) = rx.recv() {
-        if let Some(ref mut r) = recorder {
-            r.add_audio(&chunk);
-        }
+    let drain_timeout = Duration::from_millis(100);
 
-        unsafe {
-            moonshine_transcribe_add_audio_to_stream(
-                handle, stream_handle,
-                chunk.as_ptr(), chunk.len() as u64,
-                SAMPLE_RATE, 0,
-            );
-        }
+    // Main loop: process audio until stop_signal
+    while !stop_signal.load(Ordering::SeqCst) {
+        match rx.recv_timeout(drain_timeout) {
+            Ok(chunk) => {
+                if let Some(ref mut r) = recorder {
+                    r.add_audio(&chunk);
+                }
 
-        let mut t_ptr: *mut transcript_t = std::ptr::null_mut();
-        let rc = unsafe { moonshine_transcribe_stream(handle, stream_handle, 0, &mut t_ptr) };
-        if rc != 0 || t_ptr.is_null() {
-            continue;
-        }
+                unsafe {
+                    moonshine_transcribe_add_audio_to_stream(
+                        handle, stream_handle,
+                        chunk.as_ptr(), chunk.len() as u64,
+                        SAMPLE_RATE, 0,
+                    );
+                }
 
-        if let Some(ref mut r) = recorder {
-            log_transcript_lines(r, t_ptr);
-        }
+                let mut t_ptr: *mut transcript_t = std::ptr::null_mut();
+                let rc = unsafe { moonshine_transcribe_stream(handle, stream_handle, 0, &mut t_ptr) };
+                if rc != 0 || t_ptr.is_null() {
+                    continue;
+                }
 
-        send_new_segments(&shared, t_ptr, &mut sent_ids, "P");
+                if let Some(ref mut r) = recorder {
+                    log_transcript_lines(r, t_ptr);
+                }
+
+                send_new_segments(&shared, t_ptr, &mut sent_ids, "P");
+            }
+            Err(mpsc::RecvTimeoutError::Timeout) => continue,
+            Err(mpsc::RecvTimeoutError::Disconnected) => break,
+        }
     }
 
-    // Channel closed — all audio has been delivered and processed.
+    // Drain trailing audio from ALSA buffer (cpal stream still alive)
+    loop {
+        match rx.recv_timeout(drain_timeout) {
+            Ok(chunk) => {
+                if let Some(ref mut r) = recorder {
+                    r.add_audio(&chunk);
+                }
+
+                unsafe {
+                    moonshine_transcribe_add_audio_to_stream(
+                        handle, stream_handle,
+                        chunk.as_ptr(), chunk.len() as u64,
+                        SAMPLE_RATE, 0,
+                    );
+                }
+            }
+            Err(mpsc::RecvTimeoutError::Timeout) => break,
+            Err(mpsc::RecvTimeoutError::Disconnected) => break,
+        }
+    }
 
     // If aborted (new session took over), skip final flush entirely
     if aborted.load(Ordering::SeqCst) {
@@ -321,7 +355,7 @@ fn transcriber_loop(
         return;
     }
 
-    // Final flush — no remaining audio to drain (channel is empty by definition)
+    // Final flush
     unsafe { moonshine_stop_stream(handle, stream_handle) };
     let mut t_ptr: *mut transcript_t = std::ptr::null_mut();
     let rc = unsafe { moonshine_transcribe_stream(handle, stream_handle, 0, &mut t_ptr) };