Use CancellationToken instead

Author: alan-george-lk

.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`].

Cargo.lock

@@ -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

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 }

livekit-datatrack/Cargo.toml

@@ -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.

livekit-datatrack/src/local/events.rs

@@ -27,14 +27,14 @@ use anyhow::{anyhow, Context};
 use futures_core::Stream;
 use std::{
     collections::HashMap,
-    ops::ControlFlow,
     pin::Pin,
     sync::Arc,
     task::{Context as TaskContext, Poll},
     time::Duration,
 };
 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<InputEvent>,
     event_in_rx: mpsc::Receiver<InputEvent>,
     event_out_tx: mpsc::Sender<OutputEvent>,
-    shutdown_rx: watch::Receiver<bool>,
+    token: CancellationToken,
     handle_allocator: packet::HandleAllocator,
     descriptors: HashMap<Handle, Descriptor>,
 }
@@ -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,42 +88,30 @@ 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;
                 }
             }
         }
         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<()> {
+    /// 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<DataTrackFrame>,
     event_in_tx: mpsc::Sender<InputEvent>,
     event_out_tx: mpsc::Sender<OutputEvent>,
+    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<InputEvent>,
-    drop_guard: Arc<DropGuard>,
+    token: CancellationToken,
+    /// Cancels the manager when the last [`ManagerInput`] is dropped.
+    _drop_guard: Arc<CancelOnDrop>,
 }
 
 /// 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<bool>,
-}
+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<InputEvent>, shutdown_tx: watch::Sender<bool>) -> Self {
-        Self { event_in_tx, drop_guard: DropGuard { shutdown_tx }.into() }
+    fn new(event_in_tx: mpsc::Sender<InputEvent>, 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

livekit-datatrack/src/local/manager.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.