rvsttd: mpsc channel for audio capture, Fixed(800) buffer, deterministic shutdown

rvsttd/src/main.rs

@@ -8,9 +8,10 @@ use std::io::{BufRead, BufReader, Write};
 use std::net::{TcpListener, TcpStream};
 use std::path::PathBuf;
 use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::mpsc;
 use std::sync::{Arc, Mutex};
 use std::thread;
-use std::time::{Duration, SystemTime, UNIX_EPOCH};
+use std::time::{SystemTime, UNIX_EPOCH};
 
 include!(concat!(env!("OUT_DIR"), "/moonshine_bindings.rs"));
 
@@ -197,25 +198,24 @@ impl Shared {
 
 struct Session {
     shared: Arc<Shared>,
-    stop_signal: Arc<AtomicBool>,
     aborted: Arc<AtomicBool>,
     transcriber: thread::JoinHandle<()>,
-    cpal_stream: cpal::Stream,
+    cpal_stream: Option<cpal::Stream>,
     stream_handle: i32,
 }
 
 impl Session {
-    fn stop(self) {
-        self.stop_signal.store(true, Ordering::SeqCst);
-        drop(self.cpal_stream);
+    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();
         self.transcriber.join().ok();
         unsafe { moonshine_free_stream(self.shared.transcriber_handle, self.stream_handle) };
     }
 
-    fn abort(self) {
-        self.stop_signal.store(true, Ordering::SeqCst);
+    fn abort(mut self) {
         self.aborted.store(true, Ordering::SeqCst);
-        drop(self.cpal_stream);
+        self.cpal_stream.take();
         self.transcriber.join().ok();
         unsafe { moonshine_free_stream(self.shared.transcriber_handle, self.stream_handle) };
     }
@@ -234,12 +234,11 @@ fn start_session(shared: Arc<Shared>) -> Option<Session> {
         return None;
     }
 
-    let audio_buf: Arc<Mutex<Vec<f32>>> = Arc::new(Mutex::new(Vec::new()));
-    let stop_signal = Arc::new(AtomicBool::new(false));
+    let (tx, rx) = mpsc::channel::<Vec<f32>>();
     let aborted = Arc::new(AtomicBool::new(false));
 
-    let cpal_stream = match start_cpal(audio_buf.clone(), stop_signal.clone()) {
-        Ok(s) => s,
+    let cpal_stream = match start_cpal(tx) {
+        Ok(s) => Some(s),
         Err(e) => {
             log(&format!("cpal failed: {}", e));
             unsafe { moonshine_free_stream(shared.transcriber_handle, stream_handle) };
@@ -248,16 +247,14 @@ 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, audio_buf, stop_signal_clone, aborted_clone, stream_handle);
+        transcriber_loop(shared_clone, rx, aborted_clone, stream_handle);
     });
 
     Some(Session {
         shared,
-        stop_signal,
         aborted,
         transcriber,
         cpal_stream,
@@ -267,8 +264,7 @@ fn start_session(shared: Arc<Shared>) -> Option<Session> {
 
 fn transcriber_loop(
     shared: Arc<Shared>,
-    audio_buf: Arc<Mutex<Vec<f32>>>,
-    stop_signal: Arc<AtomicBool>,
+    rx: mpsc::Receiver<Vec<f32>>,
     aborted: Arc<AtomicBool>,
     stream_handle: i32,
 ) {
@@ -284,17 +280,9 @@ fn transcriber_loop(
         r.log_event(&format!("session {} started", shared.session_id));
     }
 
-    while !stop_signal.load(Ordering::SeqCst) {
-        let chunk = {
-            let mut buf = audio_buf.lock().unwrap();
-            if buf.is_empty() {
-                drop(buf);
-                thread::sleep(Duration::from_millis(5));
-                continue;
-            }
-            std::mem::take(&mut *buf)
-        };
-
+    // 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);
         }
@@ -320,6 +308,8 @@ fn transcriber_loop(
         send_new_segments(&shared, t_ptr, &mut sent_ids, "P");
     }
 
+    // Channel closed — all audio has been delivered and processed.
+
     // If aborted (new session took over), skip final flush entirely
     if aborted.load(Ordering::SeqCst) {
         unsafe { moonshine_stop_stream(handle, stream_handle) };
@@ -331,27 +321,7 @@ fn transcriber_loop(
         return;
     }
 
-    // Drain remaining audio
-    let remaining = {
-        let mut buf = audio_buf.lock().unwrap();
-        std::mem::take(&mut *buf)
-    };
-
-    if let Some(ref mut r) = recorder {
-        r.add_audio(&remaining);
-    }
-
-    if !remaining.is_empty() {
-        unsafe {
-            moonshine_transcribe_add_audio_to_stream(
-                handle, stream_handle,
-                remaining.as_ptr(), remaining.len() as u64,
-                SAMPLE_RATE, 0,
-            );
-        }
-    }
-
-    // Final flush
+    // Final flush — no remaining audio to drain (channel is empty by definition)
     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) };
@@ -441,8 +411,7 @@ fn send_new_segments(
 // ─── cpal ─────────────────────────────────────────────────────────────────
 
 fn start_cpal(
-    audio_buf: Arc<Mutex<Vec<f32>>>,
-    stop_signal: Arc<AtomicBool>,
+    tx: mpsc::Sender<Vec<f32>>,
 ) -> Result<cpal::Stream> {
     let host = cpal::default_host();
     let dev = host
@@ -466,6 +435,7 @@ fn start_cpal(
         config.channels = 1;
     }
     config.sample_rate = SampleRate(16000);
+    config.buffer_size = cpal::BufferSize::Fixed(800);
 
     let err_fn = |e: cpal::StreamError| log(&format!("cpal error: {}", e));
 
@@ -473,9 +443,7 @@ fn start_cpal(
         SampleFormat::F32 => dev.build_input_stream(
             &config,
             move |data: &[f32], _: &_| {
-                if !stop_signal.load(Ordering::Relaxed) {
-                    audio_buf.lock().unwrap().extend_from_slice(data);
-                }
+                let _ = tx.send(data.to_vec());
             },
             err_fn,
             None,
@@ -483,9 +451,7 @@ fn start_cpal(
         SampleFormat::I16 => dev.build_input_stream(
             &config,
             move |data: &[i16], _: &_| {
-                if !stop_signal.load(Ordering::Relaxed) {
-                    audio_buf.lock().unwrap().extend(data.iter().map(|&x| x as f32 / 32768.0));
-                }
+                let _ = tx.send(data.iter().map(|&x| x as f32 / 32768.0).collect());
             },
             err_fn,
             None,
@@ -493,9 +459,7 @@ fn start_cpal(
         SampleFormat::U8 => dev.build_input_stream(
             &config,
             move |data: &[u8], _: &_| {
-                if !stop_signal.load(Ordering::Relaxed) {
-                    audio_buf.lock().unwrap().extend(data.iter().map(|&x| (x as f32 - 128.0) / 128.0));
-                }
+                let _ = tx.send(data.iter().map(|&x| (x as f32 - 128.0) / 128.0).collect());
             },
             err_fn,
             None,