tensamin/omega
Watch
6
0
Fork
You've already forked omega
 0
Code Issues Pull requests 1 Projects Releases Packages Wiki Activity Actions

[Add] Connection Tracking

Browse source
This commit is contained in:
Alex Emmet 2026-01-29 10:58:43 +01:00
commit d2a4d0c7ae
7 changed files with 143 additions and 138 deletions
Download patch file Download diff file Expand all files Collapse all files

27
src/sql/iota_omikron_tracker.rs
View file

@ -1,27 +0,0 @@
use once_cell::sync::Lazy;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
static IOTA_OMIKRON_MAP: Lazy<Arc<RwLock<HashMap<i64, i64>>>> =
Lazy::new(|| Arc::new(RwLock::new(HashMap::new())));
pub async fn track_iota_omikron(iota: i64, omikron: i64) {
let mut c = IOTA_OMIKRON_MAP.write().await;
c.insert(iota, omikron);
}
pub async fn get_omikron_for_iota(iota: i64) -> Option<i64> {
let c = IOTA_OMIKRON_MAP.read().await;
c.get(&iota).cloned()
}
pub async fn untrack_iota(iota: i64) {
let mut c = IOTA_OMIKRON_MAP.write().await;
c.remove(&iota);
}
pub async fn untrack_by_omikron(omikron: i64) {
let mut c = IOTA_OMIKRON_MAP.write().await;
c.retain(|_, v| *v != omikron);
}

1
src/sql/mod.rs
View file

@ -1,4 +1,3 @@
pub mod connection_status;
pub mod iota_omikron_tracker;
pub mod sql;
pub mod user_online_tracker;

33
src/sql/sql.rs
View file

@ -87,7 +87,6 @@ pub async fn initialize_db() -> Result<(), sqlx::Error> {
"CREATE TABLE IF NOT EXISTS
omikrons (
id BIGINT UNSIGNED NOT NULL PRIMARY KEY,
is_active INT(1) NOT NULL DEFAULT 0,
public_key VARCHAR(255) NOT NULL COLLATE utf8mb4_bin,
location VARCHAR(255) NOT NULL COLLATE utf8mb4_bin,
ip_address VARCHAR(255) NOT NULL COLLATE utf8mb4_bin
@ -589,23 +588,6 @@ pub async fn delete_iota(id: i64) -> Result<(), sqlx::Error> {
// OMIKRONS
// ==========================================================================================
pub async fn get_random_omikron() -> Result<(i64, String, String), sqlx::Error> {
let db_lock = SQL_DB.read().await;
let pool = db_lock.as_ref().expect("Database pool is not initialized");
let row = sqlx::query_as::<_, (i64, Vec<u8>, Vec<u8>)>("SELECT id, public_key, ip_address FROM omikrons WHERE is_active = 1 ORDER BY RAND() LIMIT 1")
.fetch_optional(pool)
.await?;
match row {
Some((id, public_key, ip_address)) => Ok((
id,
String::from_utf8_lossy(&public_key).to_string(),
String::from_utf8_lossy(&ip_address).to_string(),
)),
_ => Err(sqlx::Error::RowNotFound),
}
}
pub async fn get_omikron_by_id(id: i64) -> Result<(String, String), sqlx::Error> {
let db_lock = SQL_DB.read().await;
let pool = db_lock.as_ref().expect("Database pool is not initialized");
@ -625,18 +607,3 @@ pub async fn get_omikron_by_id(id: i64) -> Result<(String, String), sqlx::Error>
_ => Err(sqlx::Error::RowNotFound),
}
}
pub async fn set_omikron_active(id: i64, active: bool) -> Result<(), sqlx::Error> {
let db_lock = SQL_DB.read().await;
let pool = db_lock.as_ref().expect("Database pool is not initialized");
let active = if active { 1 } else { 0 };
sqlx::query("UPDATE omikrons SET active = ? WHERE id = CAST(? AS UNSIGNED)")
.bind(active)
.bind(id)
.execute(pool)
.await?;
Ok(())
}

94
src/sql/user_online_tracker.rs
View file

@ -1,9 +1,7 @@
use crate::sql;
use crate::sql::connection_status::ConnectionType;
use dashmap::DashMap;
use once_cell::sync::Lazy;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
#[derive(Debug, Clone)]
pub struct UserStatus {
@ -11,42 +9,59 @@ pub struct 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<Arc<RwLock<HashMap<i64, Vec<i64>>>>> =
Lazy::new(|| Arc::new(RwLock::new(HashMap::new())));
static IOTA_OMIKRON_CONNECTIONS: Lazy<DashMap<i64, Vec<i64>>> = Lazy::new(DashMap::new);
// UserID -> UserStatus
static USER_STATUS_MAP: Lazy<Arc<RwLock<HashMap<i64, UserStatus>>>> =
Lazy::new(|| Arc::new(RwLock::new(HashMap::new())));
static USER_STATUS_MAP: Lazy<DashMap<i64, UserStatus>> = Lazy::new(DashMap::new);
pub async fn track_iota_connection(iota_id: i64, omikron_id: i64) {
let mut iota_map = IOTA_OMIKRON_CONNECTIONS.write().await;
let connections = iota_map.entry(iota_id).or_default();
if !connections.contains(&omikron_id) {
connections.push(omikron_id);
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 async fn untrack_iota_connection(iota_id: i64, omikron_id: i64) -> bool {
let mut iota_map = IOTA_OMIKRON_CONNECTIONS.write().await;
if let Some(connections) = iota_map.get_mut(&iota_id) {
pub fn untrack_iota_connection(iota_id: i64, omikron_id: i64) -> bool {
if let Some(mut connections) = IOTA_OMIKRON_CONNECTIONS.get_mut(&iota_id) {
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() {
iota_map.remove(&iota_id);
return true; // Iota is now offline
IOTA_OMIKRON_CONNECTIONS.remove(&iota_id);
return true;
}
}
false
}
pub async fn get_iota_omikron_connections(iota_id: i64) -> Option<Vec<i64>> {
let iota_map = IOTA_OMIKRON_CONNECTIONS.read().await;
iota_map.get(&iota_id).cloned()
pub fn get_iota_primary_omikron_connection(iota_id: i64) -> Option<i64> {
IOTA_PRIMARY_OMIKRON_CONNECTION.get(&iota_id).map(|v| *v)
}
pub async fn track_user_status(user_id: i64, status: ConnectionType, omikron_id: i64) {
let mut user_map = USER_STATUS_MAP.write().await;
user_map.insert(
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: ConnectionType, omikron_id: i64) {
USER_STATUS_MAP.insert(
user_id,
UserStatus {
connection_type: status,
@ -55,31 +70,34 @@ pub async fn track_user_status(user_id: i64, status: ConnectionType, omikron_id:
);
}
pub async fn get_user_status(user_id: i64) -> Option<UserStatus> {
let user_map = USER_STATUS_MAP.read().await;
user_map.get(&user_id).cloned()
pub fn get_user_status(user_id: i64) -> Option<UserStatus> {
USER_STATUS_MAP.get(&user_id).map(|v| v.clone())
}
pub async fn untrack_user(user_id: i64) {
let mut user_map = USER_STATUS_MAP.write().await;
user_map.remove(&user_id);
pub fn untrack_user(user_id: i64) {
USER_STATUS_MAP.remove(&user_id);
}
pub async fn untrack_many_users(user_ids: &[i64]) {
let mut user_map = USER_STATUS_MAP.write().await;
pub fn untrack_many_users(user_ids: &[i64]) {
for user_id in user_ids {
user_map.remove(user_id);
USER_STATUS_MAP.remove(user_id);
}
}
pub async fn untrack_omikron(omikron_id: i64) {
let mut iota_map = IOTA_OMIKRON_CONNECTIONS.write().await;
let mut user_map = USER_STATUS_MAP.write().await;
let mut offline_iotas = Vec::new();
iota_map.retain(|iota_id, connections| {
IOTA_OMIKRON_CONNECTIONS.retain(|iota_id, connections| {
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
@ -88,12 +106,12 @@ pub async fn untrack_omikron(omikron_id: i64) {
}
});
user_map.retain(|_, status| status.omikron_id != omikron_id);
USER_STATUS_MAP.retain(|_, status| status.omikron_id != omikron_id);
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_map.remove(&user.0);
USER_STATUS_MAP.remove(&user.0);
}
}
}