diff --git a/.changeset/fix_data_track_shutdown_deadlock.md b/.changeset/fix_data_track_shutdown_deadlock.md new file mode 100644 index 000000000..ecceb225f --- /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 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 8b3830c33..1bdc26120 100644 --- a/livekit-datatrack/src/local/manager.rs +++ b/livekit-datatrack/src/local/manager.rs @@ -34,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)] @@ -51,6 +52,7 @@ pub struct Manager { event_in_tx: mpsc::Sender, event_in_rx: mpsc::Receiver, event_out_tx: mpsc::Sender, + token: CancellationToken, 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 token = CancellationToken::new(); - let event_in = ManagerInput::new(event_in_tx.clone()); + 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, + token, handle_allocator: packet::HandleAllocator::default(), descriptors: HashMap::new(), }; @@ -84,29 +88,42 @@ 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"); - 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 shutdown ends event processing immediately. + biased; + _ = self.token.cancelled() => { + break; + } + event = self.event_in_rx.recv() => { + let Some(event) = event else { break }; + self.handle_event(event).await; } - InputEvent::RepublishTracks => self.on_republish_tracks().await, - InputEvent::Shutdown => break, } } self.shutdown().await; log::debug!("Task ended"); } + /// 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, + 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, + } + } + async fn on_publish_request(&mut self, event: PublishRequest) { if let Err(error) = crate::schema::validate_schema( event.options.frame_encoding.as_ref(), @@ -255,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()); @@ -279,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); } } @@ -308,18 +328,27 @@ 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; + // `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 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; + } } /// Maximum number of outgoing frames to buffer per track. @@ -337,6 +366,7 @@ struct TrackTask { frame_rx: mpsc::Receiver, event_in_tx: mpsc::Sender, event_out_tx: mpsc::Sender, + token: CancellationToken, } impl TrackTask { @@ -347,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(); } @@ -360,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); } @@ -415,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`]. @@ -430,21 +467,33 @@ impl Stream for ManagerOutput { } } -/// Guard that sends shutdown event when the last reference is dropped. +/// Cancels a [`CancellationToken`] when dropped. #[derive(Debug)] -struct DropGuard { - event_in_tx: mpsc::Sender, -} +struct CancelOnDrop(CancellationToken); -impl Drop for DropGuard { +impl Drop for CancelOnDrop { fn drop(&mut self) { - _ = self.event_in_tx.try_send(InputEvent::Shutdown); + self.0.cancel(); } } 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, 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. @@ -522,8 +571,25 @@ 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(); + } + #[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. + for _ in 0..Manager::EVENT_BUFFER_COUNT { + let (result_tx, _result_rx) = oneshot::channel(); + input.send(QueryPublished { result_tx }.into()).unwrap(); + } + input.shutdown(); + + let join_handle = livekit_runtime::spawn(manager.run()); timeout(Duration::from_secs(1), join_handle).await.unwrap(); } @@ -819,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 909d79c75..b06126270 100644 --- a/livekit-datatrack/src/remote/manager.rs +++ b/livekit-datatrack/src/remote/manager.rs @@ -39,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)] @@ -57,6 +58,7 @@ pub struct Manager { event_in_tx: mpsc::Sender, event_in_rx: mpsc::Receiver, event_out_tx: mpsc::Sender, + token: CancellationToken, /// 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 token = CancellationToken::new(); - let event_in = ManagerInput::new(event_in_tx.clone()); + 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, + token, descriptors: HashMap::default(), sub_handles: HashMap::default(), }; @@ -98,30 +102,43 @@ 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"); - 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 shutdown ends event processing immediately. + biased; + _ = self.token.cancelled() => { + 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 + event = self.event_in_rx.recv() => { + let Some(event) = event else { break }; + self.handle_event(event).await; } - InputEvent::Shutdown => break, } } self.shutdown().await; log::debug!("Task ended"); } + /// 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, + 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, + } + } + async fn on_subscribe_request(&mut self, event: SubscribeRequest) { let Some(descriptor) = self.descriptors.get_mut(&event.sid) else { let error = DataTrackSubscribeError::Internal( @@ -366,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()); @@ -421,8 +439,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,9 +450,15 @@ 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 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; + } } /// Maximum number of incoming packets to buffer per track to be sent @@ -482,6 +507,7 @@ struct TrackTask { packet_rx: mpsc::Receiver, frame_tx: broadcast::Sender, event_in_tx: mpsc::Sender, + token: CancellationToken, } impl TrackTask { @@ -491,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() => { @@ -523,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`]. @@ -538,21 +571,33 @@ impl Stream for ManagerOutput { } } -/// Guard that sends shutdown event when the last reference is dropped. +/// Cancels a [`CancellationToken`] when dropped. #[derive(Debug)] -struct DropGuard { - event_in_tx: mpsc::Sender, -} +struct CancelOnDrop(CancellationToken); -impl Drop for DropGuard { +impl Drop for CancelOnDrop { fn drop(&mut self) { - _ = self.event_in_tx.try_send(InputEvent::Shutdown); + self.0.cancel(); } } 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, 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. @@ -595,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(); } @@ -626,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 { @@ -788,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(_))); } @@ -831,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; }, }