Added websocket server
This commit is contained in:
parent
b9d5425ac5
commit
683233a4b5
@ -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"]
|
||||
|
@ -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);
|
||||
|
@ -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;
|
||||
}
|
||||
}
|
||||
|
@ -1 +1,2 @@
|
||||
pub mod surrealdb;
|
||||
pub mod websocket;
|
||||
|
2
application/src/util/websocket.rs
Normal file
2
application/src/util/websocket.rs
Normal file
@ -0,0 +1,2 @@
|
||||
#[cfg(feature = "ssr")]
|
||||
pub mod server;
|
80
application/src/util/websocket/server.rs
Normal file
80
application/src/util/websocket/server.rs
Normal 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);
|
||||
}
|
Loading…
Reference in New Issue
Block a user