From 2f13aac0509cf136b84923a6f88fc7c6b2365e6f Mon Sep 17 00:00:00 2001
From: Alex Emmet <111742636+Alex-Emmet@users.noreply.github.com>
Date: Wed, 25 Feb 2026 09:08:24 +0100
Subject: [PATCH] [Fix] (stable) Omikron - Omega connection,

---
 src/server/omikron_connection.rs | 24 ++++++++++-------------
 src/server/omikron_manager.rs    | 25 +++++++++++++++---------
 src/sql/user_online_tracker.rs   | 33 ++++++++++++++++++--------------
 3 files changed, 45 insertions(+), 37 deletions(-)

diff --git a/src/server/omikron_connection.rs b/src/server/omikron_connection.rs
index f1634fc..9ddbb5b 100644
--- a/src/server/omikron_connection.rs
+++ b/src/server/omikron_connection.rs
@@ -958,23 +958,19 @@ 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");
-                omikron_manager::remove_omikron(omikron_id).await;
-                user_online_tracker::untrack_omikron(omikron_id).await;
-            }
+        let omikron_id = self.get_omikron_id().await;
+        if omikron_id != 0 {
+            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");
-                omikron_manager::remove_omikron(omikron_id).await;
-                user_online_tracker::untrack_omikron(omikron_id).await;
-            }
+        let omikron_id = self.get_omikron_id().await;
+        if omikron_id != 0 {
+            log_in!(omikron_id, PrintType::Omega, "Omikron Disconnected");
+            omikron_manager::remove_omikron(omikron_id).await;
+            user_online_tracker::untrack_omikron(omikron_id).await;
         }
     }
 
diff --git a/src/server/omikron_manager.rs b/src/server/omikron_manager.rs
index 864c9d6..f481702 100644
--- a/src/server/omikron_manager.rs
+++ b/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>) {
-    OMIKRON_CONNECTIONS.insert(omikron_conn.get_omikron_id().await, omikron_conn);
+pub async fn add_omikron(conn: Arc<OmikronConnection>) {
+    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);
-    } else {
-        return Err(());
+
+    let keys: Vec<_> = OMIKRON_CONNECTIONS.iter().map(|e| *e.key()).collect();
+
+    if let Some(key) = keys.into_iter().choose(&mut rng) {
+        if let Some(entry) = OMIKRON_CONNECTIONS.get(&key) {
+            return Ok(entry.clone());
+        }
     }
+
+    Err(())
 }
diff --git a/src/sql/user_online_tracker.rs b/src/sql/user_online_tracker.rs
index 8cde17b..eee1700 100644
--- a/src/sql/user_online_tracker.rs
+++ b/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
-            .get(iota_id)
-            .map(|p| *p == omikron_id)
-            .unwrap_or(false)
-        {
-            IOTA_PRIMARY_OMIKRON_CONNECTION.remove(iota_id);
-        }
-
         if connections.is_empty() {
-            offline_iotas.push(*iota_id);
-            false
-        } else {
-            true
+            offline_iotas.push(*entry.key());
         }
-    });
+    }
+
+    for iota_id in &offline_iotas {
+        IOTA_OMIKRON_CONNECTIONS.remove(iota_id);
+    }
 
     USER_STATUS_MAP.retain(|_, status| status.omikron_id != omikron_id);