From 4b76baba23905eca1fe8c09d9e9211b605cb2d68 Mon Sep 17 00:00:00 2001 From: Matt Mastracci Date: Tue, 19 Sep 2023 21:16:42 -0600 Subject: [PATCH] [WIP] split websocket --- Cargo.lock | 2 -- Cargo.toml | 2 +- cli/tests/unit/websocket_test.ts | 11 +++++--- ext/websocket/lib.rs | 46 ++++++++++++++++++++------------ 4 files changed, 38 insertions(+), 23 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e2f97747ac2a64..48b3ef744b6ad2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2292,8 +2292,6 @@ checksum = "25cbce373ec4653f1a01a31e8a5e5ec0c622dc27ff9c4e6606eefef5cbbed4a5" [[package]] name = "fastwebsockets" version = "0.4.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9e6185b6dc9dddc4db0dedd2e213047e93bcbf7a0fb092abc4c4e4f3195efdb4" dependencies = [ "base64 0.21.4", "hyper 0.14.27", diff --git a/Cargo.toml b/Cargo.toml index c978464d9f3382..594f28a7027fb2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -87,7 +87,7 @@ data-encoding = "2.3.3" dlopen = "0.1.8" encoding_rs = "=0.8.33" ecb = "=0.1.2" -fastwebsockets = "=0.4.4" +fastwebsockets = { path = "../fastwebsockets" } filetime = "0.2.16" flate2 = { version = "1.0.26", features = ["zlib-ng"], default-features = false } fs3 = "0.5.0" diff --git a/cli/tests/unit/websocket_test.ts b/cli/tests/unit/websocket_test.ts index 11f0fd7dc314f0..f2b6e08e145359 100644 --- a/cli/tests/unit/websocket_test.ts +++ b/cli/tests/unit/websocket_test.ts @@ -1,4 +1,5 @@ // Copyright 2018-2023 the Deno authors. All rights reserved. MIT license. +import console from "node:console"; import { assert, assertEquals, @@ -21,7 +22,7 @@ Deno.test(async function websocketConstructorTakeURLObjectAsParameter() { const promise = deferred(); const ws = new WebSocket(new URL("ws://localhost:4242/")); assertEquals(ws.url, "ws://localhost:4242/"); - ws.onerror = () => fail(); + ws.onerror = (e) => fail(Deno.inspect(e)); ws.onopen = () => ws.close(); ws.onclose = () => { promise.resolve(); @@ -123,8 +124,12 @@ Deno.test({ const { response, socket } = Deno.upgradeWebSocket(req); let called = false; socket.onopen = () => socket.send("Hello"); - socket.onmessage = () => { - assert(!called); + socket.onmessage = (msg: MessageEvent) => { + console.log(msg); + if (called) { + return; + } + // assert(!called); called = true; socket.send("bye"); socket.close(); diff --git a/ext/websocket/lib.rs b/ext/websocket/lib.rs index 479cae7ecc07c0..f6693d635fa704 100644 --- a/ext/websocket/lib.rs +++ b/ext/websocket/lib.rs @@ -41,17 +41,21 @@ use std::rc::Rc; use std::sync::Arc; use tokio::io::AsyncRead; use tokio::io::AsyncWrite; +use tokio::io::ReadHalf; +use tokio::io::WriteHalf; use tokio::net::TcpStream; use tokio_rustls::rustls::RootCertStore; use tokio_rustls::rustls::ServerName; use tokio_rustls::TlsConnector; use fastwebsockets::CloseCode; -use fastwebsockets::FragmentCollector; +use fastwebsockets::FragmentCollectorRead; use fastwebsockets::Frame; use fastwebsockets::OpCode; use fastwebsockets::Role; use fastwebsockets::WebSocket; +use fastwebsockets::WebSocketWrite; + mod stream; static USE_WRITEV: Lazy = Lazy::new(|| { @@ -330,12 +334,13 @@ pub struct ServerWebSocket { closed: Cell, buffer: Cell>>, string: Cell>, - ws: AsyncRefCell>, - tx_lock: AsyncRefCell<()>, + ws_read: AsyncRefCell>>, + ws_write: AsyncRefCell>>, } impl ServerWebSocket { fn new(ws: WebSocket) -> Self { + let (ws_read, ws_write) = ws.split(|s| tokio::io::split(s)); Self { buffered: Cell::new(0), error: Cell::new(None), @@ -343,8 +348,8 @@ impl ServerWebSocket { closed: Cell::new(false), buffer: Cell::new(None), string: Cell::new(None), - ws: AsyncRefCell::new(FragmentCollector::new(ws)), - tx_lock: AsyncRefCell::new(()), + ws_read: AsyncRefCell::new(FragmentCollectorRead::new(ws_read)), + ws_write: AsyncRefCell::new(ws_write), } } @@ -359,22 +364,22 @@ impl ServerWebSocket { } /// Reserve a lock, but don't wait on it. This gets us our place in line. - pub fn reserve_lock(self: &Rc) -> AsyncMutFuture<()> { - RcRef::map(self, |r| &r.tx_lock).borrow_mut() + fn reserve_lock( + self: &Rc, + ) -> AsyncMutFuture>> { + RcRef::map(self, |r| &r.ws_write).borrow_mut() } #[inline] - pub async fn write_frame( + async fn write_frame( self: &Rc, - lock: AsyncMutFuture<()>, + lock: AsyncMutFuture>>, frame: Frame<'_>, ) -> Result<(), AnyError> { - lock.await; - - // SAFETY: fastwebsockets only needs a mutable reference to the WebSocket - // to populate the write buffer. We encounter an await point when writing - // to the socket after the frame has already been written to the buffer. - let ws = unsafe { &mut *self.ws.as_ptr() }; + let mut ws = lock.await; + if ws.is_closed() { + return Ok(()); + } ws.write_frame(frame) .await .map_err(|err| type_error(err.to_string()))?; @@ -403,6 +408,7 @@ pub fn ws_create_server_stream( ws.set_writev(*USE_WRITEV); ws.set_auto_close(true); ws.set_auto_pong(true); + let rid = state.resource_table.add(ServerWebSocket::new(ws)); Ok(rid) } @@ -594,9 +600,15 @@ pub async fn op_ws_next_event( return MessageKind::Error as u16; } - let mut ws = RcRef::map(&resource, |r| &r.ws).borrow_mut().await; + let mut ws = RcRef::map(&resource, |r| &r.ws_read).borrow_mut().await; + let writer = RcRef::map(&resource, |r| &r.ws_write); + let mut sender = move |frame| { + let writer = writer.clone(); + async move { writer.borrow_mut().await.write_frame(frame).await } + }; loop { - let val = match ws.read_frame().await { + let res = ws.read_frame(&mut sender).await; + let val = match res { Ok(val) => val, Err(err) => { // No message was received, socket closed while we waited.