Split websocket state

src/lib.rs

@@ -79,7 +79,7 @@
 
 mod utils;
 mod version;
-mod websocket;
+pub mod websocket;
 
 #[cfg(feature = "enable-tokio-tungstenite")]
 pub mod tokio_tungstenite;

src/tokio_tungstenite.rs

@@ -1,5 +1,6 @@
 use crate::websocket::Stream;
-use futures_util::SinkExt;
+use futures_util::stream::{SplitSink, SplitStream};
+use futures_util::{SinkExt, StreamExt};
 use tokio::io::{AsyncRead, AsyncWrite};
 use tokio::net::TcpStream;
 use tokio_tungstenite::{
@@ -15,7 +16,14 @@ pub struct BinanceWebSocketClient;
 impl BinanceWebSocketClient {
     pub async fn connect_async(
         url: &str,
-    ) -> Result<(WebSocketState<MaybeTlsStream<TcpStream>>, Response), Error> {
+    ) -> Result<
+        (
+            WebSocketState<MaybeTlsStream<TcpStream>>,
+            WebSocketReader<MaybeTlsStream<TcpStream>>,
+            Response,
+        ),
+        Error,
+    > {
         let (socket, response) = connect_async(Url::parse(&url).unwrap()).await?;
 
         log::info!("Connected to {}", url);
@@ -25,22 +33,33 @@ impl BinanceWebSocketClient {
             log::debug!("* {}", header);
         }
 
-        Ok((WebSocketState::new(socket), response))
+        let (sink, stream) = socket.split();
+
+        Ok((WebSocketState::new(sink), stream, response))
     }
 
-    pub async fn connect_async_default(
-    ) -> Result<(WebSocketState<MaybeTlsStream<TcpStream>>, Response), Error> {
+    pub async fn connect_async_default() -> Result<
+        (
+            WebSocketState<MaybeTlsStream<TcpStream>>,
+            WebSocketReader<MaybeTlsStream<TcpStream>>,
+            Response,
+        ),
+        Error,
+    > {
         BinanceWebSocketClient::connect_async("wss://stream.binance.com:9443/stream").await
     }
 }
 
+pub type WebSocketWriter<T> = SplitSink<WebSocketStream<T>, tokio_tungstenite::tungstenite::Message>;
+pub type WebSocketReader<T> = SplitStream<WebSocketStream<T>>;
+
 pub struct WebSocketState<T> {
-    socket: WebSocketStream<T>,
+    socket: WebSocketWriter<T>,
     id: u64,
 }
 
 impl<T: AsyncRead + AsyncWrite + Unpin> WebSocketState<T> {
-    pub fn new(socket: WebSocketStream<T>) -> Self {
+    pub fn new(socket: WebSocketWriter<T>) -> Self {
         Self { socket, id: 0 }
     }
 
@@ -128,18 +147,6 @@ impl<T: AsyncRead + AsyncWrite + Unpin> WebSocketState<T> {
     }
 
     pub async fn close(mut self) -> Result<(), Error> {
-        self.socket.close(None).await
-    }
-}
-
-impl<T> From<WebSocketState<T>> for WebSocketStream<T> {
-    fn from(conn: WebSocketState<T>) -> WebSocketStream<T> {
-        conn.socket
-    }
-}
-
-impl<T> AsMut<WebSocketStream<T>> for WebSocketState<T> {
-    fn as_mut(&mut self) -> &mut WebSocketStream<T> {
-        &mut self.socket
+        self.socket.close().await
     }
 }