[Fix] (stable) Omikron - Omega connection,

src/server/omikron_connection.rs

@@ -958,25 +958,21 @@ impl OmikronConnection {
         self.send_message(&error).await;
     }
     pub async fn close(&self) {
-        if self.is_identified().await {
         let omikron_id = self.get_omikron_id().await;
         if omikron_id != 0 {
-                log_in!(PrintType::Omega, "Omikron Disconnected");
+            log_in!(omikron_id, PrintType::Omega, "Omikron Disconnected");
             omikron_manager::remove_omikron(omikron_id).await;
             user_online_tracker::untrack_omikron(omikron_id).await;
         }
     }
-    }
     pub async fn handle_close(self: Arc<Self>) {
-        if self.is_identified().await {
         let omikron_id = self.get_omikron_id().await;
         if omikron_id != 0 {
-                log_in!(PrintType::Omega, "Omikron Disconnected");
+            log_in!(omikron_id, PrintType::Omega, "Omikron Disconnected");
             omikron_manager::remove_omikron(omikron_id).await;
             user_online_tracker::untrack_omikron(omikron_id).await;
         }
     }
-    }
 
     async fn handle_ping(&self, cv: CommunicationValue) {
         if let Some(last_ping) = cv.get_data(DataTypes::last_ping) {

src/server/omikron_manager.rs

@@ -1,4 +1,4 @@
-use crate::{log_in, server::omikron_connection::OmikronConnection, util::logger::PrintType};
+use crate::server::omikron_connection::OmikronConnection;
 use dashmap::DashMap;
 use once_cell::sync::Lazy;
 use rand::prelude::IteratorRandom;
@@ -7,20 +7,27 @@ use std::sync::Arc;
 pub static OMIKRON_CONNECTIONS: Lazy<DashMap<i64, Arc<OmikronConnection>>> =
     Lazy::new(|| DashMap::new());
 
-pub async fn add_omikron(omikron_conn: Arc<OmikronConnection>) {
+pub async fn add_omikron(conn: Arc<OmikronConnection>) {
-    OMIKRON_CONNECTIONS.insert(omikron_conn.get_omikron_id().await, omikron_conn);
+    let id = conn.get_omikron_id().await;
+
+    if let Some(old) = OMIKRON_CONNECTIONS.insert(id, conn.clone()) {
+        old.close().await;
+    }
 }
 
 pub async fn remove_omikron(omikron_id: i64) {
     OMIKRON_CONNECTIONS.remove(&omikron_id);
 }
-
 pub async fn get_random_omikron() -> Result<Arc<OmikronConnection>, ()> {
-    log_in!(0, PrintType::Iota, "{}", OMIKRON_CONNECTIONS.len());
     let mut rng = rand::thread_rng();
-    if let Some((_, val)) = OMIKRON_CONNECTIONS.clone().into_iter().choose(&mut rng) {
+
-        return Ok(val);
+    let keys: Vec<_> = OMIKRON_CONNECTIONS.iter().map(|e| *e.key()).collect();
-    } else {
+
-        return Err(());
+    if let Some(key) = keys.into_iter().choose(&mut rng) {
+        if let Some(entry) = OMIKRON_CONNECTIONS.get(&key) {
+            return Ok(entry.clone());
         }
+    }
+
+    Err(())
 }

src/sql/user_online_tracker.rs

@@ -81,26 +81,31 @@ pub fn untrack_many_users(user_ids: &[i64]) {
 }
 
 pub async fn untrack_omikron(omikron_id: i64) {
+    let primary_keys_to_remove: Vec<i64> = IOTA_PRIMARY_OMIKRON_CONNECTION
+        .iter()
+        .filter(|entry| *entry.value() == omikron_id)
+        .map(|entry| *entry.key())
+        .collect();
+
+    for key in primary_keys_to_remove {
+        IOTA_PRIMARY_OMIKRON_CONNECTION.remove(&key);
+    }
+
     let mut offline_iotas = Vec::new();
 
-    IOTA_OMIKRON_CONNECTIONS.retain(|iota_id, connections| {
+    for mut entry in IOTA_OMIKRON_CONNECTIONS.iter_mut() {
+        let connections = entry.value_mut();
+
         connections.retain(|id| *id != omikron_id);
 
-        if IOTA_PRIMARY_OMIKRON_CONNECTION
+        if connections.is_empty() {
-            .get(iota_id)
+            offline_iotas.push(*entry.key());
-            .map(|p| *p == omikron_id)
+        }
-            .unwrap_or(false)
-        {
-            IOTA_PRIMARY_OMIKRON_CONNECTION.remove(iota_id);
     }
 
-        if connections.is_empty() {
+    for iota_id in &offline_iotas {
-            offline_iotas.push(*iota_id);
+        IOTA_OMIKRON_CONNECTIONS.remove(iota_id);
-            false
-        } else {
-            true
     }
-    });
 
     USER_STATUS_MAP.retain(|_, status| status.omikron_id != omikron_id);