diff --git a/.changeset/add_bounded_buffer_to_audio_stream_and_use_10_frames_as_the_default.md b/.changeset/add_bounded_buffer_to_audio_stream_and_use_10_frames_as_the_default.md new file mode 100644 index 000000000..6c9731265 --- /dev/null +++ b/.changeset/add_bounded_buffer_to_audio_stream_and_use_10_frames_as_the_default.md @@ -0,0 +1,19 @@ +--- +livekit-protocol: patch +livekit-wakeword: patch +livekit-ffi: patch +livekit-api: patch +soxr-sys: patch +webrtc-sys: patch +webrtc-sys-build: patch +imgproc: patch +yuv-sys: patch +libwebrtc: patch +livekit: patch +--- + +# add bounded buffer to audio_stream, and use 10 frames as the default + +#945 by @xianshijing-lk + + diff --git a/Cargo.lock b/Cargo.lock index 5239b6235..770df4e17 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3718,6 +3718,7 @@ dependencies = [ "livekit-runtime", "log", "parking_lot", + "rtrb", "serde", "serde_json", "thiserror 1.0.69", @@ -6187,6 +6188,12 @@ dependencies = [ "tokio", ] +[[package]] +name = "rtrb" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7204ed6420f698836b76d4d5c2ec5dec7585fd5c3a788fd1cde855d1de598239" + [[package]] name = "rustc-demangle" version = "0.1.27" diff --git a/libwebrtc/Cargo.toml b/libwebrtc/Cargo.toml index 57a4c4851..0f8ec7a5d 100644 --- a/libwebrtc/Cargo.toml +++ b/libwebrtc/Cargo.toml @@ -35,6 +35,7 @@ lazy_static = { workspace = true } parking_lot = { workspace = true } tokio = { workspace = true, default-features = false, features = ["sync", "macros"] } cxx = "1.0" +rtrb = "0.3.3" [target.'cfg(target_arch = "wasm32")'.dependencies] wasm-bindgen = "0.2" diff --git a/libwebrtc/src/audio_stream.rs b/libwebrtc/src/audio_stream.rs index ae7ed89d4..454392248 100644 --- a/libwebrtc/src/audio_stream.rs +++ b/libwebrtc/src/audio_stream.rs @@ -27,6 +27,22 @@ pub mod native { use super::stream_imp; use crate::{audio_frame::AudioFrame, audio_track::RtcAudioTrack}; + const DEFAULT_QUEUE_SIZE_FRAMES: usize = 10; + + #[derive(Clone, Debug, Default)] + pub struct NativeAudioStreamOptions { + /// Maximum number of queued WebRTC sink frames after the audio callback. + /// + /// Each queued frame corresponds to roughly 10 ms of decoded PCM audio + /// on the WebRTC sink path. + /// + /// `None` uses the default bounded queue size of 10 frames. `Some(0)` + /// opts into unbounded buffering. Positive values bound the queue, and + /// the stream drops the oldest queued frames on overflow so latency + /// stays bounded. + pub queue_size_frames: Option, + } + pub struct NativeAudioStream { pub(crate) handle: stream_imp::NativeAudioStream, } @@ -40,7 +56,28 @@ pub mod native { impl NativeAudioStream { pub fn new(audio_track: RtcAudioTrack, sample_rate: i32, num_channels: i32) -> Self { Self { - handle: stream_imp::NativeAudioStream::new(audio_track, sample_rate, num_channels), + handle: stream_imp::NativeAudioStream::new( + audio_track, + sample_rate, + num_channels, + Some(DEFAULT_QUEUE_SIZE_FRAMES), + ), + } + } + + pub fn with_options( + audio_track: RtcAudioTrack, + sample_rate: i32, + num_channels: i32, + options: NativeAudioStreamOptions, + ) -> Self { + Self { + handle: stream_imp::NativeAudioStream::new( + audio_track, + sample_rate, + num_channels, + normalize_queue_size_frames(options.queue_size_frames), + ), } } @@ -60,4 +97,12 @@ pub mod native { Pin::new(&mut self.get_mut().handle).poll_next(cx) } } + + fn normalize_queue_size_frames(queue_size_frames: Option) -> Option { + match queue_size_frames { + None => Some(DEFAULT_QUEUE_SIZE_FRAMES), + Some(0) => None, + Some(value) => Some(value), + } + } } diff --git a/libwebrtc/src/native/audio_stream.rs b/libwebrtc/src/native/audio_stream.rs index 7ab18ce81..9a16b5ff9 100644 --- a/libwebrtc/src/native/audio_stream.rs +++ b/libwebrtc/src/native/audio_stream.rs @@ -13,14 +13,19 @@ // limitations under the License. use std::{ + collections::VecDeque, pin::Pin, - sync::Arc, - task::{Context, Poll}, + sync::{ + atomic::{AtomicBool, AtomicU64, Ordering}, + Arc, + }, + task::{Context, Poll, Waker}, }; use cxx::SharedPtr; use livekit_runtime::Stream; -use tokio::sync::mpsc; +use parking_lot::Mutex; +use rtrb::{Consumer, Producer, PushError, RingBuffer}; use webrtc_sys::audio_track as sys_at; use crate::{audio_frame::AudioFrame, audio_track::RtcAudioTrack}; @@ -28,13 +33,18 @@ use crate::{audio_frame::AudioFrame, audio_track::RtcAudioTrack}; pub struct NativeAudioStream { native_sink: SharedPtr, audio_track: RtcAudioTrack, - frame_rx: mpsc::UnboundedReceiver>, + frame_queue: Arc, } impl NativeAudioStream { - pub fn new(audio_track: RtcAudioTrack, sample_rate: i32, num_channels: i32) -> Self { - let (frame_tx, frame_rx) = mpsc::unbounded_channel(); - let observer = Arc::new(AudioTrackObserver { frame_tx }); + pub fn new( + audio_track: RtcAudioTrack, + sample_rate: i32, + num_channels: i32, + queue_size_frames: Option, + ) -> Self { + let frame_queue = Arc::new(AudioFrameQueue::new(queue_size_frames)); + let observer = Arc::new(AudioTrackObserver { frame_queue: frame_queue.clone() }); let native_sink = sys_at::ffi::new_native_audio_sink( Box::new(sys_at::AudioSinkWrapper::new(observer.clone())), sample_rate, @@ -44,7 +54,7 @@ impl NativeAudioStream { let audio = unsafe { sys_at::ffi::media_to_audio(audio_track.sys_handle()) }; audio.add_sink(&native_sink); - Self { native_sink, audio_track, frame_rx } + Self { native_sink, audio_track, frame_queue } } pub fn track(&self) -> RtcAudioTrack { @@ -55,7 +65,7 @@ impl NativeAudioStream { let audio = unsafe { sys_at::ffi::media_to_audio(self.audio_track.sys_handle()) }; audio.remove_sink(&self.native_sink); - self.frame_rx.close(); + self.frame_queue.close(); } } @@ -68,18 +78,18 @@ impl Drop for NativeAudioStream { impl Stream for NativeAudioStream { type Item = AudioFrame<'static>; - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll> { - self.frame_rx.poll_recv(cx) + fn poll_next(self: Pin<&mut Self>, cx: &mut Context) -> Poll> { + self.frame_queue.poll_recv(cx) } } pub struct AudioTrackObserver { - frame_tx: mpsc::UnboundedSender>, + frame_queue: Arc, } impl sys_at::AudioSink for AudioTrackObserver { fn on_data(&self, data: &[i16], sample_rate: i32, nb_channels: usize, nb_frames: usize) { - let _ = self.frame_tx.send(AudioFrame { + self.frame_queue.push(AudioFrame { data: data.to_owned().into(), sample_rate: sample_rate as u32, num_channels: nb_channels as u32, @@ -87,3 +97,218 @@ impl sys_at::AudioSink for AudioTrackObserver { }); } } + +struct AudioFrameQueue { + kind: AudioFrameQueueKind, + closed: AtomicBool, + dropped_frames: AtomicU64, + waker: Mutex>, +} + +enum AudioFrameQueueKind { + Bounded(BoundedAudioFrameQueue), + Unbounded(UnboundedAudioFrameQueue), +} + +struct BoundedAudioFrameQueue { + producer: Mutex>>, + consumer: Mutex>>, +} + +struct UnboundedAudioFrameQueue { + frames: Mutex>>, +} + +impl AudioFrameQueue { + fn new(capacity: Option) -> Self { + let kind = match capacity.filter(|capacity| *capacity > 0) { + Some(capacity) => { + let (producer, consumer) = RingBuffer::new(capacity); + AudioFrameQueueKind::Bounded(BoundedAudioFrameQueue { + producer: Mutex::new(producer), + consumer: Mutex::new(consumer), + }) + } + None => AudioFrameQueueKind::Unbounded(UnboundedAudioFrameQueue { + frames: Mutex::new(VecDeque::new()), + }), + }; + + Self { + kind, + closed: AtomicBool::new(false), + dropped_frames: AtomicU64::new(0), + waker: Mutex::new(None), + } + } + + fn push(&self, frame: AudioFrame<'static>) { + if self.closed.load(Ordering::Acquire) { + return; + } + + match &self.kind { + AudioFrameQueueKind::Bounded(queue) => self.push_bounded(queue, frame), + AudioFrameQueueKind::Unbounded(queue) => { + queue.frames.lock().push_back(frame); + } + } + + self.wake_receiver(); + } + + fn push_bounded(&self, queue: &BoundedAudioFrameQueue, mut frame: AudioFrame<'static>) { + loop { + let push_result = queue.producer.lock().push(frame); + match push_result { + Ok(()) => return, + Err(PushError::Full(returned_frame)) => { + frame = returned_frame; + + let dropped = queue.consumer.lock().pop().is_ok(); + + if dropped { + self.record_drop(); + } else { + return; + } + } + } + } + } + + fn close(&self) { + self.closed.store(true, Ordering::Release); + self.wake_receiver(); + + match &self.kind { + AudioFrameQueueKind::Bounded(queue) => { + let mut consumer = queue.consumer.lock(); + while consumer.pop().is_ok() {} + } + AudioFrameQueueKind::Unbounded(queue) => { + queue.frames.lock().clear(); + } + } + } + + fn poll_recv(&self, cx: &mut Context<'_>) -> Poll>> { + if let Some(frame) = self.try_pop() { + return Poll::Ready(Some(frame)); + } + + if self.closed.load(Ordering::Acquire) { + return Poll::Ready(None); + } + + *self.waker.lock() = Some(cx.waker().clone()); + + if let Some(frame) = self.try_pop() { + self.waker.lock().take(); + Poll::Ready(Some(frame)) + } else if self.closed.load(Ordering::Acquire) { + Poll::Ready(None) + } else { + Poll::Pending + } + } + + fn try_pop(&self) -> Option> { + match &self.kind { + AudioFrameQueueKind::Bounded(queue) => queue.consumer.lock().pop().ok(), + AudioFrameQueueKind::Unbounded(queue) => queue.frames.lock().pop_front(), + } + } + + fn wake_receiver(&self) { + let waker = self.waker.lock().take(); + if let Some(waker) = waker { + waker.wake(); + } + } + + fn record_drop(&self) { + let dropped_frames = self.dropped_frames.fetch_add(1, Ordering::Relaxed) + 1; + if dropped_frames == 1 || dropped_frames % 100 == 0 { + log::warn!( + "native audio stream queue overflow; dropped {} queued frames", + dropped_frames + ); + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::Ordering; + + use super::AudioFrameQueue; + use crate::audio_frame::AudioFrame; + + fn test_frame(marker: i16) -> AudioFrame<'static> { + AudioFrame { + data: vec![marker].into(), + sample_rate: 48_000, + num_channels: 1, + samples_per_channel: 1, + } + } + + fn pop_marker(queue: &AudioFrameQueue) -> Option { + queue.try_pop().map(|frame| frame.data[0]) + } + + #[test] + fn bounded_queue_preserves_fifo_order_under_capacity() { + let queue = AudioFrameQueue::new(Some(3)); + + queue.push(test_frame(1)); + queue.push(test_frame(2)); + queue.push(test_frame(3)); + + assert_eq!(pop_marker(&queue), Some(1)); + assert_eq!(pop_marker(&queue), Some(2)); + assert_eq!(pop_marker(&queue), Some(3)); + assert_eq!(pop_marker(&queue), None); + } + + #[test] + fn bounded_queue_drops_oldest_when_full() { + let queue = AudioFrameQueue::new(Some(2)); + + queue.push(test_frame(1)); + queue.push(test_frame(2)); + queue.push(test_frame(3)); + + assert_eq!(queue.dropped_frames.load(Ordering::Relaxed), 1); + assert_eq!(pop_marker(&queue), Some(2)); + assert_eq!(pop_marker(&queue), Some(3)); + assert_eq!(pop_marker(&queue), None); + } + + #[test] + fn unbounded_queue_retains_all_frames() { + let queue = AudioFrameQueue::new(None); + + for marker in 1..=4 { + queue.push(test_frame(marker)); + } + + for marker in 1..=4 { + assert_eq!(pop_marker(&queue), Some(marker)); + } + assert_eq!(pop_marker(&queue), None); + assert_eq!(queue.dropped_frames.load(Ordering::Relaxed), 0); + } + + #[test] + fn close_clears_buffer_and_rejects_future_pushes() { + let queue = AudioFrameQueue::new(Some(2)); + + queue.push(test_frame(1)); + queue.close(); + queue.push(test_frame(2)); + + assert_eq!(pop_marker(&queue), None); + } +} diff --git a/livekit-ffi/protocol/audio_frame.proto b/livekit-ffi/protocol/audio_frame.proto index 8597b23ed..753de7726 100644 --- a/livekit-ffi/protocol/audio_frame.proto +++ b/livekit-ffi/protocol/audio_frame.proto @@ -30,6 +30,11 @@ message NewAudioStreamRequest { optional string audio_filter_module_id = 5; // Unique identifier passed in LoadAudioFilterPluginRequest optional string audio_filter_options = 6; optional uint32 frame_size_ms = 7; + // Maximum number of queued WebRTC sink frames. Each frame is typically 10 ms + // of decoded PCM audio on the receive path. Omit this field to use the + // default bounded queue size of 10 frames. Set it to 0 to request unbounded + // buffering. + optional uint32 queue_size_frames = 8; } message NewAudioStreamResponse { required OwnedAudioStream stream = 1; } @@ -42,6 +47,11 @@ message AudioStreamFromParticipantRequest { optional string audio_filter_module_id = 7; optional string audio_filter_options = 8; optional uint32 frame_size_ms = 9; + // Maximum number of queued WebRTC sink frames. Each frame is typically 10 ms + // of decoded PCM audio on the receive path. Omit this field to use the + // default bounded queue size of 10 frames. Set it to 0 to request unbounded + // buffering. + optional uint32 queue_size_frames = 10; } message AudioStreamFromParticipantResponse { required OwnedAudioStream stream = 1; } diff --git a/livekit-ffi/src/server/audio_stream.rs b/livekit-ffi/src/server/audio_stream.rs index 1c3e4fa84..f0f343321 100644 --- a/livekit-ffi/src/server/audio_stream.rs +++ b/livekit-ffi/src/server/audio_stream.rs @@ -17,7 +17,10 @@ use std::time::Duration; use futures_util::StreamExt; use livekit::track::Track; -use livekit::webrtc::{audio_stream::native::NativeAudioStream, prelude::*}; +use livekit::webrtc::{ + audio_stream::native::{NativeAudioStream, NativeAudioStreamOptions}, + prelude::*, +}; use livekit::{registered_audio_filter_plugin, AudioFilterAudioStream, AudioFilterStreamInfo}; use tokio::sync::{broadcast, mpsc, oneshot}; @@ -97,9 +100,18 @@ impl FfiAudioStream { let audio_stream = Self { handle_id, stream_type, self_dropped_tx }; let sample_rate = new_stream.sample_rate.unwrap_or(48000); let num_channels = new_stream.num_channels.unwrap_or(1); + let options = NativeAudioStreamOptions { + queue_size_frames: new_stream + .queue_size_frames + .map(|capacity| capacity as usize), + }; - let native_stream = - NativeAudioStream::new(rtc_track, sample_rate as i32, num_channels as i32); + let native_stream = NativeAudioStream::with_options( + rtc_track, + sample_rate as i32, + num_channels as i32, + options, + ); let stream = if let Some(audio_filter) = &audio_filter { let session = audio_filter.clone().new_session( @@ -244,6 +256,9 @@ impl FfiAudioStream { let sample_rate = request.sample_rate.unwrap_or(48000) as i32; let num_channels = request.num_channels.unwrap_or(1) as i32; let track_sid = track.sid(); + let options = NativeAudioStreamOptions { + queue_size_frames: request.queue_size_frames.map(|capacity| capacity as usize), + }; let mut track_finished_rx = track_finished_tx.subscribe(); server.async_runtime.spawn(async move { @@ -292,7 +307,8 @@ impl FfiAudioStream { None => (None, None), }; - let native_stream = NativeAudioStream::new(rtc_track, sample_rate, num_channels); + let native_stream = + NativeAudioStream::with_options(rtc_track, sample_rate, num_channels, options); let stream = if let Some(session) = audio_filter_session.take() { let stream = AudioFilterAudioStream::new(