split the ws

client/Cargo.toml

@@ -30,7 +30,7 @@ wasm-streams = "0.4.0"
 either = "1.9.0"
 tokio-util = { version = "0.7.10", features = ["io"] }
 async-compression = { version = "0.4.5", features = ["tokio", "gzip", "brotli"] }
-fastwebsockets = { version = "0.6.0" }
+fastwebsockets = { version = "0.6.0", features = ["simdutf8", "unstable-split"] }
 rand = "0.8.5"
 base64 = "0.21.7"
 

client/demo.js

@@ -73,6 +73,7 @@
         );
         while (true) {
             await ws.send("data");
+            await (new Promise((res, _) => setTimeout(res, 100)));
         }
     } else {
         let resp = await epoxy_client.fetch("https://httpbin.org/get");

client/src/websocket.rs

@@ -1,25 +1,22 @@
 use crate::*;
 
 use base64::{engine::general_purpose::STANDARD, Engine};
-use fastwebsockets::{CloseCode, Frame, OpCode, Payload, Role, WebSocket, WebSocketError};
+use fastwebsockets::{
+    CloseCode, FragmentCollectorRead, Frame, OpCode, Payload, Role, WebSocket, WebSocketWrite,
+};
 use http_body_util::Empty;
 use hyper::{
-    client::conn::http1 as hyper_conn,
     header::{CONNECTION, UPGRADE},
+    upgrade::Upgraded,
     StatusCode,
 };
 use js_sys::Function;
 use std::str::from_utf8;
-use tokio::sync::{mpsc, oneshot};
+use tokio::io::WriteHalf;
-
-enum EpxMsg {
-    SendText(String, oneshot::Sender<Result<(), WebSocketError>>),
-    Close,
-}
 
 #[wasm_bindgen]
 pub struct EpxWebSocket {
-    msg_sender: mpsc::Sender<EpxMsg>,
+    tx: WebSocketWrite<WriteHalf<TokioIo<Upgraded>>>,
     onerror: Function,
 }
 
@@ -43,14 +40,13 @@ impl EpxWebSocket {
         origin: String,
     ) -> Result<EpxWebSocket, JsError> {
         let onerr = onerror.clone();
-        let ret: Result<EpxWebSocket, JsError>  = async move {
+        let ret: Result<EpxWebSocket, JsError> = async move {
             let url = Uri::try_from(url).replace_err("Failed to parse URL")?;
             let host = url.host().replace_err("URL must have a host")?;
 
             let rand: [u8; 16] = rand::random();
             let key = STANDARD.encode(rand);
 
-
             let mut builder = Request::builder()
                 .method("GET")
                 .uri(url.clone())
@@ -71,8 +67,9 @@ impl EpxWebSocket {
 
             let (mut sender, conn) = Builder::new()
                 .title_case_headers(true)
-                .preserve_header_case(true).handshake::<TokioIo<EpxStream>, Empty<Bytes>>(TokioIo::new(stream))
+                .preserve_header_case(true)
-                    .await?;
+                .handshake::<TokioIo<EpxStream>, Empty<Bytes>>(TokioIo::new(stream))
+                .await?;
 
             wasm_bindgen_futures::spawn_local(async move {
                 if let Err(e) = conn.with_upgrades().await {
@@ -83,81 +80,52 @@ impl EpxWebSocket {
             let mut response = sender.send_request(req).await?;
             verify(&response)?;
 
-            let mut ws = WebSocket::after_handshake(
+            let ws = WebSocket::after_handshake(
                 TokioIo::new(hyper::upgrade::on(&mut response).await?),
                 Role::Client,
             );
 
-            let (msg_sender, mut rx) = mpsc::channel(1);
+            let (rx, tx) = ws.split(tokio::io::split);
+
+            let mut rx = FragmentCollectorRead::new(rx);
 
             wasm_bindgen_futures::spawn_local(async move {
-                loop {
+                while let Ok(frame) = rx
-                    tokio::select! {
+                    .read_frame(&mut |arg| async move {
-                        frame = ws.read_frame() => {
+                        error!(
-                            if let Ok(frame) = frame {
+                            "wtf is an obligated write {:?}, {:?}, {:?}",
-                                match frame.opcode {
+                            arg.fin, arg.opcode, arg.payload
-                                    OpCode::Text => {
+                        );
-                                        if let Ok(str) = from_utf8(&frame.payload) {
+                        Ok::<(), std::io::Error>(())
-                                            let _ = onmessage.call1(&JsValue::null(), &jval!(str));
+                    })
-                                        }
+                    .await
-                                    }
+                {
-                                    OpCode::Binary => {
+                    match frame.opcode {
-                                        let _ = onmessage.call1(
+                        OpCode::Text => {
-                                            &JsValue::null(),
+                            if let Ok(str) = from_utf8(&frame.payload) {
-                                            &jval!(Uint8Array::from(frame.payload.to_vec().as_slice())),
+                                let _ = onmessage.call1(&JsValue::null(), &jval!(str));
-                                        );
-                                    }
-                                    OpCode::Close => {
-                                        let _ = onclose.call0(&JsValue::null());
-                                        break;
-                                    }
-                                    _ => panic!("unknown opcode {:?}", frame.opcode),
-                                }
                             }
                         }
-                        msg = rx.recv() => {
+                        OpCode::Binary => {
-                            if let Some(msg) = msg {
+                            let _ = onmessage.call1(
-                                match msg {
+                                &JsValue::null(),
-                                    EpxMsg::SendText(payload, err) => {
+                                &jval!(Uint8Array::from(frame.payload.to_vec().as_slice())),
-                                        let _ = err.send(ws.write_frame(Frame::text(
+                            );
-                                            Payload::Owned(payload.as_bytes().to_vec()),
-                                        ))
-                                        .await);
-                                    }
-                                    EpxMsg::Close => break,
-                                }
-                            } else {
-                                break;
-                            }
                         }
+                        OpCode::Close => {
+                            let _ = onclose.call0(&JsValue::null());
+                            break;
+                        }
+                        _ => panic!("unknown opcode {:?}", frame.opcode),
                     }
                 }
-                let _ = ws
-                    .write_frame(Frame::close(CloseCode::Normal.into(), b""))
-                    .await;
             });
 
             onopen
                 .call0(&Object::default())
                 .replace_err("Failed to call onopen")?;
 
-            Ok(Self { msg_sender, onerror })
+            Ok(Self { tx, onerror })
-        }.await;
-        if let Err(ret) = ret {
-            let _ = onerr.call1(&JsValue::null(), &jval!(ret.clone()));
-            Err(ret)
-        } else {
-            ret
-        }
-    }
-
-    #[wasm_bindgen]
-    pub async fn send(&mut self, payload: String) -> Result<(), JsError> {
-        let onerr = self.onerror.clone();
-        let ret: Result<(), JsError> = async move {
-            let (tx, rx) = oneshot::channel();
-            self.msg_sender.send(EpxMsg::SendText(payload, tx)).await?;
-            Ok(rx.await??)
         }
         .await;
         if let Err(ret) = ret {
@@ -168,9 +136,26 @@ impl EpxWebSocket {
         }
     }
 
+    #[wasm_bindgen]
+    pub async fn send(&mut self, payload: String) -> Result<(), JsError> {
+        let onerr = self.onerror.clone();
+        let ret = self
+            .tx
+            .write_frame(Frame::text(Payload::Owned(payload.as_bytes().to_vec())))
+            .await;
+        if let Err(ret) = ret {
+            let _ = onerr.call1(&JsValue::null(), &jval!(format!("{}", ret)));
+            Err(ret.into())
+        } else {
+            Ok(ret?)
+        }
+    }
+
     #[wasm_bindgen]
     pub async fn close(&mut self) -> Result<(), JsError> {
-        self.msg_sender.send(EpxMsg::Close).await?;
+        self.tx
+            .write_frame(Frame::close(CloseCode::Normal.into(), b""))
+            .await?;
         Ok(())
     }
 }