144 lines
4.1 KiB
Rust
144 lines
4.1 KiB
Rust
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);
|
|
}
|
|
}
|