[Fix] More stable Connection

src/anonymous_clients/anonymous_client_connection.rs

@@ -13,7 +13,7 @@ use uuid::Uuid;
 use crate::anonymous_clients::anonymous_manager::{self, generate_username};
 use crate::calls::call_manager;
 use crate::data::communication::{CommunicationType, CommunicationValue, DataTypes};
-use crate::omega::omega_connection::{WAITING_TASKS, get_omega_connection};
+use crate::omega::omega_connection::get_omega_connection;
 use crate::rho::rho_manager;
 use crate::util::logger::PrintType;
 use crate::{log_in, log_out};
@@ -328,19 +328,17 @@ impl AnonymousClientConnection {
     }
     async fn handle_omega_forward(self: Arc<Self>, cv: CommunicationValue) {
         let client_for_closure = self.clone();
-        WAITING_TASKS.insert(
-            cv.get_id(),
-            Box::new(move |_, response_cv| {
-                let client = client_for_closure.clone();
         tokio::spawn(async move {
-                    client.send_message(&response_cv).await;
+            let response_cv = get_omega_connection()
-                });
+                .await_response(
-                true
+                    &cv.with_sender(*self.user_id.read().await),
-            }),
+                    Some(Duration::from_secs(20)),
-        );
+                )
-        get_omega_connection()
-            .send_message(&cv.with_sender(*self.user_id.read().await))
                 .await;
+            if let Ok(response_cv) = response_cv {
+                client_for_closure.send_message(&response_cv).await;
+            }
+        });
     }
 
     /// Handle ping message

src/main.rs

@@ -19,7 +19,6 @@ use crate::{
         anonymous_client_connection::AnonymousClientConnection, anonymous_manager,
     },
     calls::call_util::garbage_collect_calls,
