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

[Fix] (stable) Omikron - Omega connection,

Browse source
This commit is contained in:
Alex Emmet 2026-02-25 09:08:24 +01:00
commit 0c9ec2c021
3 changed files with 45 additions and 37 deletions
Download patch file Download diff file Expand all files Collapse all files

24
src/server/omikron_connection.rs
View file

@ -958,23 +958,19 @@ impl OmikronConnection {
self.send_message(&error).await;
}
pub async fn close(&self) {
if self.is_identified().await {
let omikron_id = self.get_omikron_id().await;
if omikron_id != 0 {
log_in!(PrintType::Omega, "Omikron Disconnected");
omikron_manager::remove_omikron(omikron_id).await;
user_online_tracker::untrack_omikron(omikron_id).await;
}
let omikron_id = self.get_omikron_id().await;
if omikron_id != 0 {
log_in!(omikron_id, PrintType::Omega, "Omikron Disconnected");
omikron_manager::remove_omikron(omikron_id).await;
user_online_tracker::untrack_omikron(omikron_id).await;
}
}
pub async fn handle_close(self: Arc<Self>) {
if self.is_identified().await {
let omikron_id = self.get_omikron_id().await;
if omikron_id != 0 {
log_in!(PrintType::Omega, "Omikron Disconnected");
omikron_manager::remove_omikron(omikron_id).await;
user_online_tracker::untrack_omikron(omikron_id).await;
}
let omikron_id = self.get_omikron_id().await;
if omikron_id != 0 {
log_in!(omikron_id, PrintType::Omega, "Omikron Disconnected");
omikron_manager::remove_omikron(omikron_id).await;
user_online_tracker::untrack_omikron(omikron_id).await;
}
}

25
src/server/omikron_manager.rs
View file

@ -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,27 @@ 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 = conn.get_omikron_id().await;
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(())
}

33
src/sql/user_online_tracker.rs
View file

@ -81,26 +81,31 @@ pub fn untrack_many_users(user_ids: &[i64]) {
}
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();
IOTA_OMIKRON_CONNECTIONS.retain(|iota_id, connections| {
for mut entry in IOTA_OMIKRON_CONNECTIONS.iter_mut() {
let connections = entry.value_mut();
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
} else {
true
offline_iotas.push(*entry.key());
}
});
}
for iota_id in &offline_iotas {
IOTA_OMIKRON_CONNECTIONS.remove(iota_id);
}
USER_STATUS_MAP.retain(|_, status| status.omikron_id != omikron_id);