cancellable SignalClient & tracing

crates/livekit-core/src/lk_runtime.rs

@@ -1,4 +1,5 @@
-use log::trace;
+use std::fmt::{Debug, Formatter};
+use tracing::{event, Level};
 
 use livekit_webrtc::peer_connection_factory::PeerConnectionFactory;
 use livekit_webrtc::webrtc::RTCRuntime;
@@ -11,9 +12,15 @@ pub struct LKRuntime {
     pub rtc_runtime: RTCRuntime,
 }
 
+impl Debug for LKRuntime {
+    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
+        write!(f, "LKRuntime")
+    }
+}
+
 impl LKRuntime {
     pub fn new() -> Self {
-        trace!("LKRuntime::new()");
+        event!(Level::TRACE, "LKRuntime::new()");
         let rtc_runtime = RTCRuntime::new();
         Self {
             pc_factory: PeerConnectionFactory::new(rtc_runtime.clone()),
@@ -24,6 +31,6 @@ impl LKRuntime {
 
 impl Drop for LKRuntime {
     fn drop(&mut self) {
-        trace!("LKRuntime::drop()");
+        event!(Level::TRACE, "LKRuntime::drop()");
     }
 }

crates/livekit-core/src/pc_transport.rs

@@ -1,9 +1,7 @@
 use std::future::Future;
 use std::pin::Pin;
 use std::time::Duration;
-
-use log::{error, trace};
-
+use tracing::{Level, event};
 use livekit_webrtc::jsep::{IceCandidate, SessionDescription};
 use livekit_webrtc::peer_connection::{
     IceConnectionState, PeerConnection, RTCOfferAnswerOptions, SignalingState,
@@ -96,7 +94,7 @@ impl PCTransport {
         }
 
         if options.ice_restart {
-            trace!("restarting ICE");
+            event!(Level::TRACE, "restarting ICE");
             self.restarting_ice = true;
         }
 
@@ -107,7 +105,7 @@ impl PCTransport {
                         .set_remote_description(remote_description)
                         .await?;
                 } else {
-                    error!("trying to ice restart when the pc doesn't have remote description");
+                    event!(Level::ERROR, "trying to restart ICE when the pc doesn't have remote description");
                 }
             } else {
                 self.renegotiate = true;
@@ -116,7 +114,6 @@ impl PCTransport {
         }
 
         let offer = self.peer_connection.create_offer(options).await?;
-        trace!("created offer {:?}", offer);
         self.peer_connection
             .set_local_description(offer.clone())
             .await?;

crates/livekit-core/src/room.rs

@@ -1,6 +1,5 @@
 use std::sync::Arc;
 use std::time::Duration;
-use log::trace;
 
 use thiserror::Error;
 use tokio::sync::Mutex;
@@ -8,7 +7,6 @@ use tokio::time::sleep;
 
 use crate::local_participant::LocalParticipant;
 use crate::proto::data_packet;
-use crate::proto::signal_request::Message::Mute;
 use crate::rtc_engine;
 use crate::rtc_engine::{EngineError, RTCEngine};
 
@@ -25,15 +23,13 @@ pub struct Room {
     engine: Arc<Mutex<RTCEngine>>,
 }
 
+#[tracing::instrument]
 pub async fn connect(url: &str, token: &str) -> Result<Room, RoomError> {
     let engine = rtc_engine::connect(url, token).await?;
 
     engine.on_data(Box::new(|packet| {
-        trace!("This is a test");
 
-        Box::pin(async move {
-            trace!("Another test");
-        })
+        Box::pin(async move {})
     })).await;
 
     let join = engine.join_response().await;
@@ -64,10 +60,13 @@ impl Room {
 }
 
 
+// eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJleHAiOjE2NjgxMzc0NDgsImlzcyI6IkFQSXpLYkFTaUNWYWtnSiIsIm5hbWUiOiJ3ZWIiLCJuYmYiOjE2NjQ1Mzc0NDgsInN1YiI6IndlYiIsInZpZGVvIjp7InJvb21DcmVhdGUiOnRydWUsInJvb21Kb2luIjp0cnVlfX0.6VMDdXJYrW3KWrEzxx4hzbmMQnjQIRILQ48Qrbx5j44
 #[tokio::test]
 async fn test_test() {
-    env_logger::init();
+   // console_subscriber::init();
+
     let mut room = connect("ws://localhost:7880", "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJleHAiOjE2NzEyMzk4NjAsImlzcyI6IkFQSXpLYkFTaUNWYWtnSiIsIm5hbWUiOiJ0ZXN0IiwibmJmIjoxNjY0MDM5ODYwLCJzdWIiOiJ0ZXN0IiwidmlkZW8iOnsicm9vbUFkbWluIjp0cnVlLCJyb29tQ3JlYXRlIjp0cnVlLCJyb29tSm9pbiI6dHJ1ZX19.0Bee2jI2cSZveAbZ8MLc-ADoMYQ4l8IRxcAxpXAS6a8").await.unwrap();
     room.local_participant().publish_data(b"This is a test", data_packet::Kind::Reliable).await.unwrap();
-    //sleep(Duration::from_secs(60)).await;
+
+    sleep(Duration::from_secs(60)).await;
 }

crates/livekit-core/src/rtc_engine.rs

@@ -1,15 +1,17 @@
+use std::fmt::{Debug, Formatter};
 use std::future::Future;
 use std::pin::Pin;
-use std::sync::{Arc, Weak};
 use std::sync::atomic::{AtomicBool, AtomicU8, Ordering};
+use std::sync::{Arc, Weak};
 use std::time::Duration;
 
 use lazy_static::lazy_static;
-use log::{error, trace};
 use prost::Message;
+use serde::{Deserialize, Serialize};
 use thiserror::Error;
 use tokio::sync::{mpsc, Mutex};
 use tokio::time;
+use tracing::{event, Level};
 
 use livekit_webrtc::data_channel::{DataChannel, DataChannelInit, DataSendError, DataState};
 use livekit_webrtc::jsep::{IceCandidate, SdpParseError, SessionDescription};
@@ -21,16 +23,15 @@ use livekit_webrtc::peer_connection_factory::{
 };
 use livekit_webrtc::rtc_error::RTCError;
 
-use crate::{proto, signal_client};
 use crate::lk_runtime::LKRuntime;
 use crate::pc_transport::PCTransport;
+use crate::proto::data_packet::Value;
 use crate::proto::{
-    data_packet, DataPacket, JoinResponse, signal_request, signal_response, SignalTarget,
+    data_packet, signal_request, signal_response, DataPacket, JoinResponse, SignalTarget,
     TrickleRequest, UserPacket,
 };
-use crate::proto::data_packet::Value;
 use crate::signal_client::{SignalClient, SignalError};
-use serde::{Deserialize, Serialize};
+use crate::{proto, signal_client};
 
 const LOSSY_DC_LABEL: &str = "_lossy";
 const RELIABLE_DC_LABEL: &str = "_reliable";
@@ -42,6 +43,7 @@ lazy_static! {
 }
 
 #[derive(Serialize, Deserialize)]
+#[allow(non_snake_case)]
 struct IceCandidateJSON {
     sdpMid: String,
     sdpMLineIndex: i32,
@@ -101,18 +103,19 @@ pub enum EngineMessage {
     Data {
         data: Vec<u8>,
         binary: bool,
-        reliable: bool,
     },
 }
 
 pub type OnDataHandler =
-Box<dyn (FnMut(Packet) -> Pin<Box<dyn Future<Output=()> + Send + 'static>>) + Send + Sync>;
+    Box<dyn (FnMut(Packet) -> Pin<Box<dyn Future<Output = ()> + Send + 'static>>) + Send + Sync>;
 
 struct EngineInternal {
     publisher_pc: Arc<Mutex<PCTransport>>,
     subscriber_pc: Arc<Mutex<PCTransport>>,
     lossy_dc: Arc<Mutex<DataChannel>>,
     reliable_dc: Arc<Mutex<DataChannel>>,
+    lossy_dc_sub: Arc<Mutex<Option<DataChannel>>>,
+    reliable_dc_sub: Arc<Mutex<Option<DataChannel>>>,
 
     msg_sender: mpsc::Sender<EngineMessage>,
     join_response: Mutex<JoinResponse>,
@@ -124,6 +127,13 @@ struct EngineInternal {
     on_data_handler: Arc<Mutex<Option<OnDataHandler>>>,
 }
 
+impl Debug for EngineInternal {
+    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
+        write!(f, "EngineInternal")
+    }
+}
+
+#[derive(Debug)]
 pub struct RTCEngine {
     signal_client: Arc<SignalClient>,
     internal: Arc<EngineInternal>,
@@ -132,6 +142,7 @@ pub struct RTCEngine {
     lk_runtime: Arc<LKRuntime>, // Keep a reference while we're using the RTCEngine
 }
 
+#[tracing::instrument]
 pub async fn connect(url: &str, token: &str) -> Result<RTCEngine, EngineError> {
     // Acquire an existing/a new LKRuntime
     let mut lk_runtime_ref = LK_RUNTIME.lock().await;
@@ -145,8 +156,8 @@ pub async fn connect(url: &str, token: &str) -> Result<RTCEngine, EngineError> {
     let lk_runtime = lk_runtime.unwrap();
     let signal_client = Arc::new(signal_client::connect(url, token).await?);
 
-    if let signal_response::Message::Join(join_response) = signal_client.recv().await? {
-        trace!("received join_response: {:?}", join_response);
+    if let Some(signal_response::Message::Join(join_response)) = signal_client.recv().await {
+        event!(Level::DEBUG, "received JoinResponse: {:?}", join_response);
         let (sender, receiver) = mpsc::channel(8);
         let internal = Arc::new(RTCEngine::configure(
             lk_runtime.clone(),
@@ -179,6 +190,7 @@ pub async fn connect(url: &str, token: &str) -> Result<RTCEngine, EngineError> {
 
 impl RTCEngine {
     /// Send data to other participants in the Room
+    #[tracing::instrument]
     pub async fn publish_data(
         &mut self,
         data: &DataPacket,
@@ -210,6 +222,7 @@ impl RTCEngine {
         }
     }
 
+    #[tracing::instrument]
     async fn ensure_publisher_connected(
         &mut self,
         kind: data_packet::Kind,
@@ -218,11 +231,12 @@ impl RTCEngine {
             return Ok(());
         }
 
+        let publisher = &self.internal.publisher_pc;
         {
-            let mut publisher = self.internal.publisher_pc.lock().await;
+            let mut publisher = publisher.lock().await;
             if !publisher.is_connected()
                 && publisher.peer_connection().ice_connection_state()
-                != IceConnectionState::IceConnectionChecking
+                    != IceConnectionState::IceConnectionChecking
             {
                 tokio::spawn({
                     let rtc_internal = self.internal.clone();
@@ -234,55 +248,44 @@ impl RTCEngine {
         }
 
         let dc = self.data_channel(kind);
-        {
-            let dc = self.data_channel(kind).lock().await;
-            if dc.state() == DataState::Open {
-                return Ok(());
-            }
+        if dc.lock().await.state() == DataState::Open {
+            return Ok(());
         }
 
-        let res = time::timeout(MAX_ICE_CONNECT_TIMEOUT, {
-            let internal = self.internal.clone();
+        let res = time::timeout(MAX_ICE_CONNECT_TIMEOUT, async move {
+            let mut interval = time::interval(Duration::from_millis(50));
 
-            async move {
-                let mut interval = time::interval(Duration::from_millis(50));
-
-                loop {
-                    if internal.publisher_pc.lock().await.is_connected() && dc.lock().await.state() == DataState::Open {
-                        break;
-                    }
-
-                    interval.tick().await;
+            loop {
+                if publisher.lock().await.is_connected()
+                    && dc.lock().await.state() == DataState::Open
+                {
+                    break;
                 }
+
+                interval.tick().await;
             }
         })
-            .await;
+        .await;
 
         if res.is_err() {
-            Err(EngineError::Connection(
-                "could not establish publisher connection".to_string(),
-            ))
+            let err =
+                EngineError::Connection("could not establish publisher connection".to_string());
+            event!(Level::ERROR, error = ?err);
+            Err(err)
         } else {
             Ok(())
         }
     }
 
-    fn send_request(msg: signal_request::Message, signal_client: Arc<SignalClient>) {
-        tokio::spawn(async move {
-            if let Err(err) = signal_client.send(msg).await {
-                error!("failed to send signal: {:?}", err);
-            }
-        });
-    }
-
+    #[tracing::instrument]
     async fn handle_signal(
         signal: signal_response::Message,
-        signal_client: &Arc<SignalClient>,
-        rtc_internal: &Arc<EngineInternal>,
+        signal_client: Arc<SignalClient>,
+        rtc_internal: Arc<EngineInternal>,
     ) -> Result<(), EngineError> {
         match signal {
             signal_response::Message::Answer(answer) => {
-                trace!("received answer for publisher: {:?}", answer);
+                event!(Level::TRACE, "received answer for publisher: {:?}", answer);
                 let sdp = SessionDescription::from(answer.r#type.parse().unwrap(), &answer.sdp)?;
                 rtc_internal
                     .publisher_pc
@@ -292,35 +295,46 @@ impl RTCEngine {
                     .await?;
             }
             signal_response::Message::Offer(offer) => {
+                event!(Level::TRACE, "received offer for subscriber: {:?}", offer);
                 let sdp = SessionDescription::from(offer.r#type.parse().unwrap(), &offer.sdp)?;
-                let mut subscriber_pc = rtc_internal.subscriber_pc.lock().await;
+                let subscriber_pc = &rtc_internal.subscriber_pc;
 
-                subscriber_pc.set_remote_description(sdp).await?;
+                subscriber_pc
+                    .lock()
+                    .await
+                    .set_remote_description(sdp)
+                    .await?;
                 let answer = subscriber_pc
+                    .lock()
+                    .await
                     .peer_connection()
                     .create_answer(RTCOfferAnswerOptions::default())
                     .await?;
                 subscriber_pc
+                    .lock()
+                    .await
                     .peer_connection()
                     .set_local_description(answer.clone())
                     .await?;
 
-                Self::send_request(
-                    signal_request::Message::Answer(proto::SessionDescription {
-                        r#type: "answer".to_string(),
-                        sdp: answer.to_string(),
-                    }),
-                    signal_client.clone(),
-                );
+                tokio::spawn(async move {
+                    let _ = signal_client.send(signal_request::Message::Answer(
+                        proto::SessionDescription {
+                            r#type: "answer".to_string(),
+                            sdp: answer.to_string(),
+                        },
+                    )).await;
+                });
             }
             signal_response::Message::Trickle(trickle) => {
                 let json: IceCandidateJSON = serde_json::from_str(&trickle.candidate_init)?;
                 let ice = IceCandidate::from(&json.sdpMid, json.sdpMLineIndex, &json.candidate)?;
 
-                trace!(
-                    "received ice_candidate: {:?} (publisher: {:?})",
-                    ice,
-                    trickle.target
+                event!(
+                    Level::TRACE,
+                    "received ice_candidate ({:?}) - {:?}",
+                    SignalTarget::from_i32(trickle.target).unwrap(),
+                    ice
                 );
 
                 if trickle.target == SignalTarget::Publisher as i32 {
@@ -345,40 +359,45 @@ impl RTCEngine {
         Ok(())
     }
 
+    #[tracing::instrument]
     async fn handle_message(
         msg: EngineMessage,
-        signal_client: &Arc<SignalClient>,
-        rtc_internal: &Arc<EngineInternal>,
+        signal_client: Arc<SignalClient>,
+        rtc_internal: Arc<EngineInternal>,
     ) -> Result<(), EngineError> {
         match msg {
             EngineMessage::IceCandidate {
                 ice_candidate,
                 publisher,
             } => {
-                trace!(
-                    "sending ice_candidate: {:?} (publisher: {:?})",
-                    ice_candidate,
-                    publisher
+                let target = if publisher {
+                    SignalTarget::Publisher
+                } else {
+                    SignalTarget::Subscriber
+                };
+
+                event!(
+                    Level::TRACE,
+                    "sending ice_candidate ({:?}) - {:?}",
+                    target,
+                    ice_candidate
                 );
 
                 let json = serde_json::to_string(&IceCandidateJSON {
                     sdpMid: ice_candidate.sdp_mid(),
                     sdpMLineIndex: ice_candidate.sdp_mline_index(),
-                    candidate: ice_candidate.candidate()
+                    candidate: ice_candidate.candidate(),
                 })?;
 
                 // Send the ice_candidate to the server
-                Self::send_request(
-                    signal_request::Message::Trickle(TrickleRequest {
-                        candidate_init: json,
-                        target: if publisher {
-                            SignalTarget::Publisher
-                        } else {
-                            SignalTarget::Subscriber
-                        } as i32,
-                    }),
-                    signal_client.clone(),
-                );
+                tokio::spawn(async move {
+                    let _ = signal_client.send(signal_request::Message::Trickle(
+                        TrickleRequest {
+                            candidate_init: json,
+                            target: target as i32
+                        },
+                    )).await;
+                });
             }
             EngineMessage::ConnectionChange { state, primary } => {
                 if primary && state == PeerConnectionState::Connected {
@@ -400,37 +419,44 @@ impl RTCEngine {
             }
             EngineMessage::PrimaryDataChannel { mut data_channel } => {
                 let reliable = data_channel.label() == RELIABLE_DC_LABEL;
-                Self::configure_dc(&mut data_channel, reliable, rtc_internal.msg_sender.clone());
+                Self::configure_dc(&mut data_channel, rtc_internal.msg_sender.clone());
 
-                trace!(
-                    "received and using subscriber datachannel (reliable: {:?})",
-                    reliable
+                event!(
+                    Level::TRACE,
+                    "received subscriber data_channel - {:?}",
+                    data_channel
                 );
+
                 if reliable {
-                    *rtc_internal.reliable_dc.lock().await = data_channel;
+                    *rtc_internal.reliable_dc_sub.lock().await = Some(data_channel);
                 } else {
-                    *rtc_internal.lossy_dc.lock().await = data_channel;
+                    *rtc_internal.lossy_dc_sub.lock().await = Some(data_channel);
                 }
             }
             EngineMessage::PublisherOffer { offer } => {
-                trace!("received publisher offer: {:?}", offer);
-                // Send the offer to the server
-                Self::send_request(
-                    signal_request::Message::Offer(proto::SessionDescription {
-                        r#type: "offer".to_string(),
-                        sdp: offer.to_string(),
-                    }),
-                    signal_client.clone(),
+                event!(
+                    Level::TRACE,
+                    "sending publisher offer - {:?}",
+                    offer
                 );
+
+                // Send the offer to the server
+                tokio::spawn(async move {
+                    let _ = signal_client.send(signal_request::Message::Offer(
+                        proto::SessionDescription {
+                            r#type: "offer".to_string(),
+                            sdp: offer.to_string(),
+                        },
+                    )).await;
+                });
             }
             EngineMessage::Data {
                 data,
                 binary,
-                reliable: _,
             } => {
                 if !binary {
                     return Err(EngineError::Internal(
-                        "text message aren't supported by LiveKit".to_string(),
+                        "text messages aren't supported by LiveKit".to_string(),
                     ));
                 }
 
@@ -442,7 +468,8 @@ impl RTCEngine {
                             f(Packet {
                                 data: user,
                                 kind: data_packet::Kind::from_i32(data.kind).unwrap(),
-                            }).await;
+                            })
+                            .await;
                         }
                     }
                     Value::Speaker(_) => {
@@ -455,6 +482,7 @@ impl RTCEngine {
         Ok(())
     }
 
+    #[tracing::instrument]
     async fn handle_loop(
         mut receiver: mpsc::Receiver<EngineMessage>,
         signal_client: Arc<SignalClient>,
@@ -462,25 +490,44 @@ impl RTCEngine {
     ) {
         loop {
             tokio::select! {
-                Ok(signal) = signal_client.recv() => {
-                    trace!("received signal: {:?}", signal);
-                    if let Err(err) = Self::handle_signal(signal, &signal_client, &rtc_internal).await {
-                        error!("failed to handle signal: {:?}", err);
+                signal = signal_client.recv() => {
+                    match signal {
+                        Some(signal) => {
+                            if let Err(err) = Self::handle_signal(signal, signal_client.clone(), rtc_internal.clone()).await {
+                                event!(
+                                    Level::ERROR,
+                                    "failed to handle signal: {:?}",
+                                    err,
+                                );
+                            }
+                        }
+                        None => {
+                            // TODO{theomonnom) Trigger reconnect
+                        }
                     }
                 },
                 Some(msg) = receiver.recv() => {
-                    if let Err(err) = Self::handle_message(msg, &signal_client, &rtc_internal).await {
-                        error!("failed to handle engine message: {:?}", err);
+                    if let Err(err) = Self::handle_message(msg, signal_client.clone(), rtc_internal.clone()).await {
+                        event!(
+                            Level::ERROR,
+                            "failed to handle engine message: {:?}",
+                            err,
+                        );
                     }
-                }
+                },
             }
         }
     }
 
+    #[tracing::instrument]
     async fn negotiate_publisher(rtc_internal: Arc<EngineInternal>) -> Result<(), EngineError> {
         rtc_internal.has_published.store(true, Ordering::SeqCst);
         if let Err(err) = rtc_internal.publisher_pc.lock().await.negotiate().await {
-            error!("failed to negotiate the publisher: {:?}", err);
+            event!(
+                Level::ERROR,
+                "failed to negotiate the publisher: {:?}",
+                err,
+            );
             Err(EngineError::Rtc(err))
         } else {
             Ok(())
@@ -488,7 +535,8 @@ impl RTCEngine {
     }
 
     /// This function is called on connect & on reconnect
-    /// It creates the PeerConnections, the DataChannels & the libwebrtc listeners
+    /// It creates the PeerConnections, the DataChannels and the libwebrtc listeners
+    #[tracing::instrument]
     fn configure(
         lk_runtime: Arc<LKRuntime>,
         sender: mpsc::Sender<EngineMessage>,
@@ -543,9 +591,11 @@ impl RTCEngine {
             Box::new(move |offer| {
                 let sender = sender.clone();
 
-                Box::pin(async move {
+                tokio::spawn(async move {
                     let _ = sender.send(EngineMessage::PublisherOffer { offer }).await;
-                })
+                });
+
+                Box::pin(async move {})
             })
         });
 
@@ -605,14 +655,16 @@ impl RTCEngine {
             },
         )?;
 
-        Self::configure_dc(&mut lossy_dc, true, sender.clone());
-        Self::configure_dc(&mut reliable_dc, false, sender.clone());
+        Self::configure_dc(&mut lossy_dc, sender.clone());
+        Self::configure_dc(&mut reliable_dc, sender.clone());
 
         Ok(EngineInternal {
             publisher_pc: Arc::new(Mutex::new(publisher_pc)),
             subscriber_pc: Arc::new(Mutex::new(subscriber_pc)),
             lossy_dc: Arc::new(Mutex::new(lossy_dc)),
             reliable_dc: Arc::new(Mutex::new(reliable_dc)),
+            lossy_dc_sub: Default::default(),
+            reliable_dc_sub: Default::default(),
             msg_sender: sender,
             join_response: Mutex::new(join),
             pc_state: AtomicU8::new(PCState::New as u8),
@@ -622,15 +674,14 @@ impl RTCEngine {
     }
 
     /// Map the libwebrtc listeners to a mpsc channel
+    #[tracing::instrument]
     fn configure_dc(
         data_channel: &mut DataChannel,
-        reliable: bool,
         sender: mpsc::Sender<EngineMessage>,
     ) {
         data_channel.on_message(Box::new(move |data, binary| {
             let _ = sender.blocking_send(EngineMessage::Data {
                 data: data.to_vec(),
-                reliable,
                 binary,
             });
         }));

crates/livekit-core/src/signal_client.rs

@@ -1,16 +1,18 @@
-use futures_util::{SinkExt, StreamExt};
+use futures::future::poll_fn;
 use futures_util::stream::{SplitSink, SplitStream};
-use log::{error, info};
+use futures_util::{SinkExt, StreamExt};
 use prost::Message as ProstMessage;
+use std::fmt::{Debug, Formatter};
 use thiserror::Error;
 use tokio::net::TcpStream;
 use tokio::sync::{mpsc, oneshot};
 use tokio::task::JoinHandle;
-use tokio_tungstenite::{connect_async, MaybeTlsStream, WebSocketStream};
 use tokio_tungstenite::tungstenite::{
-    Error as WsError,
-    Message, protocol::frame::{CloseFrame, coding::CloseCode},
+    protocol::frame::{coding::CloseCode, CloseFrame},
+    Error as WsError, Message,
 };
+use tokio_tungstenite::{connect_async, MaybeTlsStream, WebSocketStream};
+use tracing::{event, span, Level};
 
 use crate::proto::{signal_request, signal_response, SignalRequest, SignalResponse};
 
@@ -31,7 +33,7 @@ type WebSocket = WebSocketStream<MaybeTlsStream<TcpStream>>;
 
 #[derive(Debug)]
 struct RecvMessage {
-    response_chn: oneshot::Sender<SignalResult<signal_response::Message>>,
+    response_chn: oneshot::Sender<Option<signal_response::Message>>,
 }
 
 #[derive(Debug)]
@@ -49,6 +51,13 @@ pub struct SignalClient {
     write_handle: JoinHandle<()>,
 }
 
+impl Debug for SignalClient {
+    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
+        write!(f, "SignalClient")
+    }
+}
+
+#[tracing::instrument]
 pub async fn connect(url: &str, token: &str) -> SignalResult<SignalClient> {
     let mut lk_url = url::Url::parse(url)?;
     lk_url.set_path("/rtc");
@@ -57,11 +66,14 @@ pub async fn connect(url: &str, token: &str) -> SignalResult<SignalClient> {
         .append_pair("access_token", token)
         .append_pair("protocol", PROTOCOL_VERSION.to_string().as_str());
 
-    let (ws_stream, _) = connect_async(lk_url).await?;
+    event!(Level::DEBUG, "connecting to websocket: {}", lk_url);
+    let (ws_stream, _) = connect_async(lk_url.clone()).await?;
+    event!(Level::DEBUG, "connected to SignalClient");
+
     let (ws_writer, ws_reader) = ws_stream.split();
 
-    let (read_tx, read_rx) = mpsc::channel::<RecvMessage>(1);
-    let (write_tx, write_rx) = mpsc::channel::<SendMessage>(1);
+    let (read_tx, read_rx) = mpsc::channel::<RecvMessage>(8);
+    let (write_tx, write_rx) = mpsc::channel::<SendMessage>(8);
     let (read_shutdown_tx, read_shutdown_rx) = oneshot::channel();
     let (write_shutdown_tx, write_shutdown_rx) = oneshot::channel();
 
@@ -90,7 +102,7 @@ impl SignalClient {
         let _ = self.read_handle.await;
     }
 
-    pub async fn recv(&self) -> SignalResult<signal_response::Message> {
+    pub async fn recv(&self) -> Option<signal_response::Message> {
         let (send, recv) = oneshot::channel();
         let msg = RecvMessage { response_chn: send };
         let _ = self.read_sender.send(msg).await;
@@ -107,6 +119,7 @@ impl SignalClient {
         recv.await.expect("channel closed")
     }
 
+    #[tracing::instrument]
     async fn ws_write(
         mut write_receiver: mpsc::Receiver<SendMessage>,
         mut ws_writer: SplitSink<WebSocket, Message>,
@@ -115,13 +128,15 @@ impl SignalClient {
         loop {
             tokio::select! {
                 Some(msg) = write_receiver.recv() => {
+                    event!(Level::TRACE, "sending: {:?}", msg.signal);
+
                     let req = SignalRequest {
                         message: Some(msg.signal),
                     };
 
                     let write_res = ws_writer.send(Message::Binary(req.encode_to_vec())).await;
                     if let Err(err) = write_res {
-                        error!("failed to send message to ws: {:?}", err);
+                        event!(Level::ERROR, "failed to send message to ws: {:?}", err);
                         let _ = msg.response_chn.send(Err(err.into()));
                         break;
                     }
@@ -140,32 +155,35 @@ impl SignalClient {
         }
     }
 
+    #[tracing::instrument]
     async fn ws_read(
-        mut write_receiver: mpsc::Receiver<RecvMessage>,
+        mut read_receiver: mpsc::Receiver<RecvMessage>,
         mut ws_reader: SplitStream<WebSocket>,
         mut shutdown_receiver: oneshot::Receiver<()>,
     ) {
         loop {
             tokio::select! {
-                Some(msg) = write_receiver.recv() => {
-                    let read = ws_reader.next().await;
-                    if read.is_none() {
-                        let _ = msg.response_chn.send(Err(SignalError::WsError(WsError::ConnectionClosed)));
-                        break;
-                    }
-                    let read = read.unwrap();
-                    match read {
-                        Ok(Message::Binary(data)) => {
-                            let res = SignalResponse::decode(data.as_slice()).expect("failed to decode incoming SignalResponse");
-
-                            // TODO(theomonnon) Handle Message::Pong
-                            let res = res.message.unwrap();
-                            let _ = msg.response_chn.send(Ok(res));
-                        }
-                        _ => {
-                            error!("unhandled websocket message: {:?}", read);
-                            let _ = msg.response_chn.send(Err(SignalError::WsError(WsError::ConnectionClosed)));
-                            break;
+                Some(mut msg) = read_receiver.recv() => {
+                    tokio::select! {
+                        Some(read) = ws_reader.next() => {
+                           match read {
+                                Ok(Message::Binary(data)) => {
+                                    let res = SignalResponse::decode(data.as_slice()).expect("failed to decode SignalResponse");
+                                    let signal = res.message.unwrap();
+                                    event!(Level::TRACE, "received: {:?}", signal);
+                                    let _ = msg.response_chn.send(Some(signal));
+                                }
+                                _ => {
+                                    event!(Level::ERROR, "unhandled websocket message {:?}", read);
+                                    let _ = msg.response_chn.send(None);
+                                }
+                            }
+                        },
+                        _ = poll_fn(|cx| msg.response_chn.poll_closed(cx)) => {
+                            continue; // Cancelled
+                        },
+                        else => {
+                            break; // Connection closed
                         }
                     }
                 },
@@ -174,13 +192,3 @@ impl SignalClient {
         }
     }
 }
-
-#[tokio::test]
-async fn test_test() {
-    env_logger::init();
-    let client = connect("ws://localhost:7880", "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJleHAiOjE2NzEyMzk4NjAsImlzcyI6IkFQSXpLYkFTaUNWYWtnSiIsIm5hbWUiOiJ0ZXN0IiwibmJmIjoxNjY0MDM5ODYwLCJzdWIiOiJ0ZXN0IiwidmlkZW8iOnsicm9vbUFkbWluIjp0cnVlLCJyb29tQ3JlYXRlIjp0cnVlLCJyb29tSm9pbiI6dHJ1ZX19.0Bee2jI2cSZveAbZ8MLc-ADoMYQ4l8IRxcAxpXAS6a8").await.unwrap();
-    let msg = client.recv().await.unwrap();
-
-    client.close().await;
-    info!("Received message {:?}", msg);
-}