Added websocket server

This commit is contained in:
xeovalyte 2024-04-06 17:24:27 +02:00
parent b9d5425ac5
commit 683233a4b5
No known key found for this signature in database
6 changed files with 107 additions and 5 deletions

View File

@ -7,7 +7,7 @@ edition = "2021"
crate-type = ["cdylib", "rlib"]
[dependencies]
axum = { version = "0.7", optional = true }
axum = { version = "0.7", optional = true, features = [ "ws", "macros" ] }
console_error_panic_hook = "0.1"
leptos = { version = "0.6", features = [] }
leptos_axum = { version = "0.6", optional = true }
@ -26,6 +26,7 @@ serde_json = "1.0.115"
cfg-if = "1.0.0"
once_cell = "1.19.0"
futures = "0.3.30"
uuid = "1.8.0"
[features]
hydrate = ["leptos/hydrate", "leptos_meta/hydrate", "leptos_router/hydrate"]

View File

@ -1,11 +1,16 @@
use application::util::websocket::{server::AppState, server::WebSocketState};
#[cfg(feature = "ssr")]
#[tokio::main]
async fn main() {
use application::app::*;
use application::fileserv::file_and_error_handler;
use axum::Router;
use application::{app::*, util::websocket::server};
use axum::{routing::get, Router};
use leptos::*;
use leptos_axum::{generate_route_list, LeptosRoutes};
use std::collections::HashSet;
use std::sync::Arc;
use tokio::sync::{broadcast, Mutex};
application::util::surrealdb::connect()
.await
@ -21,11 +26,23 @@ async fn main() {
let addr = leptos_options.site_addr;
let routes = generate_route_list(App);
let client_set = Arc::new(Mutex::new(HashSet::<uuid::Uuid>::new()));
let (tx, _rx) = broadcast::channel(100);
let websocket_state = WebSocketState { client_set, tx };
let app_state = AppState {
websocket_state: websocket_state.into(),
routes: routes.clone(),
leptos_options,
};
// build our application with a route
let app = Router::new()
.leptos_routes(&leptos_options, routes, App)
.route("/ws", get(server::websocket_handler))
.leptos_routes(&app_state, routes, App)
.fallback(file_and_error_handler)
.with_state(leptos_options);
.with_state(app_state);
let listener = tokio::net::TcpListener::bind(&addr).await.unwrap();
logging::log!("listening on http://{}", &addr);

View File

@ -4,6 +4,7 @@ use leptos_router::ActionForm;
cfg_if::cfg_if! {
if #[cfg(feature = "ssr")] {
use crate::util::surrealdb::{DB, schemas};
use crate::util::websocket::server;
use leptos::logging;
}
}

View File

@ -1 +1,2 @@
pub mod surrealdb;
pub mod websocket;

View File

@ -0,0 +1,2 @@
#[cfg(feature = "ssr")]
pub mod server;

View File

@ -0,0 +1,80 @@
use axum::{
extract::{
ws::{Message, WebSocket, WebSocketUpgrade},
State,
},
response::IntoResponse,
};
use futures::{sink::SinkExt, stream::StreamExt};
use leptos::LeptosOptions;
use leptos_router::RouteListing;
use tokio::sync::{broadcast, Mutex};
use std::{collections::HashSet, sync::Arc};
#[derive(Clone)]
pub struct WebSocketState {
pub client_set: Arc<Mutex<HashSet<uuid::Uuid>>>,
pub tx: broadcast::Sender<String>,
}
#[derive(Clone, axum::extract::FromRef)]
pub struct AppState {
pub leptos_options: LeptosOptions,
pub websocket_state: Arc<WebSocketState>,
pub routes: Vec<RouteListing>,
}
pub async fn websocket_handler(
ws: WebSocketUpgrade,
State(state): State<AppState>,
) -> impl IntoResponse {
ws.on_upgrade(|socket| websocket(socket, Arc::new(state)))
}
async fn websocket(stream: WebSocket, state: Arc<AppState>) {
let state = &state.websocket_state;
let (mut sender, mut receiver) = stream.split();
let mut client_set = state.client_set.lock().await;
let uuid = uuid::Uuid::new_v4();
client_set.insert(uuid);
drop(client_set);
let mut rx = state.tx.subscribe();
let msg = format!("{uuid} joined");
println!("{uuid} joined");
let _ = state.tx.send(msg);
let mut send_task = tokio::spawn(async move {
while let Ok(msg) = rx.recv().await {
if sender.send(Message::Text(msg)).await.is_err() {
break;
}
}
});
let tx = state.tx.clone();
let uuid_clone = uuid.clone();
let mut recv_task = tokio::spawn(async move {
while let Some(Ok(Message::Text(text))) = receiver.next().await {
let _ = tx.send(format!("{uuid_clone}: {text}"));
}
});
tokio::select! {
_ = (&mut send_task) => recv_task.abort(),
_ = (&mut recv_task) => send_task.abort(),
};
let msg = format!("{uuid} left");
println!("{uuid} left");
let _ = state.tx.send(msg);
state.client_set.lock().await.remove(&uuid);
}