Sync

src/sql/user_online_tracker.rs

@@ -0,0 +1,144 @@
+use crate::sql;
+use crate::sql::connection_status::UserStatus;
+use dashmap::DashMap;
+use once_cell::sync::Lazy;
+
+#[derive(Debug, Clone)]
+pub struct UserConnection {
+    pub connection_type: UserStatus,
+    pub omikron_id: i64,
+}
+
+// IotaID -> Primary OmikronID
+static IOTA_PRIMARY_OMIKRON_CONNECTION: Lazy<DashMap<i64, i64>> = Lazy::new(DashMap::new);
+
+// IotaID -> Vec<OmikronID>
+static IOTA_OMIKRON_CONNECTIONS: Lazy<DashMap<i64, Vec<i64>>> = Lazy::new(DashMap::new);
+
+// UserID -> UserStatus
+static USER_STATUS_MAP: Lazy<DashMap<i64, UserConnection>> = Lazy::new(DashMap::new);
+
+pub fn track_iota_connection(iota_id: i64, omikron_id: i64, primary: bool) {
+    let mut entry = IOTA_OMIKRON_CONNECTIONS
+        .entry(iota_id)
+        .or_insert_with(Vec::new);
+
+    if !entry.contains(&omikron_id) {
+        entry.push(omikron_id);
+    }
+
+    if primary {
+        IOTA_PRIMARY_OMIKRON_CONNECTION.insert(iota_id, omikron_id);
+    }
+}
+
+pub fn untrack_iota_connection(iota_id: i64, omikron_id: i64) -> bool {
+    let connections_empty = if let Some(r) = IOTA_OMIKRON_CONNECTIONS.get(&iota_id) {
+        let mut vec = r.value().clone();
+        vec.retain(|&id| id != omikron_id);
+        let empty = vec.is_empty();
+        drop(r);
+        IOTA_OMIKRON_CONNECTIONS.insert(iota_id, vec);
+        empty
+    } else {
+        false
+    };
+
+    if let Some(primary_ref) = IOTA_PRIMARY_OMIKRON_CONNECTION.get(&iota_id) {
+        let primary_id = *primary_ref.value();
+        drop(primary_ref);
+        if primary_id == omikron_id {
+            IOTA_PRIMARY_OMIKRON_CONNECTION.remove(&iota_id);
+        }
+    }
+
+    if connections_empty {
+        IOTA_OMIKRON_CONNECTIONS.remove(&iota_id);
+    }
+
+    connections_empty
+}
+
+pub fn get_iota_primary_omikron_connection(iota_id: i64) -> Option<i64> {
+    IOTA_PRIMARY_OMIKRON_CONNECTION.get(&iota_id).map(|v| *v)
+}
+
+pub fn get_iota_omikron_connections(iota_id: i64) -> Option<Vec<i64>> {
+    IOTA_OMIKRON_CONNECTIONS.get(&iota_id).map(|v| v.clone())
+}
+
+pub fn track_user_status(user_id: i64, status: UserStatus, omikron_id: i64) {
+    USER_STATUS_MAP.insert(
+        user_id,
+        UserConnection {
+            connection_type: status,
+            omikron_id,
+        },
+    );
+}
+
+pub fn get_user_status(user_id: i64) -> Option<UserConnection> {
+    USER_STATUS_MAP.get(&user_id).map(|v| v.clone())
+}
+
+pub fn untrack_many_users(user_ids: &[i64]) {
+    for user_id in user_ids {
+        USER_STATUS_MAP.remove(user_id);
+    }
+}
+
+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();
+    let mut primary_to_remove = Vec::new();
+
+    // Collect iotas and primary info first
+    for r in IOTA_OMIKRON_CONNECTIONS.iter() {
+        let iota_id = *r.key();
+        let mut connections = r.value().clone();
+        connections.retain(|&id| id != omikron_id);
+
+        if connections.is_empty() {
+            offline_iotas.push(iota_id);
+        }
+
+        if IOTA_PRIMARY_OMIKRON_CONNECTION
+            .get(&iota_id)
+            .map(|p| *p == omikron_id)
+            .unwrap_or(false)
+        {
+            primary_to_remove.push(iota_id);
+        }
+
+        // Update the connections vector after filtering
+        IOTA_OMIKRON_CONNECTIONS.insert(iota_id, connections);
+    }
+
+    // Step 2: Remove primary connections safely
+    for iota_id in primary_to_remove {
+        IOTA_PRIMARY_OMIKRON_CONNECTION.remove(&iota_id);
+    }
+
+    // Step 3: Remove users that were on this omikron
+    USER_STATUS_MAP.retain(|_, status| status.omikron_id != omikron_id);
+
+    // Step 4: For offline iotas, remove associated users from USER_STATUS_MAP
+    for iota_id in offline_iotas {
+        if let Ok(users) = sql::sql::get_users_by_iota_id(iota_id).await {
+            for user in users {
+                USER_STATUS_MAP.remove(&user.0);
+            }
+        }
+        // Finally remove the empty connections vector
+        IOTA_OMIKRON_CONNECTIONS.remove(&iota_id);
+    }
+}