ensure samples don't get dropped (#1043)
* ensure samples don't get dropped * format
This commit is contained in:
CJ Pais
GitHub
parent
184bdf4586
commit
f45ad97e51
1 changed files with 71 additions and 17 deletions
|
|
@ -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<f32>),
|
||||
EndOfStream,
|
||||
}
|
||||
|
||||
pub struct AudioRecorder {
|
||||
device: Option<Device>,
|
||||
cmd_tx: Option<mpsc::Sender<Cmd>>,
|
||||
|
|
@ -59,7 +67,7 @@ impl AudioRecorder {
|
|||
return Ok(()); // already open
|
||||
}
|
||||
|
||||
let (sample_tx, sample_rx) = mpsc::channel::<Vec<f32>>();
|
||||
let (sample_tx, sample_rx) = mpsc::channel::<AudioChunk>();
|
||||
let (cmd_tx, cmd_rx) = mpsc::channel::<Cmd>();
|
||||
let (init_tx, init_rx) = mpsc::sync_channel::<Result<(), String>>(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::<i8>(
|
||||
|
|
@ -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::<i16>(
|
||||
|
|
@ -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::<i32>(
|
||||
|
|
@ -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::<f32>(
|
||||
|
|
@ -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<T>(
|
||||
device: &cpal::Device,
|
||||
config: &cpal::SupportedStreamConfig,
|
||||
sample_tx: mpsc::Sender<Vec<f32>>,
|
||||
sample_tx: mpsc::Sender<AudioChunk>,
|
||||
channels: usize,
|
||||
stop_flag: Arc<AtomicBool>,
|
||||
) -> Result<cpal::Stream, cpal::BuildStreamError>
|
||||
where
|
||||
T: Sample + SizedSample + Send + 'static,
|
||||
f32: cpal::FromSample<T>,
|
||||
{
|
||||
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::<f32>()));
|
||||
} 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<Arc<Mutex<Box<dyn vad::VoiceActivityDetector>>>>,
|
||||
sample_rx: mpsc::Receiver<Vec<f32>>,
|
||||
sample_rx: mpsc::Receiver<AudioChunk>,
|
||||
cmd_rx: mpsc::Receiver<Cmd>,
|
||||
level_cb: Option<Arc<dyn Fn(Vec<f32>) + Send + Sync + 'static>>,
|
||||
stop_flag: Arc<AtomicBool>,
|
||||
) {
|
||||
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,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue