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

src/main.rs

@@ -27,6 +27,9 @@ async fn main() {
     startup();
     log_in!("Incoming messages");
     log_out!("Outgoing messages");
+
+    let _ = omikron_connection::start(9187).await;
+
     log!("Started");
     log!("  .env");
     if let Err(e) = initialize_db().await {
@@ -43,7 +46,8 @@ async fn main() {
     } else {
         log!("  Users");
     }
+
     let _ = server::server::start(9188).await;
-    let _ = omikron_connection::OmikronServer::start(9187).await;
+
     tokio::signal::ctrl_c().await.unwrap();
 }

src/server/api.rs

@@ -65,16 +65,20 @@ pub async fn handle(path: &str, body_string: Option<String>) -> HttpResponse {
         // ==================================================
         ["api", "get", "omikron"] => {
             if let Ok(omikron_conn) = get_random_omikron().await {
-                let id = omikron_conn.get_omikron_id().await;
+                if let Some(id) = omikron_conn.get_omikron_id().await {
+                    if let Ok((public_key, ip_address)) = sql::get_omikron_by_id(id).await {
+                        let mut res = JsonValue::new_object();
+                        res["status"] = "success".into();
+                        res["id"] = id.into();
+                        res["public_key"] = public_key.into();
+                        res["ip_address"] = ip_address.into();
 
-                if let Ok((public_key, ip_address)) = sql::get_omikron_by_id(id).await {
-                    let mut res = JsonValue::new_object();
-                    res["status"] = "success".into();
-                    res["id"] = id.into();
-                    res["public_key"] = public_key.into();
-                    res["ip_address"] = ip_address.into();
-
-                    (StatusCode::OK, res.dump())
+                        (StatusCode::OK, res.dump())
+                    } else {
+                        let mut res = JsonValue::new_object();
+                        res["status"] = "error".into();
+                        (StatusCode::INTERNAL_SERVER_ERROR, res.dump())
+                    }
                 } else {
                     let mut res = JsonValue::new_object();
                     res["status"] = "error".into();

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,33 @@ 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 = match conn.get_omikron_id().await {
+        Some(id) => id,
+        _ => {
+            conn.close().await;
+            return;
+        }
+    };
+
+    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(())
 }

src/sql/user_online_tracker.rs

@@ -88,7 +88,16 @@ pub fn untrack_many_users(user_ids: &[i64]) {
 }
 
 pub async fn untrack_omikron(omikron_id: i64) {
-    // Step 1: Collect all iota_ids where this omikron_id is connected
+    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();
     let mut primary_to_remove = Vec::new();