tensamin/omega
Watch
6
0
Fork
You've already forked omega
 0
Code Issues Pull requests 1 Projects Releases Packages Wiki Activity Actions

Merge branch 'main' of github.com:Tensamin/Omega

Browse source
This commit is contained in:
Alex Emmet 2026-02-11 22:45:16 +01:00
commit fcdd13bfc8
8 changed files with 676 additions and 185 deletions
Download patch file Download diff file Expand all files Collapse all files

52
src/server/socket.rs
View file

@ -12,9 +12,13 @@ use crate::server::omikron_connection::OmikronConnection;
pub fn handle(path: String, upgrades: OnUpgrade) {
tokio::spawn(async move {
log!(
"[ws] Spawning new task to handle WebSocket upgrade for path: {}",
path
);
match upgrades.await {
Ok(upgraded_stream) => {
log!("Valid WebSocket upgrade");
log!("[ws] WebSocket upgrade successful for path: {}", path);
let raw_stream = TokioIo::new(upgraded_stream);
let ws_stream = WebSocketStream::from_raw_socket(
@ -24,20 +28,28 @@ pub fn handle(path: String, upgrades: OnUpgrade) {
)
.await;
log!(
"WebSocket handshake successful, handling connection for {}",
"[ws] WebSocket handshake successful, handling connection for {}",
path
);
let (writer, reader) = ws_stream.split();
if path == "/ws/omikron" {
let connection = OmikronConnection::new(writer, reader);
start_connecteable_handler(connection).await;
tokio::spawn(start_connecteable_handler(connection));
}
}
Err(e) => {
log!("WebSocket upgrade failed after response: {:?}", e);
log!(
"[ERROR] WebSocket upgrade failed for path {}: {:?}",
path,
e
);
}
}
log!(
"[ws] WebSocket handling task for path: {} is finished.",
path
);
});
}
pub async fn start_connecteable_handler(connection: Arc<OmikronConnection>) {
@ -46,14 +58,13 @@ pub async fn start_connecteable_handler(connection: Arc<OmikronConnection>) {
const IDLE_TIMEOUT: Duration = Duration::from_secs(30);
log!("[ws_handler] Starting connection handler loop.");
loop {
let mut receiver = connection.receiver.write().await;
let mut receiver_guard = connection.receiver.write().await;
match tokio::time::timeout(IDLE_TIMEOUT, receiver.next()).await {
match tokio::time::timeout(IDLE_TIMEOUT, receiver_guard.next()).await {
Ok(Some(Ok(msg))) => {
// Drop the lock so other tasks can use the receiver if needed,
// and so we can handle the message without holding the lock.
drop(receiver);
drop(receiver_guard);
match msg {
Message::Text(text) => {
@ -63,28 +74,30 @@ pub async fn start_connecteable_handler(connection: Arc<OmikronConnection>) {
});
}
Message::Close(_) => {
log!("[ws_handler] Received 'Close' message. Breaking loop.");
break;
}
Message::Pong(_) => {
// Received a pong, connection is alive.
log!("[ws_handler] Received 'Pong'. Connection is alive.");
}
_ => {
log!("[ws_handler] Received unhandled message type.");
}
_ => {}
}
}
Ok(Some(Err(e))) => {
log!("WS Error: {}", e);
log!("[ERROR] WS Error: {}. Breaking loop.", e);
break;
}
Ok(None) => {
// Stream is closed
log!("[ws_handler] WebSocket stream closed by peer. Breaking loop.");
break;
}
Err(_) => {
// Timeout, we need to send a ping.
// Drop receiver lock before acquiring sender lock to avoid deadlock.
drop(receiver);
log!("WebSocket connection is idle. Sending a ping.");
drop(receiver_guard);
log!("[ws_handler] Timeout: Dropped receiver lock. Sending a ping.");
let mut sender = connection.sender.write().await;
log!("[ws_handler] Acquired sender lock for ping.");
if let Err(e) = sender
.send(Message::Text(Utf8Bytes::from(
CommunicationValue::new(CommunicationType::ping)
@ -93,11 +106,14 @@ pub async fn start_connecteable_handler(connection: Arc<OmikronConnection>) {
)))
.await
{
log!("Failed to send ping: {}. Closing connection.", e);
log!("[ERROR] Failed to send ping: {}. Closing connection.", e);
break;
}
log!("[ws_handler] Ping sent successfully.");
}
}
}
log!("[ws_handler] Connection handler loop finished.");
connection.handle_close().await;
log!("[ws_handler] Connection closed.");
}