rustfmt & signal_client improvememts

- handle ping message
- recv now only returns on SignalResponse

crates/livekit-core/src/signal_client.rs

@@ -1,20 +1,21 @@
-use futures::future::poll_fn;
-use futures_util::stream::{SplitSink, SplitStream};
-use futures_util::{SinkExt, StreamExt};
-use prost::Message as ProstMessage;
 use std::fmt::{Debug, Formatter};
+
+use futures::future::poll_fn;
+use futures_util::{SinkExt, StreamExt};
+use futures_util::stream::{SplitSink, SplitStream};
+use prost::Message as ProstMessage;
 use thiserror::Error;
 use tokio::net::TcpStream;
 use tokio::sync::{mpsc, oneshot};
 use tokio::task::JoinHandle;
-use tokio_tungstenite::tungstenite::{
-    protocol::frame::{coding::CloseCode, CloseFrame},
-    Error as WsError, Message,
-};
 use tokio_tungstenite::{connect_async, MaybeTlsStream, WebSocketStream};
-use tracing::{event, span, Level};
+use tokio_tungstenite::tungstenite::{Error as WsError, Message};
+use tokio_tungstenite::tungstenite::protocol::CloseFrame;
+use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode;
+use tracing::{event, Level};
 
 use crate::proto::{signal_request, signal_response, SignalRequest, SignalResponse};
+use crate::signal_client::SendMessage::Pong;
 
 pub const PROTOCOL_VERSION: u32 = 8;
 
@@ -32,21 +33,26 @@ type SignalResult<T> = Result<T, SignalError>;
 type WebSocket = WebSocketStream<MaybeTlsStream<TcpStream>>;
 
 #[derive(Debug)]
-struct RecvMessage {
-    response_chn: oneshot::Sender<Option<signal_response::Message>>,
+enum RecvMessage {
+    Signal {
+        response_chn: oneshot::Sender<Option<signal_response::Message>>,
+    },
 }
 
 #[derive(Debug)]
-struct SendMessage {
-    signal: signal_request::Message,
-    response_chn: oneshot::Sender<SignalResult<()>>,
+enum SendMessage {
+    Signal {
+        signal: signal_request::Message,
+        response_chn: oneshot::Sender<SignalResult<()>>,
+    },
+    Pong {
+        ping_data: Vec<u8>,
+    },
 }
 
 pub struct SignalClient {
-    read_sender: mpsc::Sender<RecvMessage>,
-    write_sender: mpsc::Sender<SendMessage>,
-    write_shutdown_sender: oneshot::Sender<()>,
-    read_shutdown_sender: oneshot::Sender<()>,
+    read_tx: mpsc::Sender<RecvMessage>,
+    write_tx: mpsc::Sender<SendMessage>,
     read_handle: JoinHandle<()>,
     write_handle: JoinHandle<()>,
 }
@@ -57,7 +63,7 @@ impl Debug for SignalClient {
     }
 }
 
-#[tracing::instrument]
+#[tracing::instrument(skip(url, token))]
 pub async fn connect(url: &str, token: &str) -> SignalResult<SignalClient> {
     let mut lk_url = url::Url::parse(url)?;
     lk_url.set_path("/rtc");
@@ -74,21 +80,13 @@ pub async fn connect(url: &str, token: &str) -> SignalResult<SignalClient> {
 
     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();
 
-    let read_handle = tokio::spawn(SignalClient::ws_read(read_rx, ws_reader, read_shutdown_rx));
-    let write_handle = tokio::spawn(SignalClient::ws_write(
-        write_rx,
-        ws_writer,
-        write_shutdown_rx,
-    ));
+    let read_handle = tokio::spawn(SignalClient::ws_read(read_rx, ws_reader, write_tx.clone()));
+    let write_handle = tokio::spawn(SignalClient::ws_write(write_rx, ws_writer));
 
     Ok(SignalClient {
-        read_sender: read_tx,
-        write_sender: write_tx,
-        write_shutdown_sender: write_shutdown_tx,
-        read_shutdown_sender: read_shutdown_tx,
+        read_tx,
+        write_tx,
         read_handle,
         write_handle,
     })
@@ -96,99 +94,115 @@ pub async fn connect(url: &str, token: &str) -> SignalResult<SignalClient> {
 
 impl SignalClient {
     pub async fn close(self) {
-        let _ = self.write_shutdown_sender.send(());
-        let _ = self.write_handle.await;
-        let _ = self.read_shutdown_sender.send(());
+        drop(self.read_tx);
+        drop(self.write_tx);
+
         let _ = self.read_handle.await;
+        let _ = self.write_handle.await;
     }
 
     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;
+        let msg = RecvMessage::Signal { response_chn: send };
+        let _ = self.read_tx.send(msg).await;
         recv.await.expect("channel closed")
     }
 
     pub async fn send(&self, signal: signal_request::Message) -> SignalResult<()> {
         let (send, recv) = oneshot::channel();
-        let msg = SendMessage {
+        let msg = SendMessage::Signal {
             signal,
             response_chn: send,
         };
-        let _ = self.write_sender.send(msg).await;
+        let _ = self.write_tx.send(msg).await;
         recv.await.expect("channel closed")
     }
 
-    #[tracing::instrument]
     async fn ws_write(
         mut write_receiver: mpsc::Receiver<SendMessage>,
         mut ws_writer: SplitSink<WebSocket, Message>,
-        mut shutdown_receiver: oneshot::Receiver<()>,
     ) {
-        loop {
-            tokio::select! {
-                Some(msg) = write_receiver.recv() => {
-                    event!(Level::TRACE, "sending: {:?}", msg.signal);
+        while let Some(msg) = write_receiver.recv().await {
+            match msg {
+                SendMessage::Signal {
+                    signal,
+                    response_chn,
+                } => {
+                    event!(Level::TRACE, "sending: {:?}", signal);
 
                     let req = SignalRequest {
-                        message: Some(msg.signal),
+                        message: Some(signal),
                     };
 
                     let write_res = ws_writer.send(Message::Binary(req.encode_to_vec())).await;
                     if let Err(err) = write_res {
-                        event!(Level::ERROR, "failed to send message to ws: {:?}", err);
-                        let _ = msg.response_chn.send(Err(err.into()));
+                        event!(Level::ERROR, "failed to send signal: {:?}", err);
+                        let _ = response_chn.send(Err(err.into()));
                         break;
                     }
 
-                    let _ = msg.response_chn.send(Ok(()));
+                    let _ = response_chn.send(Ok(()));
+                }
+                Pong { ping_data } => {
+                    if let Err(err) = ws_writer.send(Message::Pong(ping_data)).await {
+                        event!(Level::ERROR, "failed to send pong message: {:?}", err);
+                    }
+                }
+            }
+        }
+
+        let _ = ws_writer
+            .send(Message::Close(Some(CloseFrame {
+                code: CloseCode::Normal,
+                reason: "disconnected by client".into(),
+            })))
+            .await;
+        let _ = ws_writer.flush().await;
+    }
+
+    async fn ws_read(
+        mut read_receiver: mpsc::Receiver<RecvMessage>,
+        mut ws_reader: SplitStream<WebSocket>,
+        write_tx: mpsc::Sender<SendMessage>,
+    ) {
+        while let Some(RecvMessage::Signal { mut response_chn }) = read_receiver.recv().await {
+            tokio::select! {
+                read = Self::handle_msg(&mut ws_reader, &write_tx) => {
+                    let _ = response_chn.send(read);
+                }
+                _ = poll_fn(|cx| response_chn.poll_closed(cx)) => {
+                    continue; // Cancelled
                 },
-                _ = (&mut shutdown_receiver) => {
-                    let _ = ws_writer.send(Message::Close(Some(CloseFrame {
-                        code: CloseCode::Normal,
-                        reason: "disconnected by client".into()
-                    }))).await;
-                    let _ = ws_writer.flush().await;
-                    break;
+            }
+        }
+    }
+
+    async fn handle_msg(
+        ws_reader: &mut SplitStream<WebSocket>,
+        write_tx: &mpsc::Sender<SendMessage>,
+    ) -> Option<signal_response::Message> {
+        loop {
+            let read = ws_reader.next().await?;
+            match read {
+                Ok(Message::Binary(data)) => {
+                    let res = SignalResponse::decode(data.as_slice())
+                        .expect("failed to decode SignalResponse");
+                    event!(Level::TRACE, "received: {:?}", res);
+                    return Some(res.message.unwrap());
+                }
+                Ok(Message::Ping(data)) => {
+                    let _ = write_tx.send(Pong { ping_data: data });
+                    continue;
+                }
+                Ok(Message::Close(close)) => {
+                    event!(Level::DEBUG, "server closed the connection: {:?}", close);
+                    return None;
+                }
+                _ => {
+                    event!(Level::ERROR, "unhandled websocket message {:?}", read);
+                    return None;
                 }
             }
         }
     }
-
-    #[tracing::instrument]
-    async fn ws_read(
-        mut read_receiver: mpsc::Receiver<RecvMessage>,
-        mut ws_reader: SplitStream<WebSocket>,
-        mut shutdown_receiver: oneshot::Receiver<()>,
-    ) {
-        loop {
-            tokio::select! {
-                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
-                        }
-                    }
-                },
-                _ = (&mut shutdown_receiver) => break
-            }
-        }
-    }
 }