[Add] Structure

src/rho/rho_manager.rs

@@ -1,83 +1,129 @@
+use std::sync::Arc;
+
+use dashmap::DashMap;
+
 use super::rho_connection::RhoConnection;
-use crate::log_in;
-use crate::util::logger::PrintType;
-use std::{
-    collections::HashMap,
-    sync::{Arc, LazyLock},
-};
-use tokio::sync::RwLock;
 
-pub static RHO_CONNECTIONS: LazyLock<Arc<RwLock<HashMap<i64, Arc<RhoConnection>>>>> =
-    LazyLock::new(|| Arc::new(RwLock::new(HashMap::new())));
+/*
+ * Owns both indices for active Iota connections. Keeping the user index next
+ * to the Iota index makes binding updates atomic from the manager's caller's
+ * perspective and lets separate Omikron instances keep separate routing state.
+ */
+#[derive(Default)]
+pub struct RhoManager {
+    connections: DashMap<i64, Arc<RhoConnection>>,
+    users: DashMap<i64, Arc<RhoConnection>>,
+}
 
-pub async fn get_rho_con_for_user(user_id: i64) -> Option<Arc<RhoConnection>> {
-    let connections = RHO_CONNECTIONS.read().await;
-    for rho_connection in connections.values() {
-        let rho_user_ids = rho_connection.get_user_ids().await;
-        log_in!(
-            user_id,
-            PrintType::Client,
-            "Comparing user IDs: {:?}",
-            rho_user_ids
-        );
-        if rho_user_ids.contains(&user_id) {
-            return Some(Arc::clone(rho_connection));
+impl RhoManager {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    pub async fn get_for_user(&self, user_id: i64) -> Option<Arc<RhoConnection>> {
+        self.users.get(&user_id).map(|entry| entry.value().clone())
+    }
+
+    pub async fn contains_iota(&self, iota_id: i64) -> bool {
+        self.connections.contains_key(&iota_id)
+    }
+
+    pub async fn bind_user_to_iota(
+        &self,
+        user_id: i64,
+        iota_id: i64,
+    ) -> Option<Arc<RhoConnection>> {
+        let rho = self
+            .connections
+            .get(&iota_id)
+            .map(|entry| entry.value().clone());
+        if let Some(rho) = rho {
+            rho.add_user_id(user_id).await;
+            self.users.insert(user_id, rho.clone());
+            Some(rho)
+        } else {
+            None
         }
     }
-    None
-}
 
-#[allow(dead_code)]
-pub async fn contains_iota(iota_id: i64) -> bool {
-    let connections = RHO_CONNECTIONS.read().await;
-    connections.contains_key(&iota_id)
-}
+    pub async fn replace_users_for_iota(&self, iota_id: i64, user_ids: Vec<i64>) {
+        let rho = self
+            .connections
+            .get(&iota_id)
+            .map(|entry| entry.value().clone());
+        let Some(rho) = rho else {
+            return;
+        };
 
-/// Bind a user ID to an already tracked iota/rho connection.
-pub async fn bind_user_to_iota(user_id: i64, iota_id: i64) -> Option<Arc<RhoConnection>> {
-    let connections = RHO_CONNECTIONS.read().await;
-    if let Some(rho_connection) = connections.get(&iota_id) {
-        let rho = Arc::clone(rho_connection);
-        drop(connections);
+        let previous_users = rho.get_user_ids().await;
+        rho.set_user_ids(user_ids.clone()).await;
+        rho.get_iota_connection()
+            .set_user_ids(
+                user_ids
+                    .iter()
+                    .filter_map(|user_id| u64::try_from(*user_id).ok())
+                    .collect(),
+            )
+            .await;
 
-        rho.add_user_id(user_id).await;
+        for user_id in previous_users {
+            if self
+                .users
+                .get(&user_id)
+                .is_some_and(|entry| Arc::ptr_eq(entry.value(), &rho))
+            {
+                self.users.remove(&user_id);
+            }
+        }
+        for user_id in user_ids {
+            self.users.insert(user_id, rho.clone());
+        }
+    }
 
-        log_in!(
-            user_id,
-            PrintType::Client,
-            "Bound user {} to iota {}",
-            user_id,
-            iota_id
-        );
+    pub async fn remove_user_binding(&self, user_id: i64, iota_id: i64) {
+        let rho = self.users.get(&user_id).map(|entry| entry.value().clone());
+        if let Some(rho) = rho {
+            if rho.get_iota_id().await as i64 == iota_id {
+                self.users.remove(&user_id);
+            }
+        }
+    }
 
-        Some(rho)
-    } else {
-        None
+    pub async fn remove(&self, iota_id: i64) -> Option<Arc<RhoConnection>> {
+        let rho = self.connections.remove(&iota_id).map(|(_, rho)| rho);
+        if let Some(rho) = rho.as_ref() {
+            self.users.retain(|_, mapped| !Arc::ptr_eq(mapped, rho));
+        }
+        rho
+    }
+
+    pub async fn add(&self, rho_connection: Arc<RhoConnection>) {
+        let iota_id = rho_connection.get_iota_id().await as i64;
+        let user_ids = rho_connection.get_user_ids().await;
+        self.connections.insert(iota_id, rho_connection.clone());
+        for user_id in user_ids {
+            self.users.insert(user_id, rho_connection.clone());
+        }
+    }
+
+    pub async fn get_by_iota(&self, iota_id: i64) -> Option<Arc<RhoConnection>> {
+        self.connections
+            .get(&iota_id)
+            .map(|entry| entry.value().clone())
+    }
+
+    pub async fn connection_count(&self) -> usize {
+        self.connections.len()
+    }
+
+    pub async fn iota_ids(&self) -> Vec<i64> {
+        self.connections.iter().map(|entry| *entry.key()).collect()
+    }
+
+    pub async fn connections(&self) -> Vec<Arc<RhoConnection>> {
+        self.connections
+            .iter()
+            .map(|entry| entry.value().clone())
+            .collect()
     }
 }
-
-/// Remove a RhoConnection by Iota ID
-pub async fn remove_rho(iota_id: i64) -> Option<Arc<RhoConnection>> {
-    let mut connections = RHO_CONNECTIONS.write().await;
-    connections.remove(&iota_id)
-}
-
-/// Add a RhoConnection to the manager
-pub async fn add_rho(rho_connection: Arc<RhoConnection>) {
-    let mut connections = RHO_CONNECTIONS.write().await;
-    let iota_id = rho_connection.get_iota_id().await;
-    connections.insert(iota_id as i64, rho_connection);
-}
-
-/// Get a RhoConnection by Iota ID directly
-#[allow(dead_code)]
-pub async fn get_rho_by_iota(iota_id: i64) -> Option<Arc<RhoConnection>> {
-    let connections = RHO_CONNECTIONS.read().await;
-    connections.get(&iota_id).map(Arc::clone)
-}
-
-/// Get the count of active connections
-pub async fn connection_count() -> usize {
-    let connections = RHO_CONNECTIONS.read().await;
-    connections.len()
-}