From f45ad97e516d5b89a186d1013aec0d7d7b144246 Mon Sep 17 00:00:00 2001 From: CJ Pais Date: Mon, 16 Mar 2026 08:47:00 +0800 Subject: [PATCH] ensure samples don't get dropped (#1043) * ensure samples don't get dropped * format --- src-tauri/src/audio_toolkit/audio/recorder.rs | 88 +++++++++++++++---- 1 file changed, 71 insertions(+), 17 deletions(-) diff --git a/src-tauri/src/audio_toolkit/audio/recorder.rs b/src-tauri/src/audio_toolkit/audio/recorder.rs index 3d23e74..09c2701 100644 --- a/src-tauri/src/audio_toolkit/audio/recorder.rs +++ b/src-tauri/src/audio_toolkit/audio/recorder.rs @@ -1,6 +1,9 @@ use std::{ io::Error, - sync::{mpsc, Arc, Mutex}, + sync::{ + atomic::{AtomicBool, Ordering}, + mpsc, Arc, Mutex, + }, time::Duration, }; @@ -22,6 +25,11 @@ enum Cmd { Shutdown, } +enum AudioChunk { + Samples(Vec), + EndOfStream, +} + pub struct AudioRecorder { device: Option, cmd_tx: Option>, @@ -59,7 +67,7 @@ impl AudioRecorder { return Ok(()); // already open } - let (sample_tx, sample_rx) = mpsc::channel::>(); + let (sample_tx, sample_rx) = mpsc::channel::(); let (cmd_tx, cmd_rx) = mpsc::channel::(); let (init_tx, init_rx) = mpsc::sync_channel::>(1); @@ -77,6 +85,8 @@ impl AudioRecorder { let level_cb = self.level_cb.clone(); let worker = std::thread::spawn(move || { + let stop_flag = Arc::new(AtomicBool::new(false)); + let stop_flag_for_stream = stop_flag.clone(); let init_result = (|| -> Result<(cpal::Stream, u32), String> { let config = AudioRecorder::get_preferred_config(&thread_device) .map_err(|e| format!("Failed to fetch preferred config: {e}"))?; @@ -98,6 +108,7 @@ impl AudioRecorder { &config, sample_tx, channels, + stop_flag_for_stream, ) .map_err(|e| format!("Failed to build input stream: {e}"))?, cpal::SampleFormat::I8 => AudioRecorder::build_stream::( @@ -105,6 +116,7 @@ impl AudioRecorder { &config, sample_tx, channels, + stop_flag_for_stream, ) .map_err(|e| format!("Failed to build input stream: {e}"))?, cpal::SampleFormat::I16 => AudioRecorder::build_stream::( @@ -112,6 +124,7 @@ impl AudioRecorder { &config, sample_tx, channels, + stop_flag_for_stream, ) .map_err(|e| format!("Failed to build input stream: {e}"))?, cpal::SampleFormat::I32 => AudioRecorder::build_stream::( @@ -119,6 +132,7 @@ impl AudioRecorder { &config, sample_tx, channels, + stop_flag_for_stream, ) .map_err(|e| format!("Failed to build input stream: {e}"))?, cpal::SampleFormat::F32 => AudioRecorder::build_stream::( @@ -126,6 +140,7 @@ impl AudioRecorder { &config, sample_tx, channels, + stop_flag_for_stream, ) .map_err(|e| format!("Failed to build input stream: {e}"))?, sample_format => { @@ -144,7 +159,7 @@ impl AudioRecorder { Ok((stream, sample_rate)) => { let _ = init_tx.send(Ok(())); // Keep the stream alive while we process samples. - run_consumer(sample_rate, vad, sample_rx, cmd_rx, level_cb); + run_consumer(sample_rate, vad, sample_rx, cmd_rx, level_cb, stop_flag); drop(stream); } Err(error_message) => { @@ -209,23 +224,32 @@ impl AudioRecorder { fn build_stream( device: &cpal::Device, config: &cpal::SupportedStreamConfig, - sample_tx: mpsc::Sender>, + sample_tx: mpsc::Sender, channels: usize, + stop_flag: Arc, ) -> Result where T: Sample + SizedSample + Send + 'static, f32: cpal::FromSample, { let mut output_buffer = Vec::new(); + let mut eos_sent = false; let stream_cb = move |data: &[T], _: &cpal::InputCallbackInfo| { + if stop_flag.load(Ordering::Relaxed) { + if !eos_sent { + let _ = sample_tx.send(AudioChunk::EndOfStream); + eos_sent = true; + } + return; + } + eos_sent = false; + output_buffer.clear(); if channels == 1 { - // Direct conversion without intermediate Vec output_buffer.extend(data.iter().map(|&sample| sample.to_sample::())); } else { - // Convert to mono directly let frame_count = data.len() / channels; output_buffer.reserve(frame_count); @@ -239,7 +263,10 @@ impl AudioRecorder { } } - if sample_tx.send(output_buffer.clone()).is_err() { + if sample_tx + .send(AudioChunk::Samples(output_buffer.clone())) + .is_err() + { log::error!("Failed to send samples"); } }; @@ -326,9 +353,10 @@ mod tests { fn run_consumer( in_sample_rate: u32, vad: Option>>>, - sample_rx: mpsc::Receiver>, + sample_rx: mpsc::Receiver, cmd_rx: mpsc::Receiver, level_cb: Option) + Send + Sync + 'static>>, + stop_flag: Arc, ) { let mut frame_resampler = FrameResampler::new( in_sample_rate as usize, @@ -372,11 +400,16 @@ fn run_consumer( } loop { - let raw = match sample_rx.recv() { - Ok(s) => s, + let chunk = match sample_rx.recv() { + Ok(c) => c, Err(_) => break, // stream closed }; + let raw = match chunk { + AudioChunk::Samples(s) => s, + AudioChunk::EndOfStream => continue, + }; + // ---------- spectrum processing ---------------------------------- // if let Some(buckets) = visualizer.feed(&raw) { if let Some(cb) = &level_cb { @@ -393,21 +426,35 @@ fn run_consumer( while let Ok(cmd) = cmd_rx.try_recv() { match cmd { Cmd::Start => { + stop_flag.store(false, Ordering::Relaxed); processed_samples.clear(); recording = true; - visualizer.reset(); // Reset visualization buffer + visualizer.reset(); if let Some(v) = &vad { v.lock().unwrap().reset(); } } Cmd::Stop(reply_tx) => { recording = false; + stop_flag.store(true, Ordering::Relaxed); - // Drain any audio chunks that were captured but not yet consumed - while let Ok(remaining) = sample_rx.try_recv() { - frame_resampler.push(&remaining, &mut |frame: &[f32]| { - handle_frame(frame, true, &vad, &mut processed_samples) - }); + // Drain all remaining audio until the producer confirms end-of-stream. + // The cpal callback sees the stop flag, sends EndOfStream, and goes + // silent — guaranteeing every captured sample is in the channel + // ahead of the sentinel. + loop { + match sample_rx.recv_timeout(Duration::from_secs(2)) { + Ok(AudioChunk::Samples(remaining)) => { + frame_resampler.push(&remaining, &mut |frame: &[f32]| { + handle_frame(frame, true, &vad, &mut processed_samples) + }); + } + Ok(AudioChunk::EndOfStream) => break, + Err(_) => { + log::warn!("Timed out waiting for EndOfStream from audio callback"); + break; + } + } } frame_resampler.finish(&mut |frame: &[f32]| { @@ -415,8 +462,15 @@ fn run_consumer( }); let _ = reply_tx.send(std::mem::take(&mut processed_samples)); + + // Resume the audio callback so the consumer loop can continue + // receiving chunks (important for always-on microphone mode). + stop_flag.store(false, Ordering::Relaxed); + } + Cmd::Shutdown => { + stop_flag.store(true, Ordering::Relaxed); + return; } - Cmd::Shutdown => return, } } }