From 6b70ba7cd2cb38a14c52f804df1ba9933677e8e6 Mon Sep 17 00:00:00 2001 From: Evan Almloff Date: Wed, 14 May 2025 01:28:43 -0500 Subject: [PATCH] Fix websocket server functions and add an example (#4107) * Fix websocket server function * Add websocket test * add a websocket example * fix clippy --- Cargo.lock | 11 ++ Cargo.toml | 1 + examples/fullstack-websockets/.gitignore | 4 + examples/fullstack-websockets/Cargo.toml | 16 ++ examples/fullstack-websockets/src/main.rs | 69 ++++++++ packages/playwright-tests/fullstack.spec.js | 8 + .../playwright-tests/fullstack/Cargo.toml | 1 + .../playwright-tests/fullstack/src/main.rs | 49 +++++- packages/server/Cargo.toml | 1 + packages/server/src/context.rs | 149 ++++++++++++++++-- 10 files changed, 297 insertions(+), 12 deletions(-) create mode 100644 examples/fullstack-websockets/.gitignore create mode 100644 examples/fullstack-websockets/Cargo.toml create mode 100644 examples/fullstack-websockets/src/main.rs diff --git a/Cargo.lock b/Cargo.lock index 13d329560..70b9b8183 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4334,6 +4334,7 @@ name = "dioxus-playwright-fullstack-test" version = "0.1.0" dependencies = [ "dioxus", + "futures", "serde", "tokio", ] @@ -4469,6 +4470,7 @@ dependencies = [ "dioxus-lib", "dioxus-router", "dioxus-ssr", + "enumset", "futures-channel", "futures-util", "generational-box", @@ -5565,6 +5567,15 @@ dependencies = [ "tokio", ] +[[package]] +name = "fullstack-websocket-example" +version = "0.1.0" +dependencies = [ + "dioxus", + "futures", + "tokio", +] + [[package]] name = "funty" version = "2.0.0" diff --git a/Cargo.toml b/Cargo.toml index 5072c22bd..88e294507 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -109,6 +109,7 @@ members = [ "examples/fullstack-streaming", "examples/fullstack-desktop", "examples/fullstack-auth", + "examples/fullstack-websockets", # Playwright tests "packages/playwright-tests/liveview", diff --git a/examples/fullstack-websockets/.gitignore b/examples/fullstack-websockets/.gitignore new file mode 100644 index 000000000..21fff11dd --- /dev/null +++ b/examples/fullstack-websockets/.gitignore @@ -0,0 +1,4 @@ +dist +target +static +.dioxus \ No newline at end of file diff --git a/examples/fullstack-websockets/Cargo.toml b/examples/fullstack-websockets/Cargo.toml new file mode 100644 index 000000000..bcbd49fdc --- /dev/null +++ b/examples/fullstack-websockets/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "fullstack-websocket-example" +version = "0.1.0" +edition = "2021" +publish = false + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +dioxus = { workspace = true, features = ["fullstack"] } +futures.workspace = true +tokio = { workspace = true, features = ["full"], optional = true } + +[features] +server = ["dioxus/server", "dep:tokio"] +web = ["dioxus/web"] diff --git a/examples/fullstack-websockets/src/main.rs b/examples/fullstack-websockets/src/main.rs new file mode 100644 index 000000000..fc755f739 --- /dev/null +++ b/examples/fullstack-websockets/src/main.rs @@ -0,0 +1,69 @@ +#![allow(non_snake_case)] +use dioxus::prelude::{ + server_fn::{codec::JsonEncoding, BoxedStream, Websocket}, + *, +}; +use futures::{channel::mpsc, SinkExt, StreamExt}; + +fn main() { + launch(app); +} + +fn app() -> Element { + let mut uppercase = use_signal(String::new); + let mut uppercase_channel = use_signal(|| None); + + // Start the websocket connection in a background task + use_future(move || async move { + let (tx, rx) = mpsc::channel(1); + let mut receiver = uppercase_ws(rx.into()).await.unwrap(); + // Store the channel in a signal for use in the input handler + uppercase_channel.set(Some(tx)); + // Whenever we get a message from the server, update the uppercase signal + while let Some(Ok(msg)) = receiver.next().await { + uppercase.set(msg); + } + }); + + rsx! { + input { + oninput: move |e| async move { + if let Some(mut uppercase_channel) = uppercase_channel() { + let msg = e.value(); + uppercase_channel.send(Ok(msg)).await.unwrap(); + } + }, + } + "Uppercase: {uppercase}" + } +} + +// The server macro accepts a protocol parameter which implements the protocol trait. The protocol +// controls how the inputs and outputs are encoded when handling the server function. In this case, +// the websocket protocol can encode a stream input and stream output where messages are +// serialized as JSON +#[server(protocol = Websocket)] +async fn uppercase_ws( + input: BoxedStream, +) -> Result, ServerFnError> { + let mut input = input; + + // Create a channel with the output of the websocket + let (mut tx, rx) = mpsc::channel(1); + + // Spawn a task that processes the input stream and sends any new messages to the output + tokio::spawn(async move { + while let Some(msg) = input.next().await { + if tx + .send(msg.map(|msg| msg.to_ascii_uppercase())) + .await + .is_err() + { + break; + } + } + }); + + // Return the output stream + Ok(rx.into()) +} diff --git a/packages/playwright-tests/fullstack.spec.js b/packages/playwright-tests/fullstack.spec.js index 07d3468fa..3c688ba57 100644 --- a/packages/playwright-tests/fullstack.spec.js +++ b/packages/playwright-tests/fullstack.spec.js @@ -68,3 +68,11 @@ test("document elements", async ({ page }) => { const main = page.locator("#main"); await expect(main).toHaveCSS("font-family", "Roboto"); }); + +test("websockets", async ({ page }) => { + await page.goto("http://localhost:3333"); + // wait until the websocket div is mounted + const wsDiv = page.locator("div#websocket-div"); + await expect(wsDiv).toHaveText("Received: HELLO WORLD"); +}); + diff --git a/packages/playwright-tests/fullstack/Cargo.toml b/packages/playwright-tests/fullstack/Cargo.toml index 6fc29ab58..9f96d5863 100644 --- a/packages/playwright-tests/fullstack/Cargo.toml +++ b/packages/playwright-tests/fullstack/Cargo.toml @@ -8,6 +8,7 @@ publish = false [dependencies] dioxus = { workspace = true, features = ["fullstack"] } +futures.workspace = true serde = "1.0.218" tokio = { workspace = true, features = ["full"], optional = true } diff --git a/packages/playwright-tests/fullstack/src/main.rs b/packages/playwright-tests/fullstack/src/main.rs index 46311d8b9..f5689758d 100644 --- a/packages/playwright-tests/fullstack/src/main.rs +++ b/packages/playwright-tests/fullstack/src/main.rs @@ -5,7 +5,14 @@ // - Hydration #![allow(non_snake_case)] -use dioxus::{prelude::*, CapturedError}; +use dioxus::{ + prelude::{ + server_fn::{codec::JsonEncoding, BoxedStream, Websocket}, + *, + }, + CapturedError, +}; +use futures::{channel::mpsc, SinkExt, StreamExt}; fn main() { dioxus::LaunchBuilder::new() @@ -43,6 +50,7 @@ fn app() -> Element { OnMounted {} DefaultServerFnCodec {} DocumentElements {} + WebSockets {} } } @@ -156,3 +164,42 @@ fn DocumentElements() -> Element { document::Style { id: "style-head", "body {{ font-family: 'Roboto'; }}" } } } + +#[server(protocol = Websocket)] +async fn echo_ws( + input: BoxedStream, +) -> Result, ServerFnError> { + let mut input = input; + + let (mut tx, rx) = mpsc::channel(1); + + tokio::spawn(async move { + while let Some(msg) = input.next().await { + let _ = tx.send(msg.map(|msg| msg.to_ascii_uppercase())).await; + } + }); + + Ok(rx.into()) +} + +/// This component tests websocket server functions +#[component] +fn WebSockets() -> Element { + let mut received = use_signal(String::new); + use_future(move || async move { + let (mut tx, rx) = mpsc::channel(1); + let mut receiver = echo_ws(rx.into()).await.unwrap(); + tx.send(Ok("hello world".to_string())).await.unwrap(); + while let Some(Ok(msg)) = receiver.next().await { + println!("Received: {}", msg); + received.set(msg); + } + }); + + rsx! { + div { + id: "websocket-div", + "Received: {received}" + } + } +} diff --git a/packages/server/Cargo.toml b/packages/server/Cargo.toml index f320df25b..8560bbb5b 100644 --- a/packages/server/Cargo.toml +++ b/packages/server/Cargo.toml @@ -42,6 +42,7 @@ tracing-futures = { workspace = true } once_cell = { workspace = true } async-trait = { workspace = true } serde = { workspace = true } +enumset = "1.1.5" futures-util = { workspace = true } futures-channel = { workspace = true } diff --git a/packages/server/src/context.rs b/packages/server/src/context.rs index be5902d7d..3b51f73b4 100644 --- a/packages/server/src/context.rs +++ b/packages/server/src/context.rs @@ -1,10 +1,44 @@ +use enumset::{EnumSet, EnumSetType}; use parking_lot::RwLock; use std::any::Any; use std::collections::HashMap; +use std::sync::atomic::AtomicU32; use std::sync::Arc; type SendSyncAnyMap = std::collections::HashMap; +#[derive(EnumSetType)] +enum ResponsePartsModified { + Version, + Headers, + Status, + Extensions, + Body, +} + +struct AtomicResponsePartsModified { + modified: AtomicU32, +} + +impl AtomicResponsePartsModified { + fn new() -> Self { + Self { + modified: AtomicU32::new(EnumSet::::empty().as_u32()), + } + } + + fn set(&self, part: ResponsePartsModified) { + let modified = + EnumSet::from_u32(self.modified.load(std::sync::atomic::Ordering::Relaxed)) | part; + self.modified + .store(modified.as_u32(), std::sync::atomic::Ordering::Relaxed); + } + + fn is_modified(&self, part: ResponsePartsModified) -> bool { + self.modified.load(std::sync::atomic::Ordering::Relaxed) & (1 << part as usize) != 0 + } +} + /// A shared context for server functions that contains information about the request and middleware state. /// /// You should not construct this directly inside components or server functions. Instead use [`server_context()`] to get the server context from the current request. @@ -24,6 +58,7 @@ type SendSyncAnyMap = std::collections::HashMap; #[derive(Clone)] pub struct DioxusServerContext { shared_context: Arc>, + response_parts_modified: Arc, response_parts: Arc>, pub(crate) parts: Arc>, response_sent: Arc, @@ -48,6 +83,7 @@ impl Default for DioxusServerContext { fn default() -> Self { Self { shared_context: Arc::new(RwLock::new(HashMap::new())), + response_parts_modified: Arc::new(AtomicResponsePartsModified::new()), response_parts: Arc::new(RwLock::new( http::response::Response::new(()).into_parts().0, )), @@ -59,7 +95,7 @@ impl Default for DioxusServerContext { mod server_fn_impl { use super::*; - use parking_lot::{RwLockReadGuard, RwLockWriteGuard}; + use parking_lot::{MappedRwLockWriteGuard, RwLockReadGuard, RwLockWriteGuard}; use std::any::{Any, TypeId}; impl DioxusServerContext { @@ -68,6 +104,7 @@ mod server_fn_impl { Self { parts: Arc::new(RwLock::new(parts)), shared_context: Arc::new(RwLock::new(SendSyncAnyMap::new())), + response_parts_modified: Arc::new(AtomicResponsePartsModified::new()), response_parts: std::sync::Arc::new(RwLock::new( http::response::Response::new(()).into_parts().0, )), @@ -81,6 +118,7 @@ mod server_fn_impl { Self { parts, shared_context: Arc::new(RwLock::new(SendSyncAnyMap::new())), + response_parts_modified: Arc::new(AtomicResponsePartsModified::new()), response_parts: std::sync::Arc::new(RwLock::new( http::response::Response::new(()).into_parts().0, )), @@ -178,7 +216,7 @@ mod server_fn_impl { self.response_parts.read() } - /// Get the response parts from the server context + /// Get the headers from the server context mutably /// #[doc = include_str!("../docs/request_origin.md")] /// @@ -189,13 +227,82 @@ mod server_fn_impl { /// #[server] /// async fn set_headers() -> Result<(), ServerFnError> { /// let server_context = server_context(); - /// server_context.response_parts_mut() - /// .headers + /// server_context.headers_mut() /// .insert("Cookie", http::HeaderValue::from_static("dioxus=fullstack")); /// Ok(()) /// } /// ``` - pub fn response_parts_mut(&self) -> RwLockWriteGuard<'_, http::response::Parts> { + pub fn headers_mut(&self) -> MappedRwLockWriteGuard<'_, http::HeaderMap> { + self.response_parts_modified + .set(ResponsePartsModified::Headers); + RwLockWriteGuard::map(self.response_parts_mut(), |parts| &mut parts.headers) + } + + /// Get the status from the server context mutably + /// + #[doc = include_str!("../docs/request_origin.md")] + /// + /// # Example + /// + /// ```rust, no_run + /// # use dioxus::prelude::*; + /// #[server] + /// async fn set_status() -> Result<(), ServerFnError> { + /// let server_context = server_context(); + /// *server_context.status_mut() = http::StatusCode::INTERNAL_SERVER_ERROR; + /// Ok(()) + /// } + /// ``` + pub fn status_mut(&self) -> MappedRwLockWriteGuard<'_, http::StatusCode> { + self.response_parts_modified + .set(ResponsePartsModified::Status); + RwLockWriteGuard::map(self.response_parts_mut(), |parts| &mut parts.status) + } + + /// Get the version from the server context mutably + /// + #[doc = include_str!("../docs/request_origin.md")] + /// + /// # Example + /// + /// ```rust, no_run + /// # use dioxus::prelude::*; + /// #[server] + /// async fn set_version() -> Result<(), ServerFnError> { + /// let server_context = server_context(); + /// *server_context.version_mut() = http::Version::HTTP_2; + /// Ok(()) + /// } + /// ``` + pub fn version_mut(&self) -> MappedRwLockWriteGuard<'_, http::Version> { + self.response_parts_modified + .set(ResponsePartsModified::Version); + RwLockWriteGuard::map(self.response_parts_mut(), |parts| &mut parts.version) + } + + /// Get the extensions from the server context mutably + /// + #[doc = include_str!("../docs/request_origin.md")] + /// + /// # Example + /// + /// ```rust, no_run + /// # use dioxus::prelude::*; + /// #[server] + /// async fn set_version() -> Result<(), ServerFnError> { + /// let server_context = server_context(); + /// *server_context.version_mut() = http::Version::HTTP_2; + /// Ok(()) + /// } + /// ``` + pub fn extensions_mut(&self) -> MappedRwLockWriteGuard<'_, http::Extensions> { + self.response_parts_modified + .set(ResponsePartsModified::Extensions); + RwLockWriteGuard::map(self.response_parts_mut(), |parts| &mut parts.extensions) + } + + /// Get the response parts mutably. This does not track what parts have been written to so it should not be exposed publicly. + fn response_parts_mut(&self) -> RwLockWriteGuard<'_, http::response::Parts> { if self .response_sent .load(std::sync::atomic::Ordering::Relaxed) @@ -280,13 +387,33 @@ mod server_fn_impl { .store(true, std::sync::atomic::Ordering::Relaxed); let parts = self.response_parts.read(); - let mut_headers = response.headers_mut(); - for (key, value) in parts.headers.iter() { - mut_headers.insert(key, value.clone()); + if self + .response_parts_modified + .is_modified(ResponsePartsModified::Headers) + { + let mut_headers = response.headers_mut(); + for (key, value) in parts.headers.iter() { + mut_headers.insert(key, value.clone()); + } + } + if self + .response_parts_modified + .is_modified(ResponsePartsModified::Status) + { + *response.status_mut() = parts.status; + } + if self + .response_parts_modified + .is_modified(ResponsePartsModified::Version) + { + *response.version_mut() = parts.version; + } + if self + .response_parts_modified + .is_modified(ResponsePartsModified::Extensions) + { + response.extensions_mut().extend(parts.extensions.clone()); } - *response.status_mut() = parts.status; - *response.version_mut() = parts.version; - response.extensions_mut().extend(parts.extensions.clone()); } } }