use std::sync::{Arc, Weak}; use dashmap::DashMap; use crate::db::{Db, UserIdInDb}; // use this to maintain Storage #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] pub struct StorageToken { pub token: String, pub client_url: url::Url, pub machine_id: uuid::Uuid, pub user_id: UserIdInDb, } #[derive(Debug, Clone)] struct ClientInfo { storage_token: StorageToken, report_time: i64, } #[derive(Debug)] pub struct StorageInner { // some map for indexing user_clients_map: DashMap>, pub db: Db, } #[derive(Debug, Clone)] pub struct Storage(Arc); pub type WeakRefStorage = Weak; impl TryFrom for Storage { type Error = (); fn try_from(weak: Weak) -> Result { weak.upgrade().map(Storage).ok_or(()) } } impl Storage { pub fn new(db: Db) -> Self { Storage(Arc::new(StorageInner { user_clients_map: DashMap::new(), db, })) } fn remove_mid_to_client_info_map( map: &DashMap, machine_id: &uuid::Uuid, client_url: &url::Url, ) { map.remove_if(machine_id, |_, v| v.storage_token.client_url == *client_url); } fn update_mid_to_client_info_map( map: &DashMap, client_info: &ClientInfo, ) { map.entry(client_info.storage_token.machine_id) .and_modify(|e| { if e.report_time < client_info.report_time { assert_eq!( e.storage_token.machine_id, client_info.storage_token.machine_id ); *e = client_info.clone(); } }) .or_insert(client_info.clone()); } pub fn update_client(&self, stoken: StorageToken, report_time: i64) { let inner = self.0.user_clients_map.entry(stoken.user_id).or_default(); let client_info = ClientInfo { storage_token: stoken.clone(), report_time, }; Self::update_mid_to_client_info_map(&inner, &client_info); } pub fn remove_client(&self, stoken: &StorageToken) { self.0 .user_clients_map .remove_if(&stoken.user_id, |_, set| { Self::remove_mid_to_client_info_map(set, &stoken.machine_id, &stoken.client_url); set.is_empty() }); } pub fn weak_ref(&self) -> WeakRefStorage { Arc::downgrade(&self.0) } pub fn get_client_url_by_machine_id( &self, user_id: UserIdInDb, machine_id: &uuid::Uuid, ) -> Option { self.0.user_clients_map.get(&user_id).and_then(|info_map| { info_map .get(machine_id) .map(|info| info.storage_token.client_url.clone()) }) } pub fn list_user_clients(&self, user_id: UserIdInDb) -> Vec { self.0 .user_clients_map .get(&user_id) .map(|info_map| { info_map .iter() .map(|info| info.value().storage_token.client_url.clone()) .collect() }) .unwrap_or_default() } pub fn db(&self) -> &Db { &self.0.db } }