From f3bdba68c27959487681d60310f27b3961510c53 Mon Sep 17 00:00:00 2001 From: Alan George Date: Tue, 4 Aug 2026 21:33:45 -0600 Subject: [PATCH 1/2] Shutdown no longer travels through event channel, resolving data tracks deadlock on shutdown under high load --- .../fix_data_track_shutdown_deadlock.md | 11 ++ livekit-datatrack/src/local/manager.rs | 125 +++++++++++++++--- livekit-datatrack/src/remote/manager.rs | 101 +++++++++++--- 3 files changed, 193 insertions(+), 44 deletions(-) create mode 100644 .changeset/fix_data_track_shutdown_deadlock.md diff --git a/.changeset/fix_data_track_shutdown_deadlock.md b/.changeset/fix_data_track_shutdown_deadlock.md new file mode 100644 index 000000000..23e2a7758 --- /dev/null +++ b/.changeset/fix_data_track_shutdown_deadlock.md @@ -0,0 +1,11 @@ +--- +livekit-datatrack: patch +livekit-ffi: patch +livekit-uniffi: patch +livekit: patch +--- + +Fix a data track manager deadlock during room disconnect. Shutdown was delivered +through the bounded event channel with `try_send`, so it was silently dropped +whenever in-flight track events had saturated the channel, leaving the manager +task and every caller awaiting disconnect stranded. diff --git a/livekit-datatrack/src/local/manager.rs b/livekit-datatrack/src/local/manager.rs index 8b3830c33..4af7faa38 100644 --- a/livekit-datatrack/src/local/manager.rs +++ b/livekit-datatrack/src/local/manager.rs @@ -27,6 +27,7 @@ use anyhow::{anyhow, Context}; use futures_core::Stream; use std::{ collections::HashMap, + ops::ControlFlow, pin::Pin, sync::Arc, task::{Context as TaskContext, Poll}, @@ -51,6 +52,7 @@ pub struct Manager { event_in_tx: mpsc::Sender, event_in_rx: mpsc::Receiver, event_out_tx: mpsc::Sender, + shutdown_rx: watch::Receiver, handle_allocator: packet::HandleAllocator, descriptors: HashMap, } @@ -67,13 +69,15 @@ impl Manager { pub fn new(options: ManagerOptions) -> (Self, ManagerInput, ManagerOutput) { let (event_in_tx, event_in_rx) = mpsc::channel(Self::EVENT_BUFFER_COUNT); let (event_out_tx, event_out_rx) = mpsc::channel(Self::EVENT_BUFFER_COUNT); + let (shutdown_tx, shutdown_rx) = watch::channel(false); - let event_in = ManagerInput::new(event_in_tx.clone()); + let event_in = ManagerInput::new(event_in_tx.clone(), shutdown_tx); let manager = Manager { encryption_provider: options.encryption_provider, event_in_tx, event_in_rx, event_out_tx, + shutdown_rx, handle_allocator: packet::HandleAllocator::default(), descriptors: HashMap::new(), }; @@ -88,25 +92,52 @@ impl Manager { /// pub async fn run(mut self) { log::debug!("Task started"); - while let Some(event) = self.event_in_rx.recv().await { - log::debug!("Input event: {:?}", event); - match event { - InputEvent::PublishRequest(event) => self.on_publish_request(event).await, - InputEvent::PublishCancelled(event) => self.on_publish_cancelled(event).await, - InputEvent::QueryPublished(event) => self.on_query_published(event).await, - InputEvent::UnpublishRequest(event) => self.on_unpublish_request(event).await, - InputEvent::SfuPublishResponse(event) => self.on_sfu_publish_response(event).await, - InputEvent::SfuUnpublishResponse(event) => { - self.on_sfu_unpublish_response(event).await + loop { + tokio::select! { + // Biased so queued events are still processed in order once + // shutdown has been signalled out-of-band; see `ManagerInput::send`. + biased; + event = self.event_in_rx.recv() => { + let Some(event) = event else { break }; + if self.handle_event(event).await.is_break() { + break; + } + } + _ = self.shutdown_rx.changed() => { + self.drain_pending().await; + break; } - InputEvent::RepublishTracks => self.on_republish_tracks().await, - InputEvent::Shutdown => break, } } self.shutdown().await; log::debug!("Task ended"); } + /// Drains events that were queued ahead of an out-of-band shutdown signal. + async fn drain_pending(&mut self) { + while let Ok(event) = self.event_in_rx.try_recv() { + if self.handle_event(event).await.is_break() { + break; + } + } + } + + /// Handles a single input event, reporting whether the task should stop. + async fn handle_event(&mut self, event: InputEvent) -> ControlFlow<()> { + log::debug!("Input event: {:?}", event); + match event { + InputEvent::PublishRequest(event) => self.on_publish_request(event).await, + InputEvent::PublishCancelled(event) => self.on_publish_cancelled(event).await, + InputEvent::QueryPublished(event) => self.on_query_published(event).await, + InputEvent::UnpublishRequest(event) => self.on_unpublish_request(event).await, + InputEvent::SfuPublishResponse(event) => self.on_sfu_publish_response(event).await, + InputEvent::SfuUnpublishResponse(event) => self.on_sfu_unpublish_response(event).await, + InputEvent::RepublishTracks => self.on_republish_tracks().await, + InputEvent::Shutdown => return ControlFlow::Break(()), + } + ControlFlow::Continue(()) + } + async fn on_publish_request(&mut self, event: PublishRequest) { if let Err(error) = crate::schema::validate_schema( event.options.frame_encoding.as_ref(), @@ -308,18 +339,37 @@ impl Manager { } /// Performs cleanup before the task ends. - async fn shutdown(self) { - for (_, descriptor) in self.descriptors { + async fn shutdown(mut self) { + let mut task_handles = Vec::new(); + for (_, descriptor) in std::mem::take(&mut self.descriptors) { match descriptor { Descriptor::Pending(result_tx) => { _ = result_tx.send(Err(PublishError::Disconnected)) } Descriptor::Active { state_tx, task_handle, .. } => { _ = state_tx.send(PublishState::Unpublished); - task_handle.await; + task_handles.push(task_handle); } } } + + // Track tasks emit a final unpublish request as they end, so the input + // channel has to keep draining while they are joined. Joining without + // draining deadlocks as soon as more tasks are ending than the channel + // can buffer. + let join_tasks = async { + for task_handle in task_handles { + task_handle.await; + } + }; + tokio::pin!(join_tasks); + loop { + tokio::select! { + _ = &mut join_tasks => break, + // Never yields `None`: the manager owns a sender for its own lifetime. + _ = self.event_in_rx.recv() => {} + } + } } /// Maximum number of outgoing frames to buffer per track. @@ -415,7 +465,7 @@ pub(crate) enum PublishState { #[derive(Debug, Clone)] pub struct ManagerInput { event_in_tx: mpsc::Sender, - _drop_guard: Arc, + drop_guard: Arc, } /// Stream of [`OutputEvent`]s produced by [`Manager`]. @@ -430,25 +480,32 @@ impl Stream for ManagerOutput { } } -/// Guard that sends shutdown event when the last reference is dropped. +/// Guard that signals shutdown when the last reference is dropped. #[derive(Debug)] struct DropGuard { - event_in_tx: mpsc::Sender, + shutdown_tx: watch::Sender, } impl Drop for DropGuard { fn drop(&mut self) { - _ = self.event_in_tx.try_send(InputEvent::Shutdown); + _ = self.shutdown_tx.send(true); } } impl ManagerInput { - fn new(event_in_tx: mpsc::Sender) -> Self { - Self { event_in_tx: event_in_tx.clone(), _drop_guard: DropGuard { event_in_tx }.into() } + fn new(event_in_tx: mpsc::Sender, shutdown_tx: watch::Sender) -> Self { + Self { event_in_tx, drop_guard: DropGuard { shutdown_tx }.into() } } /// Sends an input event to the manager's task to be processed. pub fn send(&self, event: InputEvent) -> Result<(), InternalError> { + // Shutdown bypasses the bounded event channel. In-flight track events + // routinely saturate it, and a shutdown dropped for lack of capacity + // strands the manager task along with everyone awaiting its completion. + if matches!(event, InputEvent::Shutdown) { + _ = self.drop_guard.shutdown_tx.send(true); + return Ok(()); + } Ok(self.event_in_tx.try_send(event).context("Failed to handle input event")?) } @@ -527,6 +584,30 @@ mod tests { timeout(Duration::from_secs(1), join_handle).await.unwrap(); } + #[tokio::test] + async fn test_task_shutdown_with_saturated_event_channel() { + let options = ManagerOptions { encryption_provider: None }; + let (manager, input, _output) = Manager::new(options); + + // Fill the event channel before the manager starts draining it so that + // shutdown cannot depend on any remaining capacity. + let mut result_rxs = Vec::new(); + for _ in 0..Manager::EVENT_BUFFER_COUNT { + let (result_tx, result_rx) = oneshot::channel(); + input.send(QueryPublished { result_tx }.into()).unwrap(); + result_rxs.push(result_rx); + } + input.send(InputEvent::Shutdown).unwrap(); + + let join_handle = livekit_runtime::spawn(manager.run()); + timeout(Duration::from_secs(1), join_handle).await.unwrap(); + + // Events queued ahead of the shutdown signal are still processed. + for result_rx in result_rxs { + assert!(result_rx.await.is_ok()); + } + } + #[tokio::test] async fn test_publish() { let payload_size = 256; diff --git a/livekit-datatrack/src/remote/manager.rs b/livekit-datatrack/src/remote/manager.rs index 909d79c75..1330ead9c 100644 --- a/livekit-datatrack/src/remote/manager.rs +++ b/livekit-datatrack/src/remote/manager.rs @@ -30,6 +30,7 @@ use bytes::Bytes; use std::{ collections::{HashMap, HashSet}, mem, + ops::ControlFlow, pin::Pin, sync::{ atomic::{AtomicUsize, Ordering}, @@ -57,6 +58,7 @@ pub struct Manager { event_in_tx: mpsc::Sender, event_in_rx: mpsc::Receiver, event_out_tx: mpsc::Sender, + shutdown_rx: watch::Receiver, /// Mapping between track SID and descriptor. descriptors: HashMap, @@ -81,13 +83,15 @@ impl Manager { pub fn new(options: ManagerOptions) -> (Self, ManagerInput, ManagerOutput) { let (event_in_tx, event_in_rx) = mpsc::channel(Self::EVENT_BUFFER_COUNT); let (event_out_tx, event_out_rx) = mpsc::channel(Self::EVENT_BUFFER_COUNT); + let (shutdown_tx, shutdown_rx) = watch::channel(false); - let event_in = ManagerInput::new(event_in_tx.clone()); + let event_in = ManagerInput::new(event_in_tx.clone(), shutdown_tx); let manager = Manager { decryption_provider: options.decryption_provider, event_in_tx, event_in_rx, event_out_tx, + shutdown_rx, descriptors: HashMap::default(), sub_handles: HashMap::default(), }; @@ -102,26 +106,53 @@ impl Manager { /// pub async fn run(mut self) { log::debug!("Task started"); - while let Some(event) = self.event_in_rx.recv().await { - match event { - InputEvent::SubscribeRequest(event) => self.on_subscribe_request(event).await, - InputEvent::UnsubscribeRequest(event) => self.on_unsubscribe_request(event).await, - InputEvent::SfuPublicationUpdates(event) => { - self.on_sfu_publication_updates(event).await + loop { + tokio::select! { + // Biased so queued events are still processed in order once + // shutdown has been signalled out-of-band; see `ManagerInput::send`. + biased; + event = self.event_in_rx.recv() => { + let Some(event) = event else { break }; + if self.handle_event(event).await.is_break() { + break; + } } - InputEvent::SfuSubscriberHandles(event) => self.on_sfu_subscriber_handles(event), - InputEvent::SetPipelineOptions(event) => self.on_set_pipeline_options(event), - InputEvent::PacketReceived(bytes) => self.on_packet_received(bytes), - InputEvent::ResendSubscriptionUpdates => { - self.on_resend_subscription_updates().await + _ = self.shutdown_rx.changed() => { + self.drain_pending().await; + break; } - InputEvent::Shutdown => break, } } self.shutdown().await; log::debug!("Task ended"); } + /// Drains events that were queued ahead of an out-of-band shutdown signal. + async fn drain_pending(&mut self) { + while let Ok(event) = self.event_in_rx.try_recv() { + if self.handle_event(event).await.is_break() { + break; + } + } + } + + /// Handles a single input event, reporting whether the task should stop. + async fn handle_event(&mut self, event: InputEvent) -> ControlFlow<()> { + match event { + InputEvent::SubscribeRequest(event) => self.on_subscribe_request(event).await, + InputEvent::UnsubscribeRequest(event) => self.on_unsubscribe_request(event).await, + InputEvent::SfuPublicationUpdates(event) => { + self.on_sfu_publication_updates(event).await + } + InputEvent::SfuSubscriberHandles(event) => self.on_sfu_subscriber_handles(event), + InputEvent::SetPipelineOptions(event) => self.on_set_pipeline_options(event), + InputEvent::PacketReceived(bytes) => self.on_packet_received(bytes), + InputEvent::ResendSubscriptionUpdates => self.on_resend_subscription_updates().await, + InputEvent::Shutdown => return ControlFlow::Break(()), + } + ControlFlow::Continue(()) + } + async fn on_subscribe_request(&mut self, event: SubscribeRequest) { let Some(descriptor) = self.descriptors.get_mut(&event.sid) else { let error = DataTrackSubscribeError::Internal( @@ -421,8 +452,9 @@ impl Manager { } /// Performs cleanup before the task ends. - async fn shutdown(self) { - for (_, descriptor) in self.descriptors { + async fn shutdown(mut self) { + let mut task_handles = Vec::new(); + for (_, descriptor) in mem::take(&mut self.descriptors) { _ = descriptor.published_tx.send(false); match descriptor.subscription { SubscriptionState::None => {} @@ -431,7 +463,25 @@ impl Manager { _ = result_tx.send(Err(DataTrackSubscribeError::Disconnected)); } } - SubscriptionState::Active { task_handle, .. } => task_handle.await, + SubscriptionState::Active { task_handle, .. } => task_handles.push(task_handle), + } + } + + // Track tasks emit a final unsubscribe request as they end, so the input + // channel has to keep draining while they are joined. Joining without + // draining deadlocks as soon as more tasks are ending than the channel + // can buffer. + let join_tasks = async { + for task_handle in task_handles { + task_handle.await; + } + }; + tokio::pin!(join_tasks); + loop { + tokio::select! { + _ = &mut join_tasks => break, + // Never yields `None`: the manager owns a sender for its own lifetime. + _ = self.event_in_rx.recv() => {} } } } @@ -523,7 +573,7 @@ impl TrackTask { #[derive(Debug, Clone)] pub struct ManagerInput { event_in_tx: mpsc::Sender, - _drop_guard: Arc, + drop_guard: Arc, } /// Stream of [`OutputEvent`]s produced by [`Manager`]. @@ -538,25 +588,32 @@ impl Stream for ManagerOutput { } } -/// Guard that sends shutdown event when the last reference is dropped. +/// Guard that signals shutdown when the last reference is dropped. #[derive(Debug)] struct DropGuard { - event_in_tx: mpsc::Sender, + shutdown_tx: watch::Sender, } impl Drop for DropGuard { fn drop(&mut self) { - _ = self.event_in_tx.try_send(InputEvent::Shutdown); + _ = self.shutdown_tx.send(true); } } impl ManagerInput { - fn new(event_in_tx: mpsc::Sender) -> Self { - Self { event_in_tx: event_in_tx.clone(), _drop_guard: DropGuard { event_in_tx }.into() } + fn new(event_in_tx: mpsc::Sender, shutdown_tx: watch::Sender) -> Self { + Self { event_in_tx, drop_guard: DropGuard { shutdown_tx }.into() } } /// Sends an input event to the manager's task to be processed. pub fn send(&self, event: InputEvent) -> Result<(), InternalError> { + // Shutdown bypasses the bounded event channel. In-flight track events + // routinely saturate it, and a shutdown dropped for lack of capacity + // strands the manager task along with everyone awaiting its completion. + if matches!(event, InputEvent::Shutdown) { + _ = self.drop_guard.shutdown_tx.send(true); + return Ok(()); + } Ok(self.event_in_tx.try_send(event).context("Failed to send input event")?) } } From 4bab7d6efcd834fa9e4286496ad08257496c98f0 Mon Sep 17 00:00:00 2001 From: Alan George Date: Wed, 5 Aug 2026 12:25:31 -0600 Subject: [PATCH 2/2] Use CancellationToken instead --- .../fix_data_track_shutdown_deadlock.md | 8 +- Cargo.lock | 1 + Cargo.toml | 1 + livekit-datatrack/Cargo.toml | 1 + livekit-datatrack/src/local/events.rs | 2 - livekit-datatrack/src/local/manager.rs | 135 ++++++++--------- livekit-datatrack/src/remote/events.rs | 2 - livekit-datatrack/src/remote/manager.rs | 140 +++++++++--------- livekit-uniffi/Cargo.toml | 2 +- livekit-uniffi/src/data_track/local.rs | 19 +-- livekit-uniffi/src/data_track/remote.rs | 19 +-- livekit/src/room/mod.rs | 4 +- 12 files changed, 150 insertions(+), 184 deletions(-) diff --git a/.changeset/fix_data_track_shutdown_deadlock.md b/.changeset/fix_data_track_shutdown_deadlock.md index 23e2a7758..ecceb225f 100644 --- a/.changeset/fix_data_track_shutdown_deadlock.md +++ b/.changeset/fix_data_track_shutdown_deadlock.md @@ -5,7 +5,7 @@ livekit-uniffi: patch livekit: patch --- -Fix a data track manager deadlock during room disconnect. Shutdown was delivered -through the bounded event channel with `try_send`, so it was silently dropped -whenever in-flight track events had saturated the channel, leaving the manager -task and every caller awaiting disconnect stranded. +Fix a data track manager deadlock during room disconnect. Shutdown is now signaled +via a `CancellationToken` (with child tokens for track tasks) instead of the bounded +event channel, so it cannot be dropped when in-flight track events saturate the +channel. `InputEvent::Shutdown` is removed; use [`ManagerInput::shutdown`]. diff --git a/Cargo.lock b/Cargo.lock index 9984098ea..310802e46 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3994,6 +3994,7 @@ dependencies = [ "thiserror 2.0.19", "tokio", "tokio-stream", + "tokio-util", "uniffi", ] diff --git a/Cargo.toml b/Cargo.toml index 961d9d1e8..bfdd0f3ea 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -87,6 +87,7 @@ serde_json = "1.0" thiserror = "2" tokio = { version = "1", default-features = false } tokio-stream = "0.1" +tokio-util = "0.7" uniffi = "0.31" # For examples diff --git a/livekit-datatrack/Cargo.toml b/livekit-datatrack/Cargo.toml index 7ca2ac952..f3f39215b 100644 --- a/livekit-datatrack/Cargo.toml +++ b/livekit-datatrack/Cargo.toml @@ -13,6 +13,7 @@ livekit-runtime = { workspace = true, features = ["tokio"] } log = { workspace = true } thiserror = { workspace = true } tokio = { workspace = true, default-features = false, features = ["macros", "sync"] } +tokio-util = { workspace = true } futures-util = { workspace = true, default-features = false, features = ["sink"] } futures-core = { workspace = true } bytes = { workspace = true } diff --git a/livekit-datatrack/src/local/events.rs b/livekit-datatrack/src/local/events.rs index 9400e3f07..fa5d87aa8 100644 --- a/livekit-datatrack/src/local/events.rs +++ b/livekit-datatrack/src/local/events.rs @@ -37,8 +37,6 @@ pub enum InputEvent { /// to be recognized by the SFU. Each republished track will be assigned a new SID. /// RepublishTracks, - /// Shutdown the manager and all associated tracks. - Shutdown, } /// An event produced by [`Manager`](super::manager::Manager) requiring external action. diff --git a/livekit-datatrack/src/local/manager.rs b/livekit-datatrack/src/local/manager.rs index 4af7faa38..1bdc26120 100644 --- a/livekit-datatrack/src/local/manager.rs +++ b/livekit-datatrack/src/local/manager.rs @@ -27,7 +27,6 @@ use anyhow::{anyhow, Context}; use futures_core::Stream; use std::{ collections::HashMap, - ops::ControlFlow, pin::Pin, sync::Arc, task::{Context as TaskContext, Poll}, @@ -35,6 +34,7 @@ use std::{ }; use tokio::sync::{mpsc, oneshot, watch}; use tokio_stream::wrappers::ReceiverStream; +use tokio_util::sync::CancellationToken; /// Options for creating a [`Manager`]. #[derive(Debug)] @@ -52,7 +52,7 @@ pub struct Manager { event_in_tx: mpsc::Sender, event_in_rx: mpsc::Receiver, event_out_tx: mpsc::Sender, - shutdown_rx: watch::Receiver, + token: CancellationToken, handle_allocator: packet::HandleAllocator, descriptors: HashMap, } @@ -69,15 +69,15 @@ impl Manager { pub fn new(options: ManagerOptions) -> (Self, ManagerInput, ManagerOutput) { let (event_in_tx, event_in_rx) = mpsc::channel(Self::EVENT_BUFFER_COUNT); let (event_out_tx, event_out_rx) = mpsc::channel(Self::EVENT_BUFFER_COUNT); - let (shutdown_tx, shutdown_rx) = watch::channel(false); + let token = CancellationToken::new(); - let event_in = ManagerInput::new(event_in_tx.clone(), shutdown_tx); + let event_in = ManagerInput::new(event_in_tx.clone(), token.clone()); let manager = Manager { encryption_provider: options.encryption_provider, event_in_tx, event_in_rx, event_out_tx, - shutdown_rx, + token, handle_allocator: packet::HandleAllocator::default(), descriptors: HashMap::new(), }; @@ -88,24 +88,21 @@ impl Manager { /// Run the manager task, consuming self. /// - /// The manager will continue running until receiving [`InputEvent::Shutdown`]. + /// The manager continues until [`ManagerInput::shutdown`] is called, the last + /// [`ManagerInput`] is dropped, or the input channel closes. /// pub async fn run(mut self) { log::debug!("Task started"); loop { tokio::select! { - // Biased so queued events are still processed in order once - // shutdown has been signalled out-of-band; see `ManagerInput::send`. + // Biased so shutdown ends event processing immediately. biased; + _ = self.token.cancelled() => { + break; + } event = self.event_in_rx.recv() => { let Some(event) = event else { break }; - if self.handle_event(event).await.is_break() { - break; - } - } - _ = self.shutdown_rx.changed() => { - self.drain_pending().await; - break; + self.handle_event(event).await; } } } @@ -113,17 +110,8 @@ impl Manager { log::debug!("Task ended"); } - /// Drains events that were queued ahead of an out-of-band shutdown signal. - async fn drain_pending(&mut self) { - while let Ok(event) = self.event_in_rx.try_recv() { - if self.handle_event(event).await.is_break() { - break; - } - } - } - - /// Handles a single input event, reporting whether the task should stop. - async fn handle_event(&mut self, event: InputEvent) -> ControlFlow<()> { + /// Handles a single input event. + async fn handle_event(&mut self, event: InputEvent) { log::debug!("Input event: {:?}", event); match event { InputEvent::PublishRequest(event) => self.on_publish_request(event).await, @@ -133,9 +121,7 @@ impl Manager { InputEvent::SfuPublishResponse(event) => self.on_sfu_publish_response(event).await, InputEvent::SfuUnpublishResponse(event) => self.on_sfu_unpublish_response(event).await, InputEvent::RepublishTracks => self.on_republish_tracks().await, - InputEvent::Shutdown => return ControlFlow::Break(()), } - ControlFlow::Continue(()) } async fn on_publish_request(&mut self, event: PublishRequest) { @@ -286,6 +272,7 @@ impl Manager { frame_rx, event_in_tx: self.event_in_tx.clone(), event_out_tx: self.event_out_tx.clone(), + token: self.token.child_token(), }; let task_handle = livekit_runtime::spawn(track_task.run()); @@ -310,7 +297,9 @@ impl Manager { return; }; if *state_tx.borrow() != PublishState::Unpublished { - _ = state_tx.send(PublishState::Unpublished); + // `send_replace` updates even if the track task already dropped its receiver + // after observing manager cancellation. + _ = state_tx.send_replace(PublishState::Unpublished); } } @@ -347,28 +336,18 @@ impl Manager { _ = result_tx.send(Err(PublishError::Disconnected)) } Descriptor::Active { state_tx, task_handle, .. } => { - _ = state_tx.send(PublishState::Unpublished); + // `send_replace` updates even if the track task already dropped its + // receiver after observing manager cancellation. + _ = state_tx.send_replace(PublishState::Unpublished); task_handles.push(task_handle); } } } - // Track tasks emit a final unpublish request as they end, so the input - // channel has to keep draining while they are joined. Joining without - // draining deadlocks as soon as more tasks are ending than the channel - // can buffer. - let join_tasks = async { - for task_handle in task_handles { - task_handle.await; - } - }; - tokio::pin!(join_tasks); - loop { - tokio::select! { - _ = &mut join_tasks => break, - // Never yields `None`: the manager owns a sender for its own lifetime. - _ = self.event_in_rx.recv() => {} - } + // Track tasks observe the parent cancellation token via child tokens and + // skip their final unpublish request, so joining alone is sufficient. + for task_handle in task_handles { + task_handle.await; } } @@ -387,6 +366,7 @@ struct TrackTask { frame_rx: mpsc::Receiver, event_in_tx: mpsc::Sender, event_out_tx: mpsc::Sender, + token: CancellationToken, } impl TrackTask { @@ -397,6 +377,8 @@ impl TrackTask { let mut state = *self.state_rx.borrow(); while state != PublishState::Unpublished { tokio::select! { + biased; + _ = self.token.cancelled() => break, _ = self.state_rx.changed() => { state = *self.state_rx.borrow(); } @@ -410,8 +392,11 @@ impl TrackTask { } } - let event = UnpublishRequest { handle: self.info.pub_handle }; - _ = self.event_in_tx.send(event.into()).await; + // Manager-wide shutdown already owns cleanup; only notify for per-track unpublish. + if !self.token.is_cancelled() { + let event = UnpublishRequest { handle: self.info.pub_handle }; + _ = self.event_in_tx.send(event.into()).await; + } log::debug!("Track task ended: sid={}", sid); } @@ -465,7 +450,9 @@ pub(crate) enum PublishState { #[derive(Debug, Clone)] pub struct ManagerInput { event_in_tx: mpsc::Sender, - drop_guard: Arc, + token: CancellationToken, + /// Cancels the manager when the last [`ManagerInput`] is dropped. + _drop_guard: Arc, } /// Stream of [`OutputEvent`]s produced by [`Manager`]. @@ -480,32 +467,37 @@ impl Stream for ManagerOutput { } } -/// Guard that signals shutdown when the last reference is dropped. +/// Cancels a [`CancellationToken`] when dropped. #[derive(Debug)] -struct DropGuard { - shutdown_tx: watch::Sender, -} +struct CancelOnDrop(CancellationToken); -impl Drop for DropGuard { +impl Drop for CancelOnDrop { fn drop(&mut self) { - _ = self.shutdown_tx.send(true); + self.0.cancel(); } } impl ManagerInput { - fn new(event_in_tx: mpsc::Sender, shutdown_tx: watch::Sender) -> Self { - Self { event_in_tx, drop_guard: DropGuard { shutdown_tx }.into() } + fn new(event_in_tx: mpsc::Sender, token: CancellationToken) -> Self { + Self { event_in_tx, token: token.clone(), _drop_guard: Arc::new(CancelOnDrop(token)) } + } + + /// Shuts down the manager, ending all event processing. + /// + /// Unlike [`Self::send`], this does not use the bounded event channel, so it + /// cannot be dropped when the channel is saturated. + /// + pub fn shutdown(&self) { + self.token.cancel(); + } + + /// Returns a clone of the manager's cancellation token. + pub fn cancellation_token(&self) -> CancellationToken { + self.token.clone() } /// Sends an input event to the manager's task to be processed. pub fn send(&self, event: InputEvent) -> Result<(), InternalError> { - // Shutdown bypasses the bounded event channel. In-flight track events - // routinely saturate it, and a shutdown dropped for lack of capacity - // strands the manager task along with everyone awaiting its completion. - if matches!(event, InputEvent::Shutdown) { - _ = self.drop_guard.shutdown_tx.send(true); - return Ok(()); - } Ok(self.event_in_tx.try_send(event).context("Failed to handle input event")?) } @@ -579,7 +571,7 @@ mod tests { let (manager, input, _) = Manager::new(options); let join_handle = livekit_runtime::spawn(manager.run()); - _ = input.send(InputEvent::Shutdown); + input.shutdown(); timeout(Duration::from_secs(1), join_handle).await.unwrap(); } @@ -591,21 +583,14 @@ mod tests { // Fill the event channel before the manager starts draining it so that // shutdown cannot depend on any remaining capacity. - let mut result_rxs = Vec::new(); for _ in 0..Manager::EVENT_BUFFER_COUNT { - let (result_tx, result_rx) = oneshot::channel(); + let (result_tx, _result_rx) = oneshot::channel(); input.send(QueryPublished { result_tx }.into()).unwrap(); - result_rxs.push(result_rx); } - input.send(InputEvent::Shutdown).unwrap(); + input.shutdown(); let join_handle = livekit_runtime::spawn(manager.run()); timeout(Duration::from_secs(1), join_handle).await.unwrap(); - - // Events queued ahead of the shutdown signal are still processed. - for result_rx in result_rxs { - assert!(result_rx.await.is_ok()); - } } #[tokio::test] @@ -900,7 +885,7 @@ mod tests { assert!(active_track.is_published()); // Shutdown the manager - input.send(InputEvent::Shutdown).unwrap(); + input.shutdown(); sleep(Duration::from_millis(50)).await; // Pending publish receives disconnected error diff --git a/livekit-datatrack/src/remote/events.rs b/livekit-datatrack/src/remote/events.rs index 00b8d9153..dfc7f398d 100644 --- a/livekit-datatrack/src/remote/events.rs +++ b/livekit-datatrack/src/remote/events.rs @@ -40,8 +40,6 @@ pub enum InputEvent { /// tracks are subscribed to locally. /// ResendSubscriptionUpdates, - /// Shutdown the manager, ending any subscriptions. - Shutdown, } /// An event produced by [`Manager`](super::manager::Manager) requiring external action. diff --git a/livekit-datatrack/src/remote/manager.rs b/livekit-datatrack/src/remote/manager.rs index 1330ead9c..b06126270 100644 --- a/livekit-datatrack/src/remote/manager.rs +++ b/livekit-datatrack/src/remote/manager.rs @@ -30,7 +30,6 @@ use bytes::Bytes; use std::{ collections::{HashMap, HashSet}, mem, - ops::ControlFlow, pin::Pin, sync::{ atomic::{AtomicUsize, Ordering}, @@ -40,6 +39,7 @@ use std::{ }; use tokio::sync::{broadcast, mpsc, oneshot, watch}; use tokio_stream::{wrappers::ReceiverStream, Stream}; +use tokio_util::sync::CancellationToken; /// Options for creating a [`Manager`]. #[derive(Debug)] @@ -58,7 +58,7 @@ pub struct Manager { event_in_tx: mpsc::Sender, event_in_rx: mpsc::Receiver, event_out_tx: mpsc::Sender, - shutdown_rx: watch::Receiver, + token: CancellationToken, /// Mapping between track SID and descriptor. descriptors: HashMap, @@ -83,15 +83,15 @@ impl Manager { pub fn new(options: ManagerOptions) -> (Self, ManagerInput, ManagerOutput) { let (event_in_tx, event_in_rx) = mpsc::channel(Self::EVENT_BUFFER_COUNT); let (event_out_tx, event_out_rx) = mpsc::channel(Self::EVENT_BUFFER_COUNT); - let (shutdown_tx, shutdown_rx) = watch::channel(false); + let token = CancellationToken::new(); - let event_in = ManagerInput::new(event_in_tx.clone(), shutdown_tx); + let event_in = ManagerInput::new(event_in_tx.clone(), token.clone()); let manager = Manager { decryption_provider: options.decryption_provider, event_in_tx, event_in_rx, event_out_tx, - shutdown_rx, + token, descriptors: HashMap::default(), sub_handles: HashMap::default(), }; @@ -102,24 +102,21 @@ impl Manager { /// Run the manager task, consuming self. /// - /// The manager will continue running until receiving [`InputEvent::Shutdown`]. + /// The manager continues until [`ManagerInput::shutdown`] is called, the last + /// [`ManagerInput`] is dropped, or the input channel closes. /// pub async fn run(mut self) { log::debug!("Task started"); loop { tokio::select! { - // Biased so queued events are still processed in order once - // shutdown has been signalled out-of-band; see `ManagerInput::send`. + // Biased so shutdown ends event processing immediately. biased; + _ = self.token.cancelled() => { + break; + } event = self.event_in_rx.recv() => { let Some(event) = event else { break }; - if self.handle_event(event).await.is_break() { - break; - } - } - _ = self.shutdown_rx.changed() => { - self.drain_pending().await; - break; + self.handle_event(event).await; } } } @@ -127,17 +124,8 @@ impl Manager { log::debug!("Task ended"); } - /// Drains events that were queued ahead of an out-of-band shutdown signal. - async fn drain_pending(&mut self) { - while let Ok(event) = self.event_in_rx.try_recv() { - if self.handle_event(event).await.is_break() { - break; - } - } - } - - /// Handles a single input event, reporting whether the task should stop. - async fn handle_event(&mut self, event: InputEvent) -> ControlFlow<()> { + /// Handles a single input event. + async fn handle_event(&mut self, event: InputEvent) { match event { InputEvent::SubscribeRequest(event) => self.on_subscribe_request(event).await, InputEvent::UnsubscribeRequest(event) => self.on_unsubscribe_request(event).await, @@ -148,9 +136,7 @@ impl Manager { InputEvent::SetPipelineOptions(event) => self.on_set_pipeline_options(event), InputEvent::PacketReceived(bytes) => self.on_packet_received(bytes), InputEvent::ResendSubscriptionUpdates => self.on_resend_subscription_updates().await, - InputEvent::Shutdown => return ControlFlow::Break(()), } - ControlFlow::Continue(()) } async fn on_subscribe_request(&mut self, event: SubscribeRequest) { @@ -397,6 +383,7 @@ impl Manager { packet_rx, frame_tx: frame_tx.clone(), event_in_tx: self.event_in_tx.clone(), + token: self.token.child_token(), }; let task_handle = livekit_runtime::spawn(track_task.run()); @@ -467,22 +454,10 @@ impl Manager { } } - // Track tasks emit a final unsubscribe request as they end, so the input - // channel has to keep draining while they are joined. Joining without - // draining deadlocks as soon as more tasks are ending than the channel - // can buffer. - let join_tasks = async { - for task_handle in task_handles { - task_handle.await; - } - }; - tokio::pin!(join_tasks); - loop { - tokio::select! { - _ = &mut join_tasks => break, - // Never yields `None`: the manager owns a sender for its own lifetime. - _ = self.event_in_rx.recv() => {} - } + // Track tasks observe the parent cancellation token via child tokens and + // skip their final unsubscribe request, so joining alone is sufficient. + for task_handle in task_handles { + task_handle.await; } } @@ -532,6 +507,7 @@ struct TrackTask { packet_rx: mpsc::Receiver, frame_tx: broadcast::Sender, event_in_tx: mpsc::Sender, + token: CancellationToken, } impl TrackTask { @@ -541,13 +517,18 @@ impl TrackTask { let mut is_published = *self.published_rx.borrow(); while is_published { tokio::select! { - biased; // State updates take priority + // Cancellation and publication state take priority over packets. + biased; + _ = self.token.cancelled() => break, _ = self.published_rx.changed() => { is_published = *self.published_rx.borrow(); }, _ = self.frame_tx.closed() => { - let event = UnsubscribeRequest { sid: self.info.sid() }; - _ = self.event_in_tx.send(event.into()).await; + // Manager-wide shutdown already owns cleanup. + if !self.token.is_cancelled() { + let event = UnsubscribeRequest { sid: self.info.sid() }; + _ = self.event_in_tx.send(event.into()).await; + } break; // No more subscribers }, Some(packet) = self.packet_rx.recv() => { @@ -573,7 +554,9 @@ impl TrackTask { #[derive(Debug, Clone)] pub struct ManagerInput { event_in_tx: mpsc::Sender, - drop_guard: Arc, + token: CancellationToken, + /// Cancels the manager when the last [`ManagerInput`] is dropped. + _drop_guard: Arc, } /// Stream of [`OutputEvent`]s produced by [`Manager`]. @@ -588,32 +571,37 @@ impl Stream for ManagerOutput { } } -/// Guard that signals shutdown when the last reference is dropped. +/// Cancels a [`CancellationToken`] when dropped. #[derive(Debug)] -struct DropGuard { - shutdown_tx: watch::Sender, -} +struct CancelOnDrop(CancellationToken); -impl Drop for DropGuard { +impl Drop for CancelOnDrop { fn drop(&mut self) { - _ = self.shutdown_tx.send(true); + self.0.cancel(); } } impl ManagerInput { - fn new(event_in_tx: mpsc::Sender, shutdown_tx: watch::Sender) -> Self { - Self { event_in_tx, drop_guard: DropGuard { shutdown_tx }.into() } + fn new(event_in_tx: mpsc::Sender, token: CancellationToken) -> Self { + Self { event_in_tx, token: token.clone(), _drop_guard: Arc::new(CancelOnDrop(token)) } + } + + /// Shuts down the manager, ending all event processing. + /// + /// Unlike [`Self::send`], this does not use the bounded event channel, so it + /// cannot be dropped when the channel is saturated. + /// + pub fn shutdown(&self) { + self.token.cancel(); + } + + /// Returns a clone of the manager's cancellation token. + pub fn cancellation_token(&self) -> CancellationToken { + self.token.clone() } /// Sends an input event to the manager's task to be processed. pub fn send(&self, event: InputEvent) -> Result<(), InternalError> { - // Shutdown bypasses the bounded event channel. In-flight track events - // routinely saturate it, and a shutdown dropped for lack of capacity - // strands the manager task along with everyone awaiting its completion. - if matches!(event, InputEvent::Shutdown) { - _ = self.drop_guard.shutdown_tx.send(true); - return Ok(()); - } Ok(self.event_in_tx.try_send(event).context("Failed to send input event")?) } } @@ -652,7 +640,7 @@ mod tests { let (manager, input, _) = Manager::new(options); let join_handle = livekit_runtime::spawn(manager.run()); - _ = input.send(InputEvent::Shutdown); + input.shutdown(); time::timeout(Duration::from_secs(1), join_handle).await.unwrap(); } @@ -683,8 +671,15 @@ mod tests { let (frame_tx, frame_rx) = broadcast::channel(4); let (event_in_tx, mut event_in_rx) = mpsc::channel(4); - let task = - TrackTask { info: info, pipeline, published_rx, packet_rx, frame_tx, event_in_tx }; + let task = TrackTask { + info, + pipeline, + published_rx, + packet_rx, + frame_tx, + event_in_tx, + token: CancellationToken::new(), + }; let task_handle = livekit_runtime::spawn(task.run()); let trigger_shutdown = async { @@ -845,7 +840,7 @@ mod tests { expect_event!(output, OutputEvent::TrackPublished); // Drain remaining events; no second TrackAvailable should appear - input.send(InputEvent::Shutdown).unwrap(); + input.shutdown(); while let Some(event) = output.next().await { assert!(!matches!(event, OutputEvent::TrackPublished(_))); } @@ -888,8 +883,17 @@ mod tests { let event = SfuPublicationUpdates { updates: HashMap::from([("id".into(), vec![info])]) }; input.send(event.into()).unwrap(); + // SID reassignment emits no output event; wait for it before shutting down. + time::timeout(Duration::from_secs(1), async { + while track.info().sid() != new_sid { + time::sleep(Duration::from_millis(1)).await; + } + }) + .await + .unwrap(); + // Drain remaining events; no publish/unpublish should appear - input.send(InputEvent::Shutdown).unwrap(); + input.shutdown(); while let Some(event) = output.next().await { assert!(!matches!( event, diff --git a/livekit-uniffi/Cargo.toml b/livekit-uniffi/Cargo.toml index 635acab64..adde2273a 100644 --- a/livekit-uniffi/Cargo.toml +++ b/livekit-uniffi/Cargo.toml @@ -18,7 +18,7 @@ livekit-datatrack = { workspace = true, features = ["uniffi"] } uniffi = { workspace = true, features = ["scaffolding-ffi-buffer-fns", "tokio"] } log = { workspace = true } tokio = { workspace = true, features = ["sync", "rt-multi-thread"] } -tokio-util = "0.7.18" +tokio-util = { workspace = true } prost = "0.12" futures-util = { workspace = true, default-features = false, features = ["sink"] } bytes = { workspace = true } diff --git a/livekit-uniffi/src/data_track/local.rs b/livekit-uniffi/src/data_track/local.rs index e0be363d6..cdedee401 100644 --- a/livekit-uniffi/src/data_track/local.rs +++ b/livekit-uniffi/src/data_track/local.rs @@ -25,7 +25,7 @@ use livekit_datatrack::{ use livekit_protocol as proto; use prost::Message; use std::sync::Arc; -use tokio_util::sync::{CancellationToken, DropGuard}; +use tokio_util::sync::CancellationToken; /// Data track published by the local participant. #[derive(uniffi::Object)] @@ -96,7 +96,6 @@ impl From for livekit_datatrack::api::DataTrackOptions { #[derive(uniffi::Object)] struct LocalDataTrackManager { input: local::ManagerInput, - _guard: DropGuard, } /// Delegate for receiving output events from [`LocalDataTrackManager`]. @@ -116,8 +115,6 @@ impl LocalDataTrackManager { delegate: Arc, encryption_provider: Option>, ) -> Arc { - let token = CancellationToken::new(); - let encryption_provider = encryption_provider.map(|p| p as Arc); let manager_options = local::ManagerOptions { encryption_provider }; @@ -125,16 +122,13 @@ impl LocalDataTrackManager { let rt = crate::runtime::runtime(); - // TODO: in a follow-up PR, refactor manager to work with cancellation tokens directly, eliminating the - // need for this additional task. - rt.spawn(shutdown_forward_task(input.clone(), token.clone())); - - let delegate_forward = DelegateForwardTask { output, delegate, token: token.clone() }; + let delegate_forward = + DelegateForwardTask { output, delegate, token: input.cancellation_token() }; rt.spawn(delegate_forward.run()); rt.spawn(manager.run()); - Self { input, _guard: token.drop_guard() }.into() + Self { input }.into() } /// Publishes a data track with given options. @@ -238,8 +232,3 @@ impl DelegateForwardTask { self.delegate.on_signal_request(req); } } - -async fn shutdown_forward_task(input: local::ManagerInput, token: CancellationToken) { - token.cancelled().await; - _ = input.send(local::InputEvent::Shutdown); -} diff --git a/livekit-uniffi/src/data_track/remote.rs b/livekit-uniffi/src/data_track/remote.rs index 177d98f88..7be7228e8 100644 --- a/livekit-uniffi/src/data_track/remote.rs +++ b/livekit-uniffi/src/data_track/remote.rs @@ -25,7 +25,7 @@ use livekit_protocol as proto; use prost::Message; use std::sync::Arc; use tokio::sync::Mutex; -use tokio_util::sync::{CancellationToken, DropGuard}; +use tokio_util::sync::CancellationToken; /// Data track published by the remote participant. #[derive(uniffi::Object)] @@ -103,7 +103,6 @@ impl DataTrackStream { #[derive(uniffi::Object)] struct RemoteDataTrackManager { input: remote::ManagerInput, - _guard: DropGuard, } /// Delegate for receiving output events from [`RemoteDataTrackManager`]. @@ -131,8 +130,6 @@ impl RemoteDataTrackManager { delegate: Arc, decryption_provider: Option>, ) -> Arc { - let token = CancellationToken::new(); - let decryption_provider = decryption_provider.map(|p| p as Arc); let manager_options = remote::ManagerOptions { decryption_provider }; @@ -140,16 +137,13 @@ impl RemoteDataTrackManager { let rt = crate::runtime::runtime(); - // TODO: in a follow-up PR, refactor manager to work with cancellation tokens directly, eliminating the - // need for this additional task. - rt.spawn(shutdown_forward_task(input.clone(), token.clone())); - - let delegate_forward = DelegateForwardTask { output, delegate, token: token.clone() }; + let delegate_forward = + DelegateForwardTask { output, delegate, token: input.cancellation_token() }; rt.spawn(delegate_forward.run()); rt.spawn(manager.run()); - Self { input, _guard: token.drop_guard() }.into() + Self { input }.into() } /// Resend all subscription updates. @@ -263,8 +257,3 @@ impl DelegateForwardTask { self.delegate.on_signal_request(req); } } - -async fn shutdown_forward_task(input: remote::ManagerInput, token: CancellationToken) { - token.cancelled().await; - _ = input.send(remote::InputEvent::Shutdown); -} diff --git a/livekit/src/room/mod.rs b/livekit/src/room/mod.rs index eb9d64ee8..c8d8a09b3 100644 --- a/livekit/src/room/mod.rs +++ b/livekit/src/room/mod.rs @@ -2204,7 +2204,7 @@ impl RoomSession { None => break, }, _ = close_rx.recv() => { - _ = self.local_dt_input.send(dt::local::InputEvent::Shutdown); + self.local_dt_input.shutdown(); break; }, } @@ -2232,7 +2232,7 @@ impl RoomSession { None => break, }, _ = close_rx.recv() => { - _ = self.remote_dt_input.send(dt::remote::InputEvent::Shutdown); + self.remote_dt_input.shutdown(); break; }, }