From da307b57e5ba2fe4e3f672e1cce277eeeb7cc7c1 Mon Sep 17 00:00:00 2001
From: aoife cassidy <aoife@livekit.io>
Date: Sat, 13 Jul 2024 23:48:19 +0300
Subject: [PATCH] livekit: wait for disconnect before callback

---
 livekit/src/room/mod.rs       | 6 ++++--
 livekit/src/rtc_engine/mod.rs | 5 ++++-
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/livekit/src/room/mod.rs b/livekit/src/room/mod.rs
index 0561e5d7..9d0a458c 100644
--- a/livekit/src/room/mod.rs
+++ b/livekit/src/room/mod.rs
@@ -566,7 +566,7 @@ impl RoomSession {
             EngineEvent::SignalRestarted { join_response, tx } => {
                 self.handle_signal_restarted(join_response, tx)
             }
-            EngineEvent::Disconnected { reason } => self.handle_disconnected(reason),
+            EngineEvent::Disconnected { reason, tx } => self.handle_disconnected(reason, tx),
             EngineEvent::Data { payload, topic, kind, participant_sid, participant_identity } => {
                 self.handle_data(payload, topic, kind, participant_sid, participant_identity);
             }
@@ -960,7 +960,7 @@ impl RoomSession {
         let _ = tx.send(());
     }
 
-    fn handle_disconnected(self: &Arc<Self>, reason: DisconnectReason) {
+    fn handle_disconnected(self: &Arc<Self>, reason: DisconnectReason, tx: oneshot::Sender<()>) {
         if self.update_connection_state(ConnectionState::Disconnected) {
             self.dispatcher.dispatch(&RoomEvent::Disconnected { reason });
         }
@@ -975,6 +975,8 @@ impl RoomSession {
                 }
             });
         }
+
+        let _ = tx.send(());
     }
 
     fn handle_data(
diff --git a/livekit/src/rtc_engine/mod.rs b/livekit/src/rtc_engine/mod.rs
index eda7c81c..7f008536 100644
--- a/livekit/src/rtc_engine/mod.rs
+++ b/livekit/src/rtc_engine/mod.rs
@@ -129,6 +129,7 @@ pub enum EngineEvent {
     },
     Disconnected {
         reason: DisconnectReason,
+        tx: oneshot::Sender<()>,
     },
 }
 
@@ -439,10 +440,12 @@ impl EngineInner {
         };
 
         if let Some((engine_task, close_tx)) = engine_task {
+            let (tx, rx) = oneshot::channel();
             session.close().await;
             let _ = close_tx.send(());
             let _ = engine_task.await;
-            let _ = self.engine_tx.send(EngineEvent::Disconnected { reason });
+            let _ = self.engine_tx.send(EngineEvent::Disconnected { reason, tx });
+            let _ = rx.await;
         }
     }