use std::collections::HashMap; use serde::{Deserialize, Serialize}; use std::sync::atomic::{AtomicU32, Ordering}; use tokio::sync::RwLock; const ROOT_CHANNEL_ID: u32 = 0; const USER_TREE_NAME: &[u8] = b"users"; const CHANNEL_TREE_NAME: &[u8] = b"channels"; type SessionId = u32; pub struct Db { db: sled::Db, users: sled::Tree, channels: sled::Tree, connected_users: RwLock>, next_session_id: AtomicU32, } #[derive(Clone)] pub struct User { pub id: Option, pub username: String, pub channel_id: u32, pub session_id: SessionId, } #[derive(Serialize, Deserialize)] struct PersistentUserData { id: u32, username: String, channel_id: u32, } #[derive(Serialize, Deserialize)] pub struct Channel { pub id: u32, pub name: String, } impl Db { pub fn open(path_to_db_file: &str) -> Self { let db = sled::open(path_to_db_file).expect("Unable to open database"); let users = db.open_tree(USER_TREE_NAME).unwrap(); let channels = db.open_tree(CHANNEL_TREE_NAME).unwrap(); let root_channel = bincode::serialize(&Channel { id: 0, name: "Root".to_string(), }) .unwrap(); channels .compare_and_swap( ROOT_CHANNEL_ID.to_be_bytes(), Option::<&[u8]>::None, Some(root_channel), ) .unwrap(); Db { db, users, channels, connected_users: RwLock::new(HashMap::new()), next_session_id: AtomicU32::new(0), } } pub async fn add_new_user(&self, username: String) -> u32 { let session_id = self.next_session_id.fetch_add(1, Ordering::SeqCst); let mut connected_users = self.connected_users.write().await; connected_users.insert( session_id, User { id: None, username, channel_id: ROOT_CHANNEL_ID, session_id, }, ); session_id } pub async fn get_channels(&self) -> Vec { self.channels .iter() .values() .map(|channel| bincode::deserialize(&channel.unwrap()).unwrap()) .collect() } pub async fn get_connected_users(&self) -> Vec { let users = self.connected_users.read().await; users.values().cloned().collect() } pub async fn get_user_by_session_id(&self, session_id: u32) -> Option { let connected_users = self.connected_users.read().await; if let Some(user) = connected_users.get(&session_id) { return Some(user.clone()); } None } pub async fn remove_connected_user(&self, session_id: u32) { let mut connected_users = self.connected_users.write().await; connected_users.remove(&session_id); } }