Fix buffered AudioSource.capture_frame hanging when the source drain stalls

.changeset/fix-audiosource-capture-frame-drain-stall.md

@@ -0,0 +1,8 @@
+---
+libwebrtc: patch
+livekit: patch
+livekit-ffi: patch
+webrtc-sys: patch
+---
+
+Bound the buffered `AudioSource::capture_frame` completion wait and split the drain lock so a stalled or wedged source drain returns a recoverable error instead of hanging the producer (and the session) forever (#408, #420, #497) - #1289 (@sam-hark)

libwebrtc/Cargo.toml

@@ -59,3 +59,4 @@ web-sys = { version = "0.3", features = [
 
 [dev-dependencies]
 env_logger = { workspace = true }
+tokio = { workspace = true, features = ["rt-multi-thread", "macros", "time", "sync"] }

libwebrtc/src/native/audio_source.rs

@@ -18,6 +18,14 @@ use webrtc_sys::audio_track as sys_at;
 
 use crate::{audio_frame::AudioFrame, audio_source::AudioSourceOptions, RtcError, RtcErrorType};
 
+/// Floor for the per-chunk completion wait, for very small queues.
+const CAPTURE_COMPLETE_TIMEOUT_FLOOR: std::time::Duration = std::time::Duration::from_secs(1);
+
+/// Headroom multiple applied to the queue duration when bounding the completion wait: a healthy
+/// drain acknowledges a chunk within one queue-duration, so this leaves margin before the wait is
+/// treated as a stall.
+const CAPTURE_COMPLETE_TIMEOUT_QUEUE_MULTIPLE: u64 = 5;
+
 #[derive(Clone)]
 pub struct NativeAudioSource {
     sys_handle: SharedPtr<sys_at::ffi::AudioTrackSource>,
@@ -153,6 +161,20 @@ impl NativeAudioSource {
             let _ = tx.send(());
         }
 
+        // Bound the per-chunk completion wait on a multiple of the queue duration. The C++
+        // drain acknowledges a chunk from a 10ms RepeatingTask; legitimate backpressure clears
+        // within one queue-duration. A wait beyond that means the drain stalled (CPU-starved,
+        // or a sink blocked in OnData), so surface a recoverable error rather than hang the
+        // caller — and the whole session — forever (see rust-sdks #408 / #420 / #497).
+        let queue_ms = (self.queue_size_samples as u64)
+            .saturating_mul(1000)
+            .checked_div((self.sample_rate as u64) * (self.num_channels as u64))
+            .unwrap_or(0);
+        let capture_timeout = std::time::Duration::from_millis(
+            queue_ms.saturating_mul(CAPTURE_COMPLETE_TIMEOUT_QUEUE_MULTIPLE),
+        )
+        .max(CAPTURE_COMPLETE_TIMEOUT_FLOOR);
+
         // iterate over chunks of self._queue_size_samples
         for chunk in frame.data.chunks(self.queue_size_samples as usize) {
             let nb_frames = chunk.len() / self.num_channels as usize;
@@ -161,7 +183,9 @@ impl NativeAudioSource {
             let ctx_ptr = Box::into_raw(ctx) as *const sys_at::SourceContext;
 
             unsafe {
-                // In the fast path, C++ never store / invoke on_complete / ctx.
+                // SAFETY: `ctx_ptr` comes from `Box::into_raw` above and is non-null and uniquely
+                // owned here. C++ only takes ownership of `ctx` when capture_frame returns true;
+                // on a false return it has not, so reclaiming the Box exactly once is sound.
                 if !self.sys_handle.capture_frame(
                     chunk,
                     self.sample_rate,
@@ -170,14 +194,28 @@ impl NativeAudioSource {
                     ctx_ptr,
                     sys_at::CompleteCallback(lk_audio_source_complete),
                 ) {
+                    drop(Box::from_raw(ctx_ptr as *mut oneshot::Sender<()>));
                     return Err(RtcError {
                         error_type: RtcErrorType::InvalidState,
                         message: "failed to capture frame".to_owned(),
                     });
                 }
             }
 
-            let _ = rx.await;
+            // Bound the wait for the drain's completion (timeout rationale above). On a stall,
+            // clear_buffer() releases the pending completion so the source stays usable afterward.
+            // Route through livekit_runtime so capture_frame stays runtime-neutral (tokio::time
+            // would panic when awaited off a Tokio runtime).
+            match livekit_runtime::timeout(capture_timeout, rx).await {
+                Ok(_) => {}
+                Err(_) => {
+                    self.clear_buffer();
+                    return Err(RtcError {
+                        error_type: RtcErrorType::InvalidState,
+                        message: "audio capture timed out: source drain stalled".to_owned(),
+                    });
+                }
+            }
         }
 
         Ok(())

libwebrtc/src/native/drain_stall_tests.rs

@@ -0,0 +1,169 @@
+// Copyright 2026 LiveKit, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Regression tests for the buffered `NativeAudioSource::capture_frame` drain stall
+//! (rust-sdks #408 / #420 / #497): the C++ `AudioTrackSource` drain fires the per-chunk completion
+//! that `capture_frame` awaits. If the drain stops progressing — a sink blocked in `OnData`, or the
+//! drain TaskQueue CPU-starved — an unbounded await would wedge the producer (and any session built
+//! on it) forever.
+
+use crate::audio_frame::AudioFrame;
+use crate::audio_source::native::NativeAudioSource;
+use crate::audio_source::AudioSourceOptions;
+
+const SAMPLE_RATE: u32 = 48000;
+
+fn silent_frame<'a>(samples: usize) -> AudioFrame<'a> {
+    AudioFrame {
+        data: vec![0i16; samples].into(),
+        sample_rate: SAMPLE_RATE,
+        num_channels: 1,
+        samples_per_channel: samples as u32,
+    }
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn capture_frame_completes_under_normal_backpressure() {
+    // The fix must not regress the healthy path. A large queue keeps the derived timeout far above
+    // the ~one-queue-duration a healthy deferral needs, so this asserts "no false-trip on legitimate
+    // backpressure" without coupling to the production floor (which would make it flaky on a loaded
+    // CI runner). No sink/track is created, so this is safe on every platform.
+    const QUEUE_MS: u32 = 1000;
+    let q = (SAMPLE_RATE / 1000 * QUEUE_MS) as usize;
+    let source = NativeAudioSource::new(AudioSourceOptions::default(), SAMPLE_RATE, 1, QUEUE_MS);
+
+    for i in 0..3u32 {
+        // q + q/2 forces a deferral (completion routed through the drain) every call.
+        source
+            .capture_frame(&silent_frame(q + q / 2))
+            .await
+            .unwrap_or_else(|e| panic!("healthy capture {i} failed: {e:?}"));
+    }
+}
+
+// The stall reproducer needs a real audio track, which requires a `PeerConnectionFactory`. On
+// macOS/Windows that brings up the platform AudioDeviceModule, whose real-time audio thread aborts
+// the process when the test's blocking sink stalls it. The Linux CI webrtc build has no such device
+// backend, so the reproducer runs there and is deterministic. The fix itself is platform-independent
+// and the healthy-path test above runs everywhere. This test also creates a factory, so — like the
+// sibling factory tests — it relies on serial execution (CI runs `--test-threads=1`).
+#[cfg(target_os = "linux")]
+mod stall {
+    use std::sync::atomic::{AtomicBool, Ordering};
+    use std::sync::Arc;
+    use std::time::Duration;
+
+    use tokio::sync::mpsc;
+    use webrtc_sys::audio_track as sys_at;
+
+    use super::{silent_frame, SAMPLE_RATE};
+    use crate::audio_source::native::NativeAudioSource;
+    use crate::audio_source::AudioSourceOptions;
+    use crate::peer_connection_factory::native::PeerConnectionFactoryExt;
+    use crate::peer_connection_factory::PeerConnectionFactory;
+    use crate::RtcErrorType;
+
+    /// A sink whose `on_data` blocks the C++ 10ms drain task, modelling a wedged/starved consumer.
+    /// It signals `entered` on its first call (so the test awaits the stall instead of sleeping for
+    /// it) and unblocks when `released` is set.
+    struct BlockingSink {
+        entered: mpsc::UnboundedSender<()>,
+        released: Arc<AtomicBool>,
+    }
+
+    impl sys_at::AudioSink for BlockingSink {
+        fn on_data(
+            &self,
+            _data: &[i16],
+            _sample_rate: i32,
+            _num_channels: usize,
+            _num_frames: usize,
+        ) {
+            let _ = self.entered.send(());
+            while !self.released.load(Ordering::Acquire) {
+                std::thread::sleep(Duration::from_millis(5));
+            }
+        }
+    }
+
+    // worker_threads = 2: one for the (on a revert, possibly-wedged) capture task, one for the timeout.
+    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+    async fn capture_frame_recovers_when_drain_stalls() {
+        const QUEUE_MS: u32 = 200;
+        let q = (SAMPLE_RATE / 1000 * QUEUE_MS) as usize;
+
+        let factory = PeerConnectionFactory::default();
+        let source =
+            NativeAudioSource::new(AudioSourceOptions::default(), SAMPLE_RATE, 1, QUEUE_MS);
+        let track = factory.create_audio_track("drain-stall", source.clone());
+
+        let (entered_tx, mut entered_rx) = mpsc::unbounded_channel();
+        let released = Arc::new(AtomicBool::new(false));
+        let native_sink = sys_at::ffi::new_native_audio_sink(
+            Box::new(sys_at::AudioSinkWrapper::new(Arc::new(BlockingSink {
+                entered: entered_tx,
+                released: released.clone(),
+            }))),
+            SAMPLE_RATE as i32,
+            1,
+        );
+        // SAFETY: `track` is an audio track created by `create_audio_track`, so downcasting its
+        // media-stream-track handle back to an `AudioTrack` is valid.
+        let audio = unsafe { sys_at::ffi::media_to_audio(track.sys_handle()) };
+        audio.add_sink(&native_sink);
+
+        // Wait until the drain is actually stalled inside the sink (bounded, so a wiring regression
+        // fails the test cleanly instead of hanging the CI job).
+        tokio::time::timeout(Duration::from_secs(5), entered_rx.recv())
+            .await
+            .expect("drain never entered the sink within 5s")
+            .expect("sink signal channel closed");
+
+        // Feed 2x the queue so the second chunk must wait on the (now stalled) drain.
+        let src = source.clone();
+        let mut handle = tokio::spawn(async move { src.capture_frame(&silent_frame(q * 2)).await });
+
+        // Outer bound distinguishes "returned" from "hung", well above the fix's own timeout.
+        let res = tokio::time::timeout(Duration::from_secs(8), &mut handle).await;
+
+        // Release the drain regardless, so teardown and the recovery check below can proceed.
+        released.store(true, Ordering::Release);
+        if res.is_err() {
+            let _ = handle.await;
+        }
+
+        match res {
+            Ok(Ok(Err(e))) => {
+                assert_eq!(e.error_type, RtcErrorType::InvalidState);
+                assert!(
+                    e.message.contains("source drain stalled"),
+                    "unexpected error: {}",
+                    e.message
+                );
+            }
+            other => {
+                panic!("capture_frame should return an error on a stalled drain, got: {other:?}")
+            }
+        }
+
+        // The stall must leave the source usable: the timeout path releases the pending completion
+        // rather than poisoning the slot, so a fresh capture succeeds once the drain is unstuck.
+        tokio::time::timeout(Duration::from_secs(5), source.capture_frame(&silent_frame(q)))
+            .await
+            .expect("recovery capture hung")
+            .expect("source did not recover after the stall cleared");
+
+        audio.remove_sink(&native_sink);
+    }
+}

libwebrtc/src/native/mod.rs

@@ -23,6 +23,8 @@ pub mod audio_track;
 pub mod data_channel;
 #[cfg(any(target_os = "macos", target_os = "windows", target_os = "linux"))]
 pub mod desktop_capturer;
+#[cfg(test)]
+mod drain_stall_tests;
 pub mod frame_cryptor;
 pub mod ice_candidate;
 pub mod media_stream;

webrtc-sys/include/livekit/audio_track.h

@@ -133,17 +133,24 @@ class AudioTrackSource {
     bool is_external_source() const { return true; }
 
    private:
-    mutable webrtc::Mutex mutex_;
+    // Split lock. `sink_mutex_` guards `sinks_` and is held across the drain's OnData loop, so a
+    // wedged or slow sink blocks only AddSink/RemoveSink — never the producer — while still
+    // preventing a sink from being freed mid-OnData. `buffer_mutex_` guards the producer and
+    // completion state, so capture_frame never contends with a sink.
+    mutable webrtc::Mutex sink_mutex_;
+    mutable webrtc::Mutex buffer_mutex_;
     std::unique_ptr<webrtc::TaskQueueBase, webrtc::TaskQueueDeleter> audio_queue_;
     webrtc::RepeatingTaskHandle audio_task_;
 
-    std::vector<webrtc::AudioTrackSinkInterface*> sinks_ RTC_GUARDED_BY(mutex_);
-    std::vector<int16_t> buffer_ RTC_GUARDED_BY(mutex_);
+    std::vector<webrtc::AudioTrackSinkInterface*> sinks_ RTC_GUARDED_BY(sink_mutex_);
+    std::vector<int16_t> buffer_ RTC_GUARDED_BY(buffer_mutex_);
 
-    const SourceContext* capture_userdata_ RTC_GUARDED_BY(mutex_);
-    void (*on_complete_)(const SourceContext*) RTC_GUARDED_BY(mutex_);
+    const SourceContext* capture_userdata_ RTC_GUARDED_BY(buffer_mutex_);
+    void (*on_complete_)(const SourceContext*) RTC_GUARDED_BY(buffer_mutex_);
 
     std::vector<int16_t> silence_buffer_;
+    // Reusable 10ms scratch owned solely by the single-threaded drain task (no lock needed).
+    std::vector<int16_t> scratch_;
 
     int sample_rate_ = 0;
     int num_channels_ = 0;