-    omega::omega_connection::OmegaConnection,
     rho::{client_connection::ClientConnection, iota_connection::IotaConnection},
     util::{
         crypto_helper::{load_public_key, load_secret_key},
@@ -39,9 +38,6 @@ pub fn get_public_key() -> x448::PublicKey {
 #[tokio::main]
 async fn main() {
     dotenv().ok();
-    tokio::spawn(async move {
-        Arc::new(OmegaConnection::new()).connect();
-    });
     startup();
     let address = format!(
         "{}:{}",

src/omega/omega_connection.rs

@@ -31,9 +31,12 @@ use crate::{
     },
 };
 
-pub static WAITING_TASKS: Lazy<
+pub struct WaitingTask {
-    DashMap<Uuid, Box<dyn Fn(Arc<OmegaConnection>, CommunicationValue) -> bool + Send + Sync>>,
+    pub task: Box<dyn Fn(Arc<OmegaConnection>, CommunicationValue) -> bool + Send + Sync>,
-> = Lazy::new(DashMap::new);
+    pub inserted_at: Instant,
+}
+
+pub static WAITING_TASKS: Lazy<DashMap<Uuid, WaitingTask>> = Lazy::new(DashMap::new);
 
 static OMEGA_CONNECTION: Lazy<Arc<OmegaConnection>> = Lazy::new(|| {
     let conn = Arc::new(OmegaConnection::new());
@@ -41,6 +44,14 @@ static OMEGA_CONNECTION: Lazy<Arc<OmegaConnection>> = Lazy::new(|| {
     tokio::spawn(async move {
         conn_clone.connect_internal(0).await;
     });
+
+    tokio::spawn(async {
+        loop {
+            sleep(Duration::from_secs(60)).await;
+            WAITING_TASKS.retain(|_, v| v.inserted_at.elapsed() < Duration::from_secs(60));
+        }
+    });
+
     conn
 });
 
@@ -52,7 +63,7 @@ pub fn get_omega_connection() -> Arc<OmegaConnection> {
 #[derive(Clone)]
 pub struct OmegaConnection {
     write: Arc<
-        RwLock<
+        Mutex<
             Option<
                 WebSocketSender<
                     Stream<TokioAdapter<TcpStream>, TokioAdapter<TlsStream<TcpStream>>>,
@@ -61,7 +72,7 @@ pub struct OmegaConnection {
         >,
     >,
     read: Arc<
-        RwLock<
+        Mutex<
             Option<
                 WebSocketReceiver<
                     Stream<TokioAdapter<TcpStream>, TokioAdapter<TlsStream<TcpStream>>>,
@@ -72,27 +83,48 @@ pub struct OmegaConnection {
     pingpong: Arc<Mutex<Option<tokio::task::JoinHandle<()>>>>,
     pub last_ping: Arc<Mutex<i64>>,
     pub message_send_times: Arc<Mutex<HashMap<Uuid, Instant>>>,
-    pub is_connected: Arc<RwLock<bool>>,
+    state: Arc<RwLock<ConnectionState>>,
+}
+
+#[derive(Clone, PartialEq, Eq)]
+enum ConnectionState {
+    Disconnected,
+    Connecting,
+    Connected,
 }
 
 impl OmegaConnection {
     pub fn new() -> Self {
         OmegaConnection {
-            read: Arc::new(RwLock::new(None)),
+            read: Arc::new(Mutex::new(None)),
-            write: Arc::new(RwLock::new(None)),
+            write: Arc::new(Mutex::new(None)),
             pingpong: Arc::new(Mutex::new(None)),
             last_ping: Arc::new(Mutex::new(-1)),
             message_send_times: Arc::new(Mutex::new(HashMap::new())),
-            is_connected: Arc::new(RwLock::new(false)),
+            state: Arc::new(RwLock::new(ConnectionState::Disconnected)),
         }
     }
-    pub fn connect(self: Arc<OmegaConnection>) {
-        let cloned_self = self.clone();
-        tokio::spawn(async move {
-            cloned_self.connect_internal(0).await;
-        });
-    }
     async fn connect_internal(self: Arc<OmegaConnection>, mut retry: usize) {
+        if self.state.read().await.eq(&ConnectionState::Connected) {
+            return;
+        }
+        if self.state.read().await.eq(&ConnectionState::Connecting) {
+            let start = Instant::now();
+            let timeout = Duration::from_secs(10);
+            while self.state.read().await.eq(&ConnectionState::Connecting)
+                && start.elapsed() < timeout
+            {
+                sleep(Duration::from_millis(100)).await;
+            }
+            return;
+        }
+
+        *self.state.write().await = ConnectionState::Connecting;
+
+        if let Some(handle) = self.pingpong.lock().await.take() {
+            handle.abort();
+        }
+
         loop {
             if retry > 5 {
                 log_err!(
@@ -100,6 +132,7 @@ impl OmegaConnection {
                     PrintType::Omega,
                     "Max retry attempts reached, giving up."
                 );
+                *self.state.write().await = ConnectionState::Disconnected;
                 return;
             }
 
@@ -107,16 +140,17 @@ impl OmegaConnection {
                 env::var("OMEGA_HOST").unwrap_or("wss://omega.tensamin.net/ws/omikron".to_string());
             match connect_async(&url_str).await {
                 Ok((ws_stream, _)) => {
-                    *self.is_connected.write().await = true;
                     log_in!(0, PrintType::Omega, "WebSocket connected to {}", url_str);
                     retry = 0;
                     let (write, read) = ws_stream.split();
-                    *self.read.write().await = Some(read);
+                    *self.read.lock().await = Some(read);
-                    *self.write.write().await = Some(write);
+                    *self.write.lock().await = Some(write);
 
-                    let cloned_self = self.clone();
+                    *self.state.write().await = ConnectionState::Connected;
-                    tokio::spawn(async move {
+
-                        cloned_self.clone().read_loop().await;
+                    let read_loop_self = self.clone();
+                    let read_loop_handle = tokio::spawn(async move {
+                        read_loop_self.read_loop().await;
                     });
 
                     let cloned_self = self.clone();
@@ -137,7 +171,8 @@ impl OmegaConnection {
                                 );
                         WAITING_TASKS.insert(
                             id,
-                            Box::new(|selfc, cv| {
+                            WaitingTask {
+                                task: Box::new(|selfc, cv| {
                                     if cv.is_type(CommunicationType::error_not_found) {
                                         log_err!(0,
                                             PrintType::Omega,
@@ -186,7 +221,8 @@ impl OmegaConnection {
                                             let response_id = response_msg.get_id();
                                             WAITING_TASKS.insert(
                                                 response_id,
-                                            Box::new(|selfc, final_cv| {
+                                                WaitingTask {
+                                                    task: Box::new(|selfc, final_cv| {
                                                         if !final_cv
                                                             .is_type(CommunicationType::identification_response)
                                                         {
@@ -237,6 +273,8 @@ impl OmegaConnection {
                                                         );
                                                         true
                                                     }),
+                                                    inserted_at: Instant::now(),
+                                                }
                                             );
 
                                             selfc.send_message(&response_msg).await;
@@ -251,34 +289,39 @@ impl OmegaConnection {
 
                                     true
                                 }),
+                                inserted_at: Instant::now(),
+                            }
                         );
                         cloned_self.send_message(&identify_msg).await
                     });
 
-                    let cloned_self = self.clone();
+                    let ping_pong_self = self.clone();
                     let handle = tokio::spawn(async move {
                         loop {
-                            if *cloned_self.is_connected.read().await == false {
+                            if *ping_pong_self.state.read().await != ConnectionState::Connected {
                                 break;
                             }
-                            cloned_self.send_ping().await;
+                            ping_pong_self.send_ping().await;
                             sleep(Duration::from_secs(5)).await;
                         }
                     });
 
-                    *self.is_connected.write().await = true;
                     *self.pingpong.lock().await = Some(handle);
 
-                    while *self.is_connected.read().await {
+                    read_loop_handle.await.unwrap_or_else(|e| {
-                        sleep(Duration::from_secs(2)).await;
+                        log_err!(0, PrintType::Omega, "Read loop task failed: {}", e)
-                    }
+                    });
-                    *self.read.write().await = None;
+
-                    *self.write.write().await = None;
+                    *self.read.lock().await = None;
+                    *self.write.lock().await = None;
+                    *self.state.write().await = ConnectionState::Disconnected;
+
                     log_err!(0, PrintType::Omega, "Connection lost. Retrying...");
                     retry += 1;
                     sleep(Duration::from_secs(2)).await;
                 }
                 Err(e) => {
+                    *self.state.write().await = ConnectionState::Disconnected;
                     log_err!(
                         0,
                         PrintType::Omega,
@@ -295,15 +338,13 @@ impl OmegaConnection {
     }
 
     async fn read_loop(self: Arc<Self>) {
+        let mut reader = match self.read.lock().await.take() {
+            Some(reader) => reader,
+            None => return,
+        };
+
         loop {
-            let msg = {
+            let msg = reader.next().await;
-                let mut guard = self.read.write().await;
-                let ws = match guard.as_mut() {
-                    Some(ws) => ws,
-                    _ => break,
-                };
-                ws.next().await
-            };
 
             match msg {
                 Some(Ok(Message::Text(msg))) => {
@@ -314,9 +355,8 @@ impl OmegaConnection {
                     }
                     let msg_id = cv.get_id();
                     log_in!(0, PrintType::Omega, "{}", &cv.to_json().to_string());
-                    // Handle waiting tasks
                     if let Some(task) = WAITING_TASKS.remove(&msg_id) {
-                        if (task.1)(self.clone(), cv.clone()) {
+                        if (task.1.task)(self.clone(), cv.clone()) {
                             continue;
                         }
                     } else {
@@ -328,21 +368,17 @@ impl OmegaConnection {
                 _ => {}
             }
         }
-
-        *self.is_connected.write().await = false;
-        *self.read.write().await = None;
-        *self.write.write().await = None;
     }
 
     pub async fn send_message(&self, cv: &CommunicationValue) {
-        let mut guard = self.write.write().await;
-        if let Some(ws) = guard.as_mut() {
         if !cv.is_type(CommunicationType::ping) {
             log_out!(0, PrintType::Omega, "{}", &cv.to_json().to_string());
         }
-            let _ = ws
+        let msg = cv.to_json().to_string();
-                .send(Message::Text(cv.to_json().to_string().into()))
+
-                .await;
+        let mut guard = self.write.lock().await;
+        if let Some(ws) = guard.as_mut() {
+            let _ = ws.send(Message::Text(msg.into())).await;
         }
     }
 
@@ -378,20 +414,22 @@ impl OmegaConnection {
 
         WAITING_TASKS.insert(
             msg_id,
-            Box::new(
+            WaitingTask {
+                task: Box::new(
                     move |_: Arc<OmegaConnection>, response: CommunicationValue| {
-                    let _ = Box::pin(async move |_: CommunicationValue| {
+                        tokio::spawn(async move {
                             let rho = rho_manager::get_rho_con_for_user(user_id).await;
                             if let Some(rho) = rho {
                                 for client in rho.get_client_connections_for_user(user_id).await {
                                     client.send_message(&response).await;
                                 }
                             }
-                        true
                         });
                         true
                     },
                 ),
+                inserted_at: Instant::now(),
+            },
         );
 
         OmegaConnection::send_global(cv).await;
@@ -402,18 +440,16 @@ impl OmegaConnection {
     }
 
     pub async fn await_connection(&self, timeout_duration: Option<Duration>) -> Result<(), String> {
-        if *self.is_connected.read().await {
+        if *self.state.read().await == ConnectionState::Connected {
             return Ok(());
         }
         let timeout = timeout_duration.unwrap_or(Duration::from_secs(10));
 
         let start = Instant::now();
         loop {
-            {
+            if *self.state.read().await == ConnectionState::Connected {
-                if *self.is_connected.read().await {
                 return Ok(());
             }
-            }
 
             if start.elapsed() >= timeout {
                 return Err(format!(
@@ -431,27 +467,29 @@ impl OmegaConnection {
         cv: &CommunicationValue,
         timeout_duration: Option<Duration>,
     ) -> Result<CommunicationValue, String> {
-        let _ = self.await_connection(timeout_duration).await;
+        self.await_connection(timeout_duration).await?;
         let (tx, mut rx) = mpsc::channel(1);
         let msg_id = cv.get_id();
 
         let task_tx = tx.clone();
         WAITING_TASKS.insert(
             msg_id,
-            Box::new(move |_, response_cv| {
+            WaitingTask {
+                task: Box::new(move |_, response_cv| {
                     let inner_tx = task_tx.clone();
                     tokio::spawn(async move {
-                    if let Err(e) = inner_tx.send(response_cv).await {
+                        if inner_tx.send(response_cv).await.is_err() {
                             log_err!(
                                 0,
                                 PrintType::Omega,
-                            "Failed to send response back to awaiter: {}",
+                                "Failed to send response back to awaiter",
-                            e
                             );
                         }
                     });
                     true
                 }),
+                inserted_at: Instant::now(),
+            },
         );
 
         self.send_message(cv).await;

src/omega/ping_pong_task.rs

@@ -1,15 +1,21 @@
 use crate::data::communication::{CommunicationType, CommunicationValue, DataTypes};
 use crate::omega::omega_connection::OmegaConnection;
 use json::number::Number;
+use std::time::Duration;
 use tokio::time::Instant;
 use uuid::Uuid;
 
+const PING_TIMEOUT: Duration = Duration::from_secs(30);
+
 impl OmegaConnection {
     pub async fn send_ping(&self) {
         let uuid = Uuid::new_v4();
         let send_time = Instant::now();
 
-        self.message_send_times.lock().await.insert(uuid, send_time);
+        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(uuid, send_time);
+
         self.send_ping_message(uuid).await;
     }
 
@@ -27,16 +33,10 @@ impl OmegaConnection {
     /// Handles incoming pong and calculates latency
     pub async fn handle_pong(&self, cv: &CommunicationValue, _log: bool) {
         let id = cv.get_id();
-        let send_time_opt = {
+        let mut message_send_times = self.message_send_times.lock().await;
-            let queue = self.message_send_times.lock().await;
+        if let Some(send_time) = message_send_times.remove(&id) {
-            queue.get(&id).cloned()
-        };
-
-        if let Some(send_time) = send_time_opt {
             let ping = Instant::now().duration_since(send_time).as_millis() as i64;
-            self.message_send_times.lock().await.remove(&id);
+            *self.last_ping.lock().await = ping;
-
-            *self.last_ping.lock().await = ping as i64;
         }
     }
 }

src/rho/client_connection.rs

@@ -15,7 +15,7 @@ use uuid::Uuid;
 use super::{rho_connection::RhoConnection, rho_manager};
 use crate::anonymous_clients::anonymous_manager;
 use crate::calls::{call_manager, call_util};
-use crate::omega::omega_connection::{WAITING_TASKS, get_omega_connection};
+use crate::omega::omega_connection::get_omega_connection;
 use crate::util::crypto_helper::{load_public_key, public_key_to_base64};
 use crate::util::crypto_util::{DataFormat, SecurePayload};
 use crate::util::logger::PrintType;
@@ -374,19 +374,17 @@ impl ClientConnection {
     }
     async fn handle_omega_forward(self: Arc<Self>, cv: CommunicationValue) {
         let client_for_closure = self.clone();
-        WAITING_TASKS.insert(
-            cv.get_id(),
-            Box::new(move |_, response_cv| {
-                let client = client_for_closure.clone();
         tokio::spawn(async move {
-                    client.send_message(&response_cv).await;
+            let response_cv = get_omega_connection()
-                });
+                .await_response(
-                true
+                    &cv.with_sender(*self.user_id.read().await),
-            }),
+                    Some(Duration::from_secs(20)),
-        );
+                )
-        get_omega_connection()
-            .send_message(&cv.with_sender(*self.user_id.read().await))
                 .await;
+            if let Ok(response_cv) = response_cv {
+                client_for_closure.send_message(&response_cv).await;
+            }
+        });
     }
 
     /// Handle ping message

src/rho/iota_connection.rs

@@ -5,7 +5,6 @@ use crate::get_public_key;
 use crate::log_err;
 use crate::log_in;
 use crate::log_out;
-use crate::omega::omega_connection::WAITING_TASKS;
 use crate::omega::omega_connection::get_omega_connection;
 use crate::util::crypto_helper::load_public_key;
 use crate::util::crypto_helper::public_key_to_base64;
@@ -460,19 +459,17 @@ impl IotaConnection {
 
     async fn handle_omega_forward(self: Arc<Self>, cv: CommunicationValue) {
         let iota_for_closure = self.clone();
-        WAITING_TASKS.insert(
-            cv.get_id(),
-            Box::new(move |_, response_cv| {
-                let iota = iota_for_closure.clone();
         tokio::spawn(async move {
-                    iota.send_message(&response_cv).await;
+            let response_cv = get_omega_connection()
-                });
+                .await_response(
-                true
+                    &cv.with_sender(*self.iota_id.read().await),
-            }),
+                    Some(Duration::from_secs(20)),
-        );
+                )
-        get_omega_connection()
-            .send_message(&cv.with_sender(*self.iota_id.read().await))
                 .await;
+            if let Ok(response_cv) = response_cv {
+                iota_for_closure.send_message(&response_cv).await;
+            }
+        });
     }
     /// Handle ping message
     async fn handle_ping(&self, cv: CommunicationValue) {