fix(fullstack): prevent AlreadyBorrowed panic in WebSocket recv() (#5300)

Author: kyle-rader

examples/07-fullstack/websocket.rs

@@ -24,28 +24,38 @@ 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);
+            }
         }
     });
 
     rsx! {
         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; },

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}")
             }
         }
     }))

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"] }

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<Result<Websocket<In, Out, Enc>, CapturedError>>,
+    #[allow(clippy::type_complexity)]
+    connection: Resource<Result<Rc<Websocket<In, Out, Enc>>, CapturedError>>,
     waker: UseWaker<()>,
     status: Signal<WebsocketState>,
     status_read: ReadSignal<WebsocketState>,
@@ -161,38 +165,35 @@ impl<In, Out, E> UseWebsocket<In, Out, E> {
     /// 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
     ///
     /// To receive a message with a particular type, see the `.recv()` method instead.
     pub async fn recv_raw(&mut self) -> Result<Message, WebsocketError> {
         self.connect().await;
-
-        let result = self
-            .connection
-            .as_ref()
-            .as_deref()
-            .ok_or_else(WebsocketError::closed_away)?
-            .as_ref()
-            .map_err(|_| WebsocketError::AlreadyClosed)?
-            .recv_raw()
-            .await;
-
-        if let Err(WebsocketError::ConnectionClosed { .. }) = result.as_ref() {
-            self.received_shutdown();
+        let ws = self.get_connection()?;
+
+        // 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);
+
+        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<In, Out, E> UseWebsocket<In, Out, E> {
         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<In, Out, E> UseWebsocket<In, Out, E> {
             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<In, Out, E> UseWebsocket<In, Out, E> {
         _self.status.set(WebsocketState::Closed);
         _self.waker.wake(());
     }
+
+    /// Clone the `Rc<Websocket>` 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<Rc<Websocket<In, Out, E>>, 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<In, Out, E> Copy for UseWebsocket<In, Out, E> {}