reconnection wip + simulate scenario + fixed cyclic types

crates/livekit-core/src/room/mod.rs

@@ -17,9 +17,13 @@ use crate::proto::participant_info;
 use thiserror::Error;
 use tracing::{debug, error, instrument, trace_span, Level};
 
-use crate::rtc_engine::{EngineError, EngineEvent, EngineEvents, RTCEngine};
+use crate::rtc_engine::{
+    EngineError, EngineEvent, EngineEvents, EngineResult, RTCEngine,
+};
 use crate::signal_client::SignalOptions;
 
+pub use crate::rtc_engine::SimulateScenario;
+
 pub mod id;
 pub mod participant;
 pub mod publication;
@@ -375,6 +379,10 @@ impl RoomHandle {
     pub fn local_participant(&self) -> Arc<LocalParticipant> {
         self.inner.local_participant.clone()
     }
+
+    pub async fn simulate_scenario(&self, scenario: SimulateScenario) -> EngineResult<()> {
+        self.inner.rtc_engine.simulate_scenario(scenario).await
+    }
 }
 
 fn unpack_stream_id(stream_id: &str) -> Option<(&str, &str)> {

crates/livekit-core/src/rtc_engine/mod.rs

@@ -1,3 +1,4 @@
+use futures::future::BoxFuture;
 use futures::FutureExt;
 use livekit_webrtc::data_channel::DataSendError;
 use livekit_webrtc::jsep::SdpParseError;
@@ -31,6 +32,18 @@ pub(crate) type EngineEmitter = mpsc::Sender<EngineEvent>;
 pub(crate) type EngineEvents = mpsc::Receiver<EngineEvent>;
 pub(crate) type EngineResult<T> = Result<T, EngineError>;
 
+#[derive(Debug, Clone, Eq, PartialEq)]
+#[repr(u8)]
+pub enum SimulateScenario {
+    SignalReconnect,
+    Speaker,
+    NodeFailure,
+    ServerLeave,
+    Migration,
+    ForceTcp,
+    ForceTls,
+}
+
 #[derive(Error, Debug)]
 pub enum EngineError {
     #[error("signal failure: {0}")]
@@ -74,19 +87,7 @@ lazy_static! {
     // Share one LKRuntime across all RTCEngine instances
     static ref LK_RUNTIME: Mutex<Weak<LKRuntime>> = Mutex::new(Weak::new());
 }
-
+///
-#[derive(Debug, Clone, Eq, PartialEq)]
-#[repr(u8)]
-pub enum SimulateScenario {
-    SignalReconnect,
-    Speaker,
-    NodeFailure,
-    ServerLeave,
-    Migration,
-    ForceTcp,
-    ForceTls,
-}
-
 /// Represents a running RTCSession with the ability to close the session
 /// and the engine_task
 #[derive(Debug)]
@@ -150,13 +151,7 @@ impl RTCEngine {
 
     #[tracing::instrument]
     pub async fn close(&self) {
-        self.inner.opened.store(false, Ordering::SeqCst);
+        self.inner.close().await
-        self.inner.close();
-        let _ = self
-            .inner
-            .engine_emitter
-            .send(EngineEvent::Disconnected)
-            .await;
     }
 
     #[tracing::instrument(skip(data))]
@@ -164,7 +159,7 @@ impl RTCEngine {
         &self,
         data: &DataPacket,
         kind: data_packet::Kind,
-    ) -> Result<(), EngineError> {
+    ) -> EngineResult<()> {
         self.inner.wait_reconnection().await?;
         self.inner
             .running_handle
@@ -174,8 +169,20 @@ impl RTCEngine {
             .unwrap()
             .session
             .publish_data(data, kind)
-            .await?;
+            .await
+    }
 
+    pub async fn simulate_scenario(&self, scenario: SimulateScenario) -> EngineResult<()> {
+        self.inner.wait_reconnection().await?;
+        self.inner
+            .running_handle
+            .read()
+            .await
+            .as_ref()
+            .unwrap()
+            .session
+            .simulate_scenario(scenario)
+            .await;
         Ok(())
     }
 
@@ -207,16 +214,24 @@ impl EngineInner {
                 },
                  _ = &mut close_receiver => {
                     break;
-                }
+                
             }
         }
     }
 
     async fn on_session_event(self: &Arc<Self>, event: SessionEvent) -> EngineResult<()> {
         match event {
-            SessionEvent::Close { reason } => {
+            SessionEvent::Close {
-                info!("received session close: {}", reason);
+                source,
-                self.handle_disconnected().await;
+                reason,
+                can_reconnect,
+            } => {
+                info!("received session close: {}, {:?}", source, reason);
+                if can_reconnect {
+                    self.handle_disconnected().await;
+                } else {
+                    self.close().await;
+                }
             }
             SessionEvent::Data { data } => {}
             SessionEvent::MediaTrack {
@@ -238,38 +253,41 @@ impl EngineInner {
         Ok(())
     }
 
-    async fn connect(
+    fn connect<'a>(
-        self: &Arc<Self>,
+        self: &'a Arc<Self>,
-        url: &str,
+        url: &'a str,
-        token: &str,
+        token: &'a str,
         options: SignalOptions,
-    ) -> EngineResult<()> {
+    ) -> BoxFuture<'a, EngineResult<()>> {
-        let (session_emitter, session_events) = mpsc::unbounded_channel();
+        async {
-        let session = RTCSession::connect(
+            let (session_emitter, session_events) = mpsc::unbounded_channel();
-            url,
+            let session = RTCSession::connect(
-            token,
+                url,
-            options,
+                token,
-            self.lk_runtime.clone(),
+                options,
-            session_emitter,
+                self.lk_runtime.clone(),
-        )
+                session_emitter,
-        .await?;
+            )
+            .await?;
 
-        let (close_sender, close_receiver) = oneshot::channel();
+            let (close_sender, close_receiver) = oneshot::channel();
-        let engine_task = tokio::spawn(self.clone().engine_task(session_events, close_receiver));
+            let engine_task =
+                tokio::spawn(self.clone().engine_task(session_events, close_receiver));
 
-        *self.session_info.lock() = Some(session.info().clone());
+            *self.session_info.lock() = Some(session.info().clone());
-        *self.running_handle.write().await = Some(EngineHandle {
+            *self.running_handle.write().await = Some(EngineHandle {
-            session,
+                session,
-            engine_task,
+                engine_task,
-            close_sender,
+                close_sender,
-        });
+            });
 
-        self.opened.store(true, Ordering::SeqCst);
+            self.opened.store(true, Ordering::SeqCst);
-
+            Ok(())
-        Ok(())
+        }
+        .boxed()
     }
 
-    async fn close(&self) {
+    async fn terminate_session(&self) {
         if let Some(handle) = self.running_handle.write().await.take() {
             handle.session.close().await;
             let _ = handle.close_sender.send(());
@@ -277,6 +295,12 @@ impl EngineInner {
         }
     }
 
+    async fn close(&self) {
+        self.opened.store(false, Ordering::SeqCst);
+        self.terminate_session().await;
+        let _ = self.engine_emitter.send(EngineEvent::Disconnected).await;
+    }
+
     async fn wait_reconnection(&self) -> EngineResult<()> {
         if !self.opened.load(Ordering::SeqCst) {
             Err(EngineError::Connection("not opened".to_owned()))?
@@ -293,6 +317,23 @@ impl EngineInner {
         Ok(())
     }
 
+    fn try_reconnect(self: Arc<Self>) {
+        warn!("reconnecting RTCEngine...");
+
+        if !self.opened.load(Ordering::SeqCst) || self.reconnecting.load(Ordering::SeqCst) {
+            return;
+        }
+
+        let mut reconnect_task = self.reconnect_task.lock();
+        *reconnect_task = Some(tokio::spawn({
+            let inner = self.clone();
+            async move {
+                inner.handle_disconnected().await;
+                inner.reconnect_task.lock().take();
+            }
+        }));
+    }
+
     /// Called every time the PeerConnection or the SignalClient is closed
     /// We first try to resume the connection, if it fails, we start a full reconnect.
     async fn handle_disconnected(self: &Arc<Self>) {
@@ -331,6 +372,7 @@ impl EngineInner {
                         full_reconnect = true;
                     }
                 } else {
+                    info!("Connected but failed?");
                     let _ = self.engine_emitter.send(EngineEvent::Resumed).await;
                     connected = true;
                     break;
@@ -344,7 +386,6 @@ impl EngineInner {
 
         if !connected {
             error!("failed to reconnect after {} attemps", RECONNECT_ATTEMPTS);
-            let _ = self.engine_emitter.send(EngineEvent::Disconnected).await;
             self.close().await;
         }
     }
@@ -353,7 +394,7 @@ impl EngineInner {
     /// It recreates a new RTCSession
     async fn try_restart_connection(self: &Arc<Self>) -> EngineResult<()> {
         let info = self.session_info.lock().clone().unwrap();
-        self.close().await;
+        self.terminate_session().await;
         self.connect(&info.url, &info.token, info.options).await?;
         self.running_handle
             .read()

crates/livekit-core/src/rtc_engine/rtc_session.rs

@@ -12,7 +12,7 @@ use tokio::time::sleep;
 
 use prost::Message;
 use serde::{Deserialize, Serialize};
-use tracing::{debug, error, trace, warn};
+use tracing::{debug, error, info, trace, warn};
 
 use crate::{proto, signal_client};
 use livekit_webrtc::data_channel::{DataChannel, DataChannelInit, DataState};
@@ -24,15 +24,15 @@ use livekit_webrtc::peer_connection_factory::RTCConfiguration;
 
 use crate::proto::data_packet::Value;
 use crate::proto::{
-    data_packet, signal_request, signal_response, DataPacket, JoinResponse, SignalTarget,
+    data_packet, signal_request, signal_response, CandidateProtocol, DataPacket, DisconnectReason,
-    TrickleRequest,
+    JoinResponse, SignalTarget, TrickleRequest,
 };
 use crate::rtc_engine::lk_runtime::LKRuntime;
 use crate::rtc_engine::pc_transport::PCTransport;
 use crate::rtc_engine::rtc_events::{RTCEvent, RTCEvents};
 use crate::signal_client::{SignalClient, SignalEvent, SignalEvents, SignalOptions};
 
-use super::{rtc_events, EngineError, EngineResult};
+use super::{rtc_events, EngineError, EngineResult, SimulateScenario};
 
 pub const MAX_ICE_CONNECT_TIMEOUT: Duration = Duration::from_secs(15);
 pub const LOSSY_DC_LABEL: &str = "_lossy";
@@ -52,7 +52,9 @@ pub enum SessionEvent {
         receiver: RtpReceiver,
     },
     Close {
-        reason: String,
+        source: String,
+        reason: DisconnectReason,
+        can_reconnect: bool,
     },
     Connected,
 }
@@ -238,9 +240,9 @@ impl RTCSession {
     #[tracing::instrument]
     pub async fn close(self) {
         // Close the tasks
-        self.close_emitter.send(true);
+        let _ = self.close_emitter.send(true);
-        self.rtc_task.await;
+        let _ = self.rtc_task.await;
-        self.signal_task.await;
+        let _ = self.signal_task.await;
         self.inner.close().await;
     }
 
@@ -259,6 +261,10 @@ impl RTCSession {
     pub async fn wait_pc_connection(&self) -> EngineResult<()> {
         self.inner.wait_pc_connection().await
     }
+
+    pub async fn simulate_scenario(&self, scenario: SimulateScenario) {
+        self.inner.simulate_scenario(scenario).await
+    }
 }
 
 impl RTCSession {
@@ -332,7 +338,7 @@ impl SessionInner {
                                 }
                             }
                             SignalEvent::Close => {
-                                self.on_session_disconnected("SignalClient closed");
+                                self.on_session_disconnected("SignalClient closed", DisconnectReason::UnknownReason, true);
                             }
                         }
                     } else {
@@ -398,6 +404,9 @@ impl SessionInner {
                         .await?;
                 }
             }
+            signal_response::Message::Leave(leave) => {
+                self.on_session_disconnected("received leave", leave.reason(), leave.can_reconnect);
+            }
             _ => {}
         }
 
@@ -437,7 +446,11 @@ impl SessionInner {
                     self.pc_state
                         .store(PCState::Disconnected as u8, Ordering::SeqCst);
 
-                    self.on_session_disconnected("pc_state failed");
+                    self.on_session_disconnected(
+                        "pc_state failed",
+                        DisconnectReason::UnknownReason,
+                        true,
+                    );
                 }
             }
             RTCEvent::DataChannel {
@@ -494,9 +507,11 @@ impl SessionInner {
 
     /// Called when the SignalClient or one of the PeerConnection has lost the connection
     /// The RTCEngine may try a reconnect.
-    fn on_session_disconnected(&self, reason: &str) {
+    fn on_session_disconnected(&self, source: &str, reason: DisconnectReason, can_reconnect: bool) {
         let _ = self.emitter.send(SessionEvent::Close {
-            reason: reason.to_owned(),
+            source: source.to_owned(),
+            reason,
+            can_reconnect,
         });
     }
 
@@ -507,6 +522,65 @@ impl SessionInner {
         self.subscriber_pc.lock().await.close();
     }
 
+    #[tracing::instrument]
+    async fn simulate_scenario(&self, scenario: SimulateScenario) {
+        match scenario {
+            SimulateScenario::SignalReconnect => {
+                self.signal_client.close().await;
+            }
+            SimulateScenario::Speaker => {
+                self.signal_client
+                    .send(signal_request::Message::Simulate(proto::SimulateScenario {
+                        scenario: Some(proto::simulate_scenario::Scenario::SpeakerUpdate(3)),
+                    }))
+                    .await;
+            }
+            SimulateScenario::NodeFailure => {
+                self.signal_client
+                    .send(signal_request::Message::Simulate(proto::SimulateScenario {
+                        scenario: Some(proto::simulate_scenario::Scenario::NodeFailure(true)),
+                    }))
+                    .await;
+            }
+            SimulateScenario::ServerLeave => {
+                self.signal_client
+                    .send(signal_request::Message::Simulate(proto::SimulateScenario {
+                        scenario: Some(proto::simulate_scenario::Scenario::ServerLeave(true)),
+                    }))
+                    .await;
+            }
+            SimulateScenario::Migration => {
+                self.signal_client
+                    .send(signal_request::Message::Simulate(proto::SimulateScenario {
+                        scenario: Some(proto::simulate_scenario::Scenario::Migration(true)),
+                    }))
+                    .await;
+            }
+            SimulateScenario::ForceTcp => {
+                self.signal_client
+                    .send(signal_request::Message::Simulate(proto::SimulateScenario {
+                        scenario: Some(
+                            proto::simulate_scenario::Scenario::SwitchCandidateProtocol(
+                                CandidateProtocol::Tcp as i32,
+                            ),
+                        ),
+                    }))
+                    .await;
+            }
+            SimulateScenario::ForceTls => {
+                self.signal_client
+                    .send(signal_request::Message::Simulate(proto::SimulateScenario {
+                        scenario: Some(
+                            proto::simulate_scenario::Scenario::SwitchCandidateProtocol(
+                                CandidateProtocol::Tls as i32,
+                            ),
+                        ),
+                    }))
+                    .await;
+            }
+        }
+    }
+
     #[tracing::instrument(skip(data))]
     async fn publish_data(
         &self,

crates/livekit-core/src/signal_client/signal_stream.rs

@@ -4,6 +4,7 @@ use prost::Message as ProstMessage;
 use tokio::net::TcpStream;
 use tokio::sync::{mpsc, oneshot};
 use tokio::task::JoinHandle;
+use tokio_tungstenite::tungstenite::error::ProtocolError;
 use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode;
 use tokio_tungstenite::tungstenite::protocol::CloseFrame;
 use tokio_tungstenite::tungstenite::Message;
@@ -60,6 +61,7 @@ impl SignalStream {
             .append_pair("access_token", token)
             .append_pair("protocol", PROTOCOL_VERSION.to_string().as_str())
             .append_pair("reconnect", if options.reconnect { "1" } else { "0" })
+            .append_pair("sid", &options.sid)
             .append_pair(
                 "auto_subscribe",
                 if options.auto_subscribe { "1" } else { "0" },
@@ -69,9 +71,8 @@ impl SignalStream {
                 if options.adaptive_stream { "1" } else { "0" },
             );
 
-        event!(Level::DEBUG, "connecting to websocket: {}", lk_url);
+        event!(Level::INFO, "connecting to SignalClient: {}", lk_url);
         let (ws_stream, _) = connect_async(lk_url).await?;
-        event!(Level::DEBUG, "connected to websocket");
         let _ = emitter.send(SignalEvent::Open).await;
 
         let (ws_writer, ws_reader) = ws_stream.split();

examples/simple_room/src/app.rs

@@ -10,7 +10,7 @@ use std::sync::{
 };
 use tokio::sync::mpsc;
 
-use livekit::room::{ConnectionState, Room, RoomError};
+use livekit::room::{ConnectionState, Room, RoomError, SimulateScenario};
 
 // Useful default constants for developing
 const DEFAULT_URL: &str = "ws://localhost:7880";
@@ -114,6 +114,11 @@ pub fn run(rt: tokio::runtime::Runtime) {
 
                         state.connecting.store(false, Ordering::SeqCst);
                     }
+                    AsyncCmd::SimulateScenario { scenario } => {
+                        if let Some(handle) = state.room.lock().get_handle() {
+                            let _ = handle.simulate_scenario(scenario).await;
+                        }
+                    }
                 }
             }
         });
@@ -210,7 +215,43 @@ impl App {
                     if ui.button("WebRTC Stats").clicked() {}
                     if ui.button("Events").clicked() {}
                 });
-                ui.menu_button("Simulate", |ui| {});
+                ui.menu_button("Simulate", |ui| {
+                    if ui.button("SignalReconnect").clicked() {
+                        let _ = self.cmd_tx.send(AsyncCmd::SimulateScenario {
+                            scenario: SimulateScenario::SignalReconnect,
+                        });
+                    }
+                    if ui.button("Speaker").clicked() {
+                        let _ = self.cmd_tx.send(AsyncCmd::SimulateScenario {
+                            scenario: SimulateScenario::Speaker,
+                        });
+                    }
+                    if ui.button("NodeFailure").clicked() {
+                        let _ = self.cmd_tx.send(AsyncCmd::SimulateScenario {
+                            scenario: SimulateScenario::NodeFailure,
+                        });
+                    }
+                    if ui.button("ServerLeave").clicked() {
+                        let _ = self.cmd_tx.send(AsyncCmd::SimulateScenario {
+                            scenario: SimulateScenario::ServerLeave,
+                        });
+                    }
+                    if ui.button("Migration").clicked() {
+                        let _ = self.cmd_tx.send(AsyncCmd::SimulateScenario {
+                            scenario: SimulateScenario::Migration,
+                        });
+                    }
+                    if ui.button("ForceTcp").clicked() {
+                        let _ = self.cmd_tx.send(AsyncCmd::SimulateScenario {
+                            scenario: SimulateScenario::ForceTcp,
+                        });
+                    }
+                    if ui.button("ForceTls").clicked() {
+                        let _ = self.cmd_tx.send(AsyncCmd::SimulateScenario {
+                            scenario: SimulateScenario::ForceTls,
+                        });
+                    }
+                });
             });
         });
 
@@ -236,12 +277,10 @@ impl App {
 
                     if ui.button("Connect").clicked() {
                         self.connection_failure = None;
-                        self.cmd_tx
+                        let _ = self.cmd_tx.send(AsyncCmd::RoomConnect {
-                            .send(AsyncCmd::RoomConnect {
+                            url: self.lk_url.clone(),
-                                url: self.lk_url.clone(),
+                            token: self.lk_token.clone(),
-                                token: self.lk_token.clone(),
+                        });
-                            })
-                            .unwrap();
                     }
 
                     if connecting {

examples/simple_room/src/events.rs

@@ -1,8 +1,9 @@
-use livekit::events::TrackSubscribedEvent;
+use livekit::{events::TrackSubscribedEvent, room::SimulateScenario};
 
 #[derive(Debug)]
 pub enum AsyncCmd {
     RoomConnect { url: String, token: String },
+    SimulateScenario { scenario: SimulateScenario }
 }
 
 #[derive(Debug)]