From 05cb440624e520fc79e11a7657c6738438b2e51e Mon Sep 17 00:00:00 2001
From: Ayush Singh <ayushdevel1325@gmail.com>
Date: Wed, 27 Dec 2023 15:31:23 +0530
Subject: [PATCH 1/2] Split websocket state

- Allow reading and writing to socket simultaneously, i.e. allow
  subscribing to new stuff while reading stream

Signed-off-by: Ayush Singh <ayushdevel1325@gmail.com>
---
 src/tokio_tungstenite.rs | 47 +++++++++++++++++++++++-----------------
 1 file changed, 27 insertions(+), 20 deletions(-)

diff --git a/src/tokio_tungstenite.rs b/src/tokio_tungstenite.rs
index 40cc1f1..8bc563a 100644
--- a/src/tokio_tungstenite.rs
+++ b/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
     }
 }

From e1e590e38c503ae3bacc827b215d5b7661c21eb4 Mon Sep 17 00:00:00 2001
From: Ayush Singh <ayushdevel1325@gmail.com>
Date: Thu, 28 Dec 2023 21:08:14 +0530
Subject: [PATCH 2/2] Make websocket module public

- This is important since websocket::Stream is in exposed by
  WebSocketState

Signed-off-by: Ayush Singh <ayushdevel1325@gmail.com>
---
 src/lib.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/lib.rs b/src/lib.rs
index e9bfbfe..2a6ea18 100644
--- a/src/lib.rs
+++ b/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;