diff --git a/Cargo.lock b/Cargo.lock index ecfd4e4..5cea467 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -616,7 +616,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -1581,7 +1581,7 @@ dependencies = [ [[package]] name = "mtp" version = "0.2.0" -source = "git+https://git.methanium.net/Methanium/mtp.git#fa271e62bec7918a8c994eea59ecdd62be76ead4" +source = "git+https://git.methanium.net/Methanium/mtp.git#bcf8aee3716f1690f1284748ac1ff3cd22799fb0" dependencies = [ "mtp-client", "mtp-codec", @@ -1597,7 +1597,7 @@ dependencies = [ [[package]] name = "mtp-client" version = "0.2.0" -source = "git+https://git.methanium.net/Methanium/mtp.git#fa271e62bec7918a8c994eea59ecdd62be76ead4" +source = "git+https://git.methanium.net/Methanium/mtp.git#bcf8aee3716f1690f1284748ac1ff3cd22799fb0" dependencies = [ "mtp-codec", "mtp-common", @@ -1610,7 +1610,7 @@ dependencies = [ [[package]] name = "mtp-codec" version = "0.2.0" -source = "git+https://git.methanium.net/Methanium/mtp.git#fa271e62bec7918a8c994eea59ecdd62be76ead4" +source = "git+https://git.methanium.net/Methanium/mtp.git#bcf8aee3716f1690f1284748ac1ff3cd22799fb0" dependencies = [ "base64 0.22.1", "byteorder", @@ -1623,7 +1623,7 @@ dependencies = [ [[package]] name = "mtp-common" version = "0.2.0" -source = "git+https://git.methanium.net/Methanium/mtp.git#fa271e62bec7918a8c994eea59ecdd62be76ead4" +source = "git+https://git.methanium.net/Methanium/mtp.git#bcf8aee3716f1690f1284748ac1ff3cd22799fb0" dependencies = [ "quinn", "rustls", @@ -1634,7 +1634,7 @@ dependencies = [ [[package]] name = "mtp-crypto" version = "0.2.0" -source = "git+https://git.methanium.net/Methanium/mtp.git#fa271e62bec7918a8c994eea59ecdd62be76ead4" +source = "git+https://git.methanium.net/Methanium/mtp.git#bcf8aee3716f1690f1284748ac1ff3cd22799fb0" dependencies = [ "base64 0.22.1", "chacha20poly1305", @@ -1656,7 +1656,7 @@ dependencies = [ [[package]] name = "mtp-files" version = "0.2.0" -source = "git+https://git.methanium.net/Methanium/mtp.git#fa271e62bec7918a8c994eea59ecdd62be76ead4" +source = "git+https://git.methanium.net/Methanium/mtp.git#bcf8aee3716f1690f1284748ac1ff3cd22799fb0" dependencies = [ "mtp-crypto", "rand 0.10.2", @@ -1667,7 +1667,7 @@ dependencies = [ [[package]] name = "mtp-host" version = "0.2.0" -source = "git+https://git.methanium.net/Methanium/mtp.git#fa271e62bec7918a8c994eea59ecdd62be76ead4" +source = "git+https://git.methanium.net/Methanium/mtp.git#bcf8aee3716f1690f1284748ac1ff3cd22799fb0" dependencies = [ "mtp-codec", "mtp-common", @@ -1682,7 +1682,7 @@ dependencies = [ [[package]] name = "mtp-transport" version = "0.2.0" -source = "git+https://git.methanium.net/Methanium/mtp.git#fa271e62bec7918a8c994eea59ecdd62be76ead4" +source = "git+https://git.methanium.net/Methanium/mtp.git#bcf8aee3716f1690f1284748ac1ff3cd22799fb0" dependencies = [ "async-trait", "mtp-codec", @@ -1700,7 +1700,7 @@ dependencies = [ [[package]] name = "mtp-type-map" version = "0.2.0" -source = "git+https://git.methanium.net/Methanium/mtp.git#fa271e62bec7918a8c994eea59ecdd62be76ead4" +source = "git+https://git.methanium.net/Methanium/mtp.git#bcf8aee3716f1690f1284748ac1ff3cd22799fb0" dependencies = [ "serde", "serde_yaml", @@ -1709,7 +1709,7 @@ dependencies = [ [[package]] name = "mtp-webserver" version = "0.2.0" -source = "git+https://git.methanium.net/Methanium/mtp.git#fa271e62bec7918a8c994eea59ecdd62be76ead4" +source = "git+https://git.methanium.net/Methanium/mtp.git#bcf8aee3716f1690f1284748ac1ff3cd22799fb0" dependencies = [ "async-trait", "bytes", @@ -2327,7 +2327,7 @@ dependencies = [ "once_cell", "socket2", "tracing", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -2581,7 +2581,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -2640,7 +2640,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -3035,7 +3035,7 @@ dependencies = [ "getrandom 0.3.4", "once_cell", "rustix", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -3596,7 +3596,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] diff --git a/mtp-type-maps b/mtp-type-maps index f753cfd..9776f43 160000 --- a/mtp-type-maps +++ b/mtp-type-maps @@ -1 +1 @@ -Subproject commit f753cfd406fc0050a8ccd76a90897205efceacd8 +Subproject commit 9776f430510d9aaeaa163f662ef5184108b6eec8 diff --git a/src/anonymous_clients/anonymous_client_connection.rs b/src/anonymous_clients/anonymous_client_connection.rs index a333260..da2c650 100644 --- a/src/anonymous_clients/anonymous_client_connection.rs +++ b/src/anonymous_clients/anonymous_client_connection.rs @@ -20,6 +20,7 @@ pub struct AnonymousClientConnection { pub sender: Arc, pub receiver: Arc, + pub ping: Arc>, pub interested_users: Arc>>, is_open: Arc>, pub user_name: Arc>, @@ -35,6 +36,7 @@ impl AnonymousClientConnection { state: general.state.clone(), user_id: user_id, + ping: Arc::new(RwLock::new(0)), interested_users: Arc::new(RwLock::new(Vec::new())), is_open: Arc::new(RwLock::new(true)), sender: general.sender.clone(), @@ -89,7 +91,9 @@ impl AnonymousClientConnection { ); return; } - log_cv_out!(PrintType::Client, &cv); + if !cv.is_type(CommunicationType::Pong) { + log_cv_out!(PrintType::Client, &cv); + } if let Err(e) = self.sender.send(&cv).await { log_out!( self.user_id as i64, @@ -107,6 +111,10 @@ impl AnonymousClientConnection { }; tokio::spawn(async move { let _permit = permit; + if cv.is_type(CommunicationType::Ping) { + self.handle_ping(cv).await; + return; + } log_cv_in!(PrintType::Client, &cv); if cv.is_type(CommunicationType::Identification) { @@ -222,6 +230,11 @@ impl AnonymousClientConnection { .await; } + // Handle ping + if cv.is_type(CommunicationType::Ping) { + self.handle_ping(cv).await; + return; + } // Handle client status changes if cv.is_type(CommunicationType::ClientChanged) { self.handle_client_changed(cv).await; @@ -321,6 +334,22 @@ impl AnonymousClientConnection { }); } + /// Handle ping message + async fn handle_ping(self: Arc, cv: CommunicationValue) { + // Update our ping if provided + if let DataValue::SignedNumber(last_ping) = cv.get_data(DataType::LastPing) { + if let Ok(ping_val) = last_ping.to_string().parse::() { + let mut ping_guard = self.ping.write().await; + *ping_guard = ping_val; + } + } + + // Send pong response + let response = CommunicationValue::new(CommunicationType::Pong).with_id(cv.get_id()); + + self.send_message(&response).await; + } + /// Handle client status change async fn handle_client_changed(self: Arc, cv: CommunicationValue) { if let DataValue::Str(status_str) = cv.get_data(DataType::UserState) { @@ -619,6 +648,7 @@ impl Clone for AnonymousClientConnection { sender: Arc::clone(&self.sender), receiver: Arc::clone(&self.receiver), user_id: self.user_id, + ping: Arc::clone(&self.ping), interested_users: Arc::clone(&self.interested_users), is_open: Arc::clone(&self.is_open), user_name: Arc::clone(&self.user_name), diff --git a/src/omega/omega_connection.rs b/src/omega/omega_connection.rs index a9c9dbb..a7806ea 100644 --- a/src/omega/omega_connection.rs +++ b/src/omega/omega_connection.rs @@ -11,7 +11,7 @@ use mtp::{ host::{Policy, SendMode}, }; use mtp_transport::ConnectionHandle; -use std::{sync::Arc, time::Duration}; +use std::{collections::HashMap, sync::Arc, time::Duration}; use tokio::{ sync::{Mutex, RwLock, mpsc, watch}, task::JoinHandle, @@ -23,8 +23,8 @@ use uuid::Uuid; const RECONNECT_DELAY: Duration = Duration::from_secs(5); const MAX_RECONNECT_DELAY: Duration = Duration::from_secs(300); const CONNECTION_TIMEOUT: Duration = Duration::from_secs(10); -const PING_INTERVAL: Duration = Duration::from_secs(5); -const MAX_MISSED_PINGS: usize = 3; +const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(5); +const MAX_HEARTBEAT_FAILURES: usize = 3; const TASK_CLEANUP_INTERVAL: Duration = Duration::from_secs(60); const TASK_MAX_AGE: Duration = Duration::from_secs(60); const MAX_CONCURRENT_REQUESTS: usize = 128; @@ -118,6 +118,9 @@ pub struct OmegaConnection { host: String, port: u16, omikron_id: u64, + last_ping: Arc>, + heartbeat_handle: Arc>>>, + message_send_times: Arc>>, pub connection_id: Uuid, shutdown_tx: Arc>>>, // Track if we should reconnect on close @@ -144,6 +147,9 @@ impl OmegaConnection { host: config.omega_host.clone(), port: config.omega_port, omikron_id: config.omikron_id, + last_ping: Arc::new(Mutex::new(-1)), + heartbeat_handle: Arc::new(Mutex::new(None)), + message_send_times: Arc::new(Mutex::new(HashMap::new())), connection_id: Uuid::new_v4(), shutdown_tx: Arc::new(Mutex::new(Some(shutdown_tx))), reconnect_on_close: Arc::new(RwLock::new(true)), @@ -198,6 +204,10 @@ impl OmegaConnection { handle.abort(); } + if let Some(handle) = self.heartbeat_handle.lock().await.take() { + handle.abort(); + } + // Close sender if connected if let Some(sender) = self.sender.read().await.as_ref() { sender.close().await; @@ -282,9 +292,7 @@ impl OmegaConnection { .with_receiver_queue_capacity(1000) .with_max_concurrent_stream_tasks(64) .with_persistent_stream_retries(5, Duration::from_secs(5)), - ) - .with_ping_interval(PING_INTERVAL) - .with_max_missed_pings(MAX_MISSED_PINGS); + ); let host_public_key = load_public_key_bundle("./omega.mpkb") .map_err(|e| format!("Failed to load omega.mpkb: {}", e))?; @@ -315,6 +323,13 @@ impl OmegaConnection { read_self.read_loop(connection, sender_handle).await; }); + // Start heartbeat + let heartbeat_self = self.clone(); + let heartbeat_handle = tokio::spawn(async move { + heartbeat_self.heartbeat_loop().await; + }); + *self.heartbeat_handle.lock().await = Some(heartbeat_handle); + // Tell omega our current state now that we're actually connected - // doing this after teardown (as before) sent into a sender that had // already been cleared, silently dropping the sync every time. @@ -342,6 +357,10 @@ impl OmegaConnection { *self.sender.write().await = None; *self.state.write().await = ConnectionState::Disconnected; + if let Some(handle) = self.heartbeat_handle.lock().await.take() { + handle.abort(); + } + match result { Ok(()) => { // Check if we should reconnect @@ -386,7 +405,7 @@ impl OmegaConnection { } // ------------------------------------------------------------------------- - // Read Loop + // Read Loop & Heartbeat // ------------------------------------------------------------------------- async fn read_loop( @@ -402,7 +421,14 @@ impl OmegaConnection { result = connection.receive() => { match result { Ok(cv) => { - log_cv_in!(PrintType::Omega, &cv); + if !cv.is_type(CommunicationType::Pong) && !cv.is_type(CommunicationType::Ping) { + log_cv_in!(PrintType::Omega, &cv); + } + + if cv.is_type(CommunicationType::Pong) || cv.is_type(CommunicationType::Ping) { + self.handle_pong(&cv).await; + continue; + } let msg_id = cv.get_id(); if let Some((_, task)) = self.waiting_tasks.remove(&msg_id) { @@ -451,12 +477,89 @@ impl OmegaConnection { } } + async fn heartbeat_loop(self: Arc) { + let mut consecutive_failures = 0; + loop { + sleep(HEARTBEAT_INTERVAL).await; + + // Check if still connected + if !self.state.read().await.is_connected() { + break; + } + + // Check if sender is closed + if let Some(sender) = self.sender.read().await.as_ref() { + if sender.is_closed() { + log_err!(0, PrintType::Omega, "Sender closed, stopping heartbeat"); + break; + } + } else { + break; + } + + match self.send_ping().await { + Ok(()) => consecutive_failures = 0, + Err(error) => { + consecutive_failures += 1; + log_err!( + 0, + PrintType::Omega, + "Heartbeat send failed ({}/{}): {}", + consecutive_failures, + MAX_HEARTBEAT_FAILURES, + error + ); + if consecutive_failures >= MAX_HEARTBEAT_FAILURES { + log_err!( + 0, + PrintType::Omega, + "Closing unhealthy Omega connection after repeated heartbeat failures" + ); + if let Some(sender) = self.sender.read().await.as_ref().cloned() { + sender.close().await; + } + break; + } + } + } + } + } + + async fn send_ping(&self) -> Result<(), String> { + let ping = CommunicationValue::new(CommunicationType::Ping).add_typed_default( + DataType::SendTime, + DataValue::SignedNumber( + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs() as i128, + ), + ); + self.try_send_message(&ping).await + } + + async fn handle_pong(&self, cv: &CommunicationValue) { + let timestamp = cv + .get_data(DataType::SendTime) + .as_number() + .unwrap_or_else(|| { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs() as i128 + }); + + *self.last_ping.lock().await = timestamp.try_into().unwrap(); + } + // ------------------------------------------------------------------------- // Public API // ------------------------------------------------------------------------- async fn try_send_message(&self, cv: &CommunicationValue) -> Result<(), String> { - log_cv_out!(PrintType::Omega, &cv); + if !cv.is_type(CommunicationType::Pong) && !cv.is_type(CommunicationType::Ping) { + log_cv_out!(PrintType::Omega, &cv); + } let sender_guard = self.sender.read().await; if let Some(sender) = sender_guard.as_ref() { diff --git a/src/omega/ping_pong_task.rs b/src/omega/ping_pong_task.rs new file mode 100644 index 0000000..a4e4687 --- /dev/null +++ b/src/omega/ping_pong_task.rs @@ -0,0 +1,41 @@ +use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue}; +use std::time::Duration; +use tokio::time::Instant; + +use crate::omega::omega_connection::OmegaConnection; + +const PING_TIMEOUT: Duration = Duration::from_secs(30); + +impl OmegaConnection { + pub async fn send_ping(&self) { + let id = rand::random::(); + let send_time = Instant::now(); + + let mut message_send_times = self.message_send_times.lock().await; + message_send_times.retain(|_uuid, time| time.elapsed() < PING_TIMEOUT); + message_send_times.insert(id as i64, send_time); + + self.send_ping_message(id).await; + } + + pub async fn send_ping_message(&self, id: u32) { + let ping_message = CommunicationValue::new(CommunicationType::Ping) + .with_id(id) + .add_data( + DataType::LastPing, + DataValue::SignedNumber(self.last_ping.lock().await.unwrap()), + ); + + self.send_message(&ping_message).await; + } + + /// Handles incoming pong and calculates latency + pub async fn handle_pong(&self, cv: &CommunicationValue, _log: bool) { + let id = cv.get_id(); + let mut message_send_times = self.message_send_times.lock().await; + if let Some(send_time) = message_send_times.remove(&(id as i64)) { + let ping = Instant::now().duration_since(send_time).as_millis() as i64; + *self.last_ping.lock().await = ping; + } + } +} diff --git a/src/rho/app_connection.rs b/src/rho/app_connection.rs index bb71c00..855104b 100644 --- a/src/rho/app_connection.rs +++ b/src/rho/app_connection.rs @@ -6,7 +6,7 @@ use crate::util::logger::PrintType; use crate::{log_cv_in, log_cv_out, log_err, log_in, log_out}; use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue}; use std::sync::Arc; -use std::time::Duration; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; use tokio::sync::RwLock; use uuid::Uuid; @@ -19,6 +19,7 @@ pub struct AppConnection { pub sender: Arc, pub receiver: Arc, + pub ping: Arc>, pub_key: Arc>>>, pub rho_connection: Arc>>>, is_open: Arc>, @@ -28,6 +29,7 @@ impl AppConnection { pub async fn from_general(general: Arc, user_id: u64) -> Arc { Arc::new(Self { state: general.state.clone(), + ping: Arc::new(RwLock::new(0)), pub_key: Arc::new(RwLock::new(None)), rho_connection: general.rho_connection.clone(), is_open: Arc::new(RwLock::new(true)), @@ -64,6 +66,12 @@ impl AppConnection { self.user_id } + /// Get current ping + #[allow(unused)] + pub async fn get_ping(&self) -> i64 { + *self.ping.read().await + } + /// Get RhoConnection if available pub async fn get_rho_connection(&self) -> Option> { self.rho_connection.read().await.clone() @@ -79,13 +87,19 @@ impl AppConnection { ); return; } - log_cv_out!(PrintType::App, &cv); + if !cv.is_type(CommunicationType::Pong) && !cv.is_type(CommunicationType::Ping) { + log_cv_out!(PrintType::App, &cv); + } let _ = self.sender.send(&cv).await; } /// Handle incoming message from app pub async fn handle_message(self: Arc, cv: CommunicationValue) { tokio::spawn(async move { + if cv.is_type(CommunicationType::Ping) { + self.handle_ping(cv).await; + return; + } log_cv_in!(PrintType::App, cv); if cv.is_type(CommunicationType::GetUserData) { @@ -154,6 +168,36 @@ impl AppConnection { }); } + /// Handle ping message + async fn handle_ping(self: Arc, cv: CommunicationValue) { + // Update our ping if provided + if let DataValue::SignedNumber(last_ping) = cv.get_data(DataType::LastPing) { + let current = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis(); + let mut ping_guard = self.ping.write().await; + *ping_guard = (current as i128 - *last_ping) as i64; + } + + // Get Iota ping from RhoConnection + let iota_ping = if let Some(rho_conn) = self.get_rho_connection().await { + rho_conn.get_iota_connection().get_ping().await + } else { + -1 + }; + + // Send pong response + let response = CommunicationValue::new(CommunicationType::Pong) + .with_id(cv.get_id()) + .add_typed_default( + DataType::PingIota, + DataValue::SignedNumber(iota_ping.into()), + ); + + self.send_message(&response).await; + } + /// Forward message to Iota async fn forward_to_iota(self: Arc, cv: CommunicationValue) { let sender_user_id = self.get_user_id().await; @@ -309,6 +353,7 @@ impl Clone for AppConnection { app_identifier: self.app_identifier.clone(), app_session: self.app_session, client_version: self.client_version.clone(), + ping: Arc::clone(&self.ping), pub_key: Arc::clone(&self.pub_key), rho_connection: Arc::clone(&self.rho_connection), is_open: Arc::clone(&self.is_open), diff --git a/src/rho/client_connection.rs b/src/rho/client_connection.rs index 0a618d8..438abd7 100644 --- a/src/rho/client_connection.rs +++ b/src/rho/client_connection.rs @@ -9,7 +9,7 @@ use crate::{log_cv_in, log_cv_out, log_err, log_in, log_out}; use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue}; use std::str::FromStr; use std::sync::Arc; -use std::time::Duration; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; use tokio::sync::RwLock; use trust_dns_resolver::TokioAsyncResolver; use uuid::Uuid; @@ -22,6 +22,7 @@ pub struct ClientConnection { pub sender: Arc, pub receiver: Arc, + pub ping: Arc>, pub_key: Arc>>>, pub rho_connection: Arc>>>, pub interested_users: Arc>>, @@ -33,6 +34,7 @@ impl ClientConnection { pub async fn from_general(general: Arc, user_id: u64) -> Arc { Arc::new(Self { state: general.state.clone(), + ping: Arc::new(RwLock::new(0)), pub_key: Arc::new(RwLock::new(None)), rho_connection: general.rho_connection.clone(), interested_users: Arc::new(RwLock::new(Vec::new())), @@ -70,6 +72,11 @@ impl ClientConnection { self.user_id } + /// Get current ping + pub async fn get_ping(&self) -> i64 { + *self.ping.read().await + } + /// Get RhoConnection if available pub async fn get_rho_connection(&self) -> Option> { self.rho_connection.read().await.clone() @@ -85,7 +92,9 @@ impl ClientConnection { ); return; } - log_cv_out!(PrintType::Client, &cv); + if !cv.is_type(CommunicationType::Pong) && !cv.is_type(CommunicationType::Ping) { + log_cv_out!(PrintType::Client, &cv); + } let _ = self.sender.send(&cv).await; } @@ -96,6 +105,10 @@ impl ClientConnection { }; tokio::spawn(async move { let _permit = permit; + if cv.is_type(CommunicationType::ClientPing) { + self.handle_ping(cv).await; + return; + } log_cv_in!(PrintType::Client, cv); let mut cv = cv; @@ -298,6 +311,36 @@ impl ClientConnection { } } + /// Handle ping message + async fn handle_ping(self: Arc, cv: CommunicationValue) { + // Update our ping if provided + if let DataValue::SignedNumber(last_ping) = cv.get_data(DataType::LastPing) { + let current = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis(); + let mut ping_guard = self.ping.write().await; + *ping_guard = (current as i128 - *last_ping) as i64; + } + + // Get Iota ping from RhoConnection + let iota_ping = if let Some(rho_conn) = self.get_rho_connection().await { + rho_conn.get_iota_connection().get_ping().await + } else { + -1 + }; + + // Send pong response + let response = CommunicationValue::new(CommunicationType::ClientPing) + .with_id(cv.get_id()) + .add_typed_default( + DataType::PingIota, + DataValue::SignedNumber(iota_ping.into()), + ); + + self.send_message(&response).await; + } + /// Handle client status change async fn handle_client_changed(self: Arc, cv: CommunicationValue) { let user_id = self.get_user_id().await; @@ -874,6 +917,7 @@ impl Clone for ClientConnection { user_id: self.user_id, session_id: self.session_id, client_version: self.client_version.clone(), + ping: Arc::clone(&self.ping), pub_key: Arc::clone(&self.pub_key), rho_connection: Arc::clone(&self.rho_connection), interested_users: Arc::clone(&self.interested_users), diff --git a/src/rho/iota_connection.rs b/src/rho/iota_connection.rs index 7d1b78e..7f6fe4c 100755 --- a/src/rho/iota_connection.rs +++ b/src/rho/iota_connection.rs @@ -31,6 +31,7 @@ pub struct IotaConnection { pub sender: Arc, pub receiver: Arc, pub user_ids: Arc>>, + pub ping: Arc>, pub_key: Arc>>>, pub waiting_tasks: DashMap, CommunicationValue) -> bool + Send + Sync>>, @@ -43,6 +44,7 @@ impl IotaConnection { pub async fn from_general(general: Arc, iota_id: u64) -> Arc { Arc::new(Self { state: general.state.clone(), + ping: Arc::new(RwLock::new(0)), pub_key: Arc::new(RwLock::new(None)), rho_connection: general.rho_connection.clone(), user_ids: Arc::new(RwLock::new(Vec::new())), @@ -149,6 +151,11 @@ impl IotaConnection { .push(cv); } + /// Get current ping + pub async fn get_ping(&self) -> i64 { + *self.ping.read().await + } + /// Set the RhoConnection reference pub async fn set_rho_connection(&self, rho_connection: Arc) { let mut rho_ref = self.rho_connection.write().await; @@ -167,7 +174,9 @@ impl IotaConnection { /// Send a CommunicationValue to the Iota pub async fn send_message(&self, cv: &CommunicationValue) { - log_cv_out!(PrintType::Iota, cv); + if !cv.is_type(CommunicationType::Pong) { + log_cv_out!(PrintType::Iota, cv); + } if let Err(e) = self.sender.send(&cv).await { log_err!( self.iota_id as i64, @@ -191,6 +200,12 @@ impl IotaConnection { } } + // Handle ping + if cv.is_type(CommunicationType::Ping) || cv.is_type(CommunicationType::Pong) { + self.handle_ping(cv).await; + return; + } + log_cv_in!(PrintType::Iota, cv); let cv = if cv.is_type(CommunicationType::ClientStateSync) { @@ -437,6 +452,46 @@ impl IotaConnection { } } } + /// Handle ping message + async fn handle_ping(&self, cv: CommunicationValue) { + if let DataValue::SignedNumber(last_ping) = cv.get_data(DataType::LastPing) { + if let Ok(ping_val) = last_ping.to_string().parse::() { + let mut ping_guard = self.ping.write().await; + *ping_guard = ping_val; + } + } + + let client_pings = if let Some(rho_conn) = self.get_rho_connection().await { + rho_conn.get_client_pings().await + } else { + HashMap::new() + }; + + let tm = TypeMap::latest(); + let pings: Vec = client_pings + .into_iter() + .map(|(k, v)| { + let mut map = BTreeMap::new(); + if let Ok(uid) = k.parse::() { + map.insert( + data_type_id(DataType::UserId, &tm), + DataValue::SignedNumber(uid), + ); + } + map.insert( + data_type_id(DataType::LastPing, &tm), + DataValue::SignedNumber(v.into()), + ); + DataValue::container_from_map(&map) + }) + .collect(); + let response = CommunicationValue::new(CommunicationType::Pong) + .with_id(cv.get_id()) + .add_typed_default(DataType::PingClients, DataValue::Array(pings)); + + self.send_message(&response).await; + } + /// Handle message forwarding to other Iotas async fn handle_forward_message(&self, cv: CommunicationValue) { let receiver_id = cv.get_receiver(); diff --git a/src/rho/rho_connection.rs b/src/rho/rho_connection.rs index 35e259d..c7e4f65 100644 --- a/src/rho/rho_connection.rs +++ b/src/rho/rho_connection.rs @@ -3,6 +3,7 @@ use super::{client_connection::ClientConnection, iota_connection::IotaConnection use crate::{data::user::UserStatus, rho::app_connection::AppConnection}; use dashmap::DashMap; use mtp::codec::{CommunicationValue, DataType}; +use std::collections::HashMap; use std::sync::Arc; use tokio::sync::RwLock; use uuid::Uuid; @@ -262,6 +263,19 @@ impl RhoConnection { } } + /// Get ping information for all clients + pub async fn get_client_pings(&self) -> HashMap { + let connections = self.get_client_connections().await; + let mut pings = HashMap::new(); + + for connection in connections.iter() { + let user_id = connection.get_user_id().await; + pings.insert(user_id.to_string(), connection.get_ping().await); + } + + pings + } + /// Check if this RhoConnection contains a specific user ID #[allow(dead_code)] pub async fn contains_user(&self, user_id: &i64) -> bool {