use std::sync::{Arc, Mutex}; use lru_cache::LruCache; use ruma::{events::StateEventType, EventId, OwnedEventId, RoomId}; use crate::{ observability::{FoundIn, Lookup, METRICS}, utils::error::Result, }; macro_rules! short_id_type { ($name:ident) => { #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] #[repr(transparent)] pub(crate) struct $name(u64); impl $name { pub(crate) fn new(id: u64) -> Self { Self(id) } pub(crate) fn get(&self) -> u64 { self.0 } } }; } short_id_type!(ShortRoomId); short_id_type!(ShortEventId); short_id_type!(ShortStateHash); short_id_type!(ShortStateKey); mod data; pub(crate) use data::Data; pub(crate) struct Service { db: &'static dyn Data, shorteventid_cache: Mutex>>, eventidshort_cache: Mutex>, statekeyshort_cache: Mutex>, shortstatekey_cache: Mutex>, } impl Service { pub(crate) fn new( db: &'static dyn Data, shorteventid_cache_size: usize, eventidshort_cache_size: usize, statekeyshort_cache_size: usize, shortstatekey_cache_size: usize, ) -> Self { Self { db, shorteventid_cache: Mutex::new(LruCache::new( shorteventid_cache_size, )), eventidshort_cache: Mutex::new(LruCache::new( eventidshort_cache_size, )), statekeyshort_cache: Mutex::new(LruCache::new( statekeyshort_cache_size, )), shortstatekey_cache: Mutex::new(LruCache::new( shortstatekey_cache_size, )), } } pub(crate) fn get_or_create_shorteventid( &self, event_id: &EventId, ) -> Result { let lookup = Lookup::CreateEventIdToShort; if let Some(short) = self.eventidshort_cache.lock().unwrap().get_mut(event_id) { METRICS.record_lookup(lookup, FoundIn::Cache); return Ok(*short); } let (short, created) = self.db.get_or_create_shorteventid(event_id)?; if created { METRICS.record_lookup(lookup, FoundIn::Nothing); } else { METRICS.record_lookup(lookup, FoundIn::Database); } self.eventidshort_cache .lock() .unwrap() .insert(event_id.to_owned(), short); Ok(short) } pub(crate) fn get_shortstatekey( &self, event_type: &StateEventType, state_key: &str, ) -> Result> { let lookup = Lookup::StateKeyToShort; if let Some(short) = self .statekeyshort_cache .lock() .unwrap() .get_mut(&(event_type.clone(), state_key.to_owned())) { METRICS.record_lookup(lookup, FoundIn::Cache); return Ok(Some(*short)); } let short = self.db.get_shortstatekey(event_type, state_key)?; if let Some(short) = short { METRICS.record_lookup(lookup, FoundIn::Database); self.statekeyshort_cache .lock() .unwrap() .insert((event_type.clone(), state_key.to_owned()), short); } else { METRICS.record_lookup(lookup, FoundIn::Nothing); } Ok(short) } pub(crate) fn get_or_create_shortstatekey( &self, event_type: &StateEventType, state_key: &str, ) -> Result { let lookup = Lookup::CreateStateKeyToShort; if let Some(short) = self .statekeyshort_cache .lock() .unwrap() .get_mut(&(event_type.clone(), state_key.to_owned())) { METRICS.record_lookup(lookup, FoundIn::Cache); return Ok(*short); } let (short, created) = self.db.get_or_create_shortstatekey(event_type, state_key)?; if created { METRICS.record_lookup(lookup, FoundIn::Nothing); } else { METRICS.record_lookup(lookup, FoundIn::Database); } self.statekeyshort_cache .lock() .unwrap() .insert((event_type.clone(), state_key.to_owned()), short); Ok(short) } pub(crate) fn get_eventid_from_short( &self, shorteventid: ShortEventId, ) -> Result> { let lookup = Lookup::ShortToEventId; if let Some(id) = self.shorteventid_cache.lock().unwrap().get_mut(&shorteventid) { METRICS.record_lookup(lookup, FoundIn::Cache); return Ok(Arc::clone(id)); } let event_id = self.db.get_eventid_from_short(shorteventid)?; METRICS.record_lookup(lookup, FoundIn::Database); self.shorteventid_cache .lock() .unwrap() .insert(shorteventid, Arc::clone(&event_id)); Ok(event_id) } pub(crate) fn get_statekey_from_short( &self, shortstatekey: ShortStateKey, ) -> Result<(StateEventType, String)> { let lookup = Lookup::ShortToStateKey; if let Some(id) = self.shortstatekey_cache.lock().unwrap().get_mut(&shortstatekey) { METRICS.record_lookup(lookup, FoundIn::Cache); return Ok(id.clone()); } let x = self.db.get_statekey_from_short(shortstatekey)?; METRICS.record_lookup(lookup, FoundIn::Database); self.shortstatekey_cache .lock() .unwrap() .insert(shortstatekey, x.clone()); Ok(x) } /// Returns `(shortstatehash, already_existed)` pub(crate) fn get_or_create_shortstatehash( &self, state_hash: &[u8], ) -> Result<(ShortStateHash, bool)> { self.db.get_or_create_shortstatehash(state_hash) } pub(crate) fn get_shortroomid( &self, room_id: &RoomId, ) -> Result> { self.db.get_shortroomid(room_id) } pub(crate) fn get_or_create_shortroomid( &self, room_id: &RoomId, ) -> Result { self.db.get_or_create_shortroomid(room_id) } }