diff --git a/examples/07-fullstack/websocket.rs b/examples/07-fullstack/websocket.rs index e96ea7aec..1c8830b73 100644 --- a/examples/07-fullstack/websocket.rs +++ b/examples/07-fullstack/websocket.rs @@ -24,21 +24,25 @@ fn main() { } fn app() -> Element { - // Track the messages we've received from the server. let mut messages = use_signal(std::vec::Vec::new); - // The `use_websocket` wraps the `WebSocket` connection and provides a reactive handle to easily - // send and receive messages and track the connection state. - // - // We can customize the websocket connection with the `WebSocketOptions` struct, allowing us to - // set things like custom headers, protocols, reconnection strategies, etc. - let mut socket = use_websocket(|| uppercase_ws("John Doe".into(), 30, WebSocketOptions::new())); + // This signal is read inside the use_websocket closure, making it a reactive dependency. + // Whenever it changes, the websocket will automatically re-connect. + let mut name = use_signal(|| "John Doe".to_string()); + + let mut socket = + use_websocket(move || uppercase_ws(name.cloned(), 30, WebSocketOptions::new())); - // Calling `.recv()` automatically waits for the connection to be established and deserializes - // messages as they arrive. use_future(move || async move { - while let Ok(msg) = socket.recv().await { - messages.push(msg); + loop { + // Wait for the socket to connect + _ = socket.connect().await; + + // Loop poll with recv. Throws an error when the connection closes, making it possible + // to run code before the socket re-connects when the name input changes + while let Ok(msg) = socket.recv().await { + messages.push(msg); + } } }); @@ -46,6 +50,12 @@ fn app() -> Element { h1 { "WebSocket Example" } p { "Type a message and see it echoed back in uppercase!" } p { "Connection status: {socket.status():?}" } + p { "Change your name to trigger a websocket re-connect" } + input { + placeholder: "Your name", + value: "{name}", + oninput: move |e| name.set(e.value()), + } input { placeholder: "Type a message", oninput: move |e| async move { _ = socket.send(ClientEvent::TextInput(e.value())).await; }, diff --git a/packages/cli/src/serve/proxy_ws.rs b/packages/cli/src/serve/proxy_ws.rs index 2eed6ed69..517f6b0ad 100644 --- a/packages/cli/src/serve/proxy_ws.rs +++ b/packages/cli/src/serve/proxy_ws.rs @@ -29,7 +29,8 @@ pub(crate) async fn proxy_websocket( match handle_ws_connection(client_ws, proxied_request).await { Ok(()) => tracing::trace!(dx_src = ?TraceSrc::Dev, "Websocket connection closed"), Err(e) => { - tracing::error!(dx_src = ?TraceSrc::Dev, "Error proxying websocket connection: {e}") + // Connection resets during shutdown are expected and not worth logging as errors + tracing::debug!(dx_src = ?TraceSrc::Dev, "Error proxying websocket connection: {e}") } } })) diff --git a/packages/fullstack/Cargo.toml b/packages/fullstack/Cargo.toml index 765c6e248..85dd4fca1 100644 --- a/packages/fullstack/Cargo.toml +++ b/packages/fullstack/Cargo.toml @@ -84,7 +84,7 @@ content_disposition = "0.4.0" [target.'cfg(not(target_arch = "wasm32"))'.dependencies] async-tungstenite = { version = "0.31.0", default-features = false, features = ["futures-03-sink"], optional = true } tungstenite = { version = "0.27", default-features = false, features = ["handshake"], optional = true } -tokio-util = { workspace = true, features = ["codec", "compat"] } +tokio-util = { workspace = true, features = ["codec", "compat", "rt"] } [dev-dependencies] dioxus = { workspace = true, features = ["fullstack", "router"] } diff --git a/packages/fullstack/src/payloads/websocket.rs b/packages/fullstack/src/payloads/websocket.rs index 49a820a7b..5aef670c3 100644 --- a/packages/fullstack/src/payloads/websocket.rs +++ b/packages/fullstack/src/payloads/websocket.rs @@ -42,6 +42,7 @@ use std::{ marker::PhantomData, pin::Pin, prelude::rust_2024::Future, + rc::Rc, task::{ready, Context, Poll}, }; @@ -89,7 +90,9 @@ pub fn use_websocket< // Wake up the `.recv()` calls waiting for the connection to be established waker.wake(()); - connection + // Wrap in Rc so we can clone it out of the Resource without holding + // a borrow guard across await points + connection.map(Rc::new) } }); @@ -113,7 +116,8 @@ where Out: 'static, Enc: 'static, { - connection: Resource, CapturedError>>, + #[allow(clippy::type_complexity)] + connection: Resource>, CapturedError>>, waker: UseWaker<()>, status: Signal, status_read: ReadSignal, @@ -161,15 +165,7 @@ impl UseWebsocket { /// To send a message with a particular type, see the `.send()` method instead. pub async fn send_raw(&self, msg: Message) -> Result<(), WebsocketError> { self.connect().await; - - self.connection - .as_ref() - .as_deref() - .ok_or_else(WebsocketError::closed_away)? - .as_ref() - .map_err(|_| WebsocketError::AlreadyClosed)? - .send_raw(msg) - .await + self.get_connection()?.send_raw(msg).await } /// Receive a raw message from the WebSocket connection @@ -177,22 +173,27 @@ impl UseWebsocket { /// To receive a message with a particular type, see the `.recv()` method instead. pub async fn recv_raw(&mut self) -> Result { self.connect().await; + let ws = self.get_connection()?; - let result = self - .connection - .as_ref() - .as_deref() - .ok_or_else(WebsocketError::closed_away)? - .as_ref() - .map_err(|_| WebsocketError::AlreadyClosed)? - .recv_raw() - .await; + // Race the recv against the waker — if the connection is being recreated + // (e.g. a reactive dependency changed), the waker fires and we return an error + // so the caller's loop can restart and pick up the new connection. + let recv_fut = ws.recv_raw(); + let waker_fut = self.waker.wait(); + futures::pin_mut!(recv_fut, waker_fut); - if let Err(WebsocketError::ConnectionClosed { .. }) = result.as_ref() { - self.received_shutdown(); + match futures::future::select(recv_fut, waker_fut).await { + futures::future::Either::Left((recv_result, _)) => { + if let Err(WebsocketError::ConnectionClosed { .. }) = recv_result.as_ref() { + self.received_shutdown(); + } + recv_result + } + futures::future::Either::Right(_) => Err(WebsocketError::ConnectionClosed { + code: CloseCode::Away, + description: "Connection replaced by a new one".to_string(), + }), } - - result } pub async fn send(&self, msg: In) -> Result<(), WebsocketError> @@ -219,22 +220,24 @@ impl UseWebsocket { E: Encoding, { self.connect().await; + let ws = self.get_connection()?; - let result = self - .connection - .as_ref() - .as_deref() - .ok_or_else(WebsocketError::closed_away)? - .as_ref() - .map_err(|_| WebsocketError::AlreadyClosed)? - .recv() - .await; + let recv_fut = ws.recv(); + let waker_fut = self.waker.wait(); + futures::pin_mut!(recv_fut, waker_fut); - if let Err(WebsocketError::ConnectionClosed { .. }) = result.as_ref() { - self.received_shutdown(); + match futures::future::select(recv_fut, waker_fut).await { + futures::future::Either::Left((recv_result, _)) => { + if let Err(WebsocketError::ConnectionClosed { .. }) = recv_result.as_ref() { + self.received_shutdown(); + } + recv_result + } + futures::future::Either::Right(_) => Err(WebsocketError::ConnectionClosed { + code: CloseCode::Away, + description: "Connection replaced by a new one".to_string(), + }), } - - result } /// Set the WebSocket connection. @@ -247,7 +250,8 @@ impl UseWebsocket { Err(_) => self.status.set(WebsocketState::FailedToConnect), } - self.connection.set(Some(socket.map_err(|e| e.into()))); + self.connection + .set(Some(socket.map(Rc::new).map_err(|e| e.into()))); self.waker.wake(()); } @@ -257,6 +261,20 @@ impl UseWebsocket { _self.status.set(WebsocketState::Closed); _self.waker.wake(()); } + + /// Clone the `Rc` out of the Resource using peek, so we don't hold a borrow + /// guard across await points. This prevents AlreadyBorrowed panics when the Resource + /// tries to write while recv() is awaiting. + #[allow(clippy::result_large_err)] + fn get_connection(&self) -> Result>, WebsocketError> { + self.connection.with_peek(|opt| { + opt.as_ref() + .ok_or_else(WebsocketError::closed_away)? + .as_ref() + .map(Rc::clone) + .map_err(|_| WebsocketError::AlreadyClosed) + }) + } } impl Copy for UseWebsocket {}