43 lines
1.3 KiB
Rust
43 lines
1.3 KiB
Rust
pub mod args;
|
|
pub mod connection;
|
|
pub mod request;
|
|
pub mod response;
|
|
|
|
use std::{collections::HashMap, time::Duration};
|
|
|
|
use axum::extract::ws::Message;
|
|
use once_cell::sync::Lazy;
|
|
use surrealdb::channel::Sender;
|
|
use tokio::sync::RwLock;
|
|
use tokio_util::sync::CancellationToken;
|
|
use uuid::Uuid;
|
|
|
|
static CONN_CLOSED_ERR: &str = "Connection closed normally";
|
|
|
|
pub struct WebSocketRef(Sender<Message>, CancellationToken);
|
|
// Mapping of WebSocketID to WebSocket
|
|
type WebSockets = RwLock<HashMap<Uuid, WebSocketRef>>;
|
|
// Mapping of LiveQueryID to WebSocketID
|
|
type LiveQueries = RwLock<HashMap<Uuid, Uuid>>;
|
|
|
|
pub(crate) static WEBSOCKETS: Lazy<WebSockets> = Lazy::new(WebSockets::default);
|
|
pub(crate) static LIVE_QUERIES: Lazy<LiveQueries> = Lazy::new(LiveQueries::default);
|
|
|
|
pub(crate) async fn graceful_shutdown() {
|
|
// Close all WebSocket connections. Queued messages will still be processed.
|
|
for (_, WebSocketRef(_, cancel_token)) in WEBSOCKETS.read().await.iter() {
|
|
cancel_token.cancel();
|
|
}
|
|
|
|
// Wait for all existing WebSocket connections to gracefully close
|
|
while WEBSOCKETS.read().await.len() > 0 {
|
|
tokio::time::sleep(Duration::from_millis(100)).await;
|
|
}
|
|
}
|
|
|
|
pub(crate) fn shutdown() {
|
|
// Close all WebSocket connections immediately
|
|
if let Ok(mut writer) = WEBSOCKETS.try_write() {
|
|
writer.drain();
|
|
}
|
|
}
|