use std::sync::Arc; use ruma::{events::StateEventType, EventId, RoomId}; use crate::{ database::KeyValueDatabase, observability::prelude::*, service::{ self, rooms::short::{ ShortEventId, ShortRoomId, ShortStateHash, ShortStateKey, }, }, services, utils, Error, Result, }; impl service::rooms::short::Data for KeyValueDatabase { #[t::instrument(skip(self))] fn get_or_create_shorteventid( &self, event_id: &EventId, ) -> Result<(ShortEventId, bool)> { let (short, created) = if let Some(shorteventid) = self.eventid_shorteventid.get(event_id.as_bytes())? { let shorteventid = utils::u64_from_bytes(&shorteventid).map_err(|_| { Error::bad_database("Invalid shorteventid in db.") })?; (shorteventid, false) } else { let shorteventid = services().globals.next_count()?; self.eventid_shorteventid .insert(event_id.as_bytes(), &shorteventid.to_be_bytes())?; self.shorteventid_eventid .insert(&shorteventid.to_be_bytes(), event_id.as_bytes())?; (shorteventid, true) }; Ok((ShortEventId::new(short), created)) } #[t::instrument(skip(self), fields(cache_result))] fn get_shortstatekey( &self, event_type: &StateEventType, state_key: &str, ) -> Result> { let mut db_key = event_type.to_string().as_bytes().to_vec(); db_key.push(0xFF); db_key.extend_from_slice(state_key.as_bytes()); let short = self .statekey_shortstatekey .get(&db_key)? .map(|shortstatekey| { utils::u64_from_bytes(&shortstatekey) .map_err(|_| { Error::bad_database("Invalid shortstatekey in db.") }) .map(ShortStateKey::new) }) .transpose()?; Ok(short) } #[t::instrument(skip(self))] fn get_or_create_shortstatekey( &self, event_type: &StateEventType, state_key: &str, ) -> Result<(ShortStateKey, bool)> { let mut db_key = event_type.to_string().as_bytes().to_vec(); db_key.push(0xFF); db_key.extend_from_slice(state_key.as_bytes()); let (short, created) = if let Some(shortstatekey) = self.statekey_shortstatekey.get(&db_key)? { ( utils::u64_from_bytes(&shortstatekey).map_err(|_| { Error::bad_database("Invalid shortstatekey in db.") })?, false, ) } else { let shortstatekey = services().globals.next_count()?; self.statekey_shortstatekey .insert(&db_key, &shortstatekey.to_be_bytes())?; self.shortstatekey_statekey .insert(&shortstatekey.to_be_bytes(), &db_key)?; (shortstatekey, true) }; let short = ShortStateKey::new(short); Ok((short, created)) } #[t::instrument(skip(self))] fn get_eventid_from_short( &self, shorteventid: ShortEventId, ) -> Result> { let bytes = self .shorteventid_eventid .get(&shorteventid.get().to_be_bytes())? .ok_or_else(|| { Error::bad_database("Shorteventid does not exist") })?; let event_id = EventId::parse_arc( utils::string_from_bytes(&bytes).map_err(|_| { Error::bad_database( "EventID in shorteventid_eventid is invalid unicode.", ) })?, ) .map_err(|_| { Error::bad_database("EventId in shorteventid_eventid is invalid.") })?; Ok(event_id) } #[t::instrument(skip(self))] fn get_statekey_from_short( &self, shortstatekey: ShortStateKey, ) -> Result<(StateEventType, String)> { let bytes = self .shortstatekey_statekey .get(&shortstatekey.get().to_be_bytes())? .ok_or_else(|| { Error::bad_database("Shortstatekey does not exist") })?; let mut parts = bytes.splitn(2, |&b| b == 0xFF); let eventtype_bytes = parts.next().expect("split always returns one entry"); let statekey_bytes = parts.next().ok_or_else(|| { Error::bad_database("Invalid statekey in shortstatekey_statekey.") })?; let event_type = StateEventType::from( utils::string_from_bytes(eventtype_bytes).map_err(|_| { Error::bad_database( "Event type in shortstatekey_statekey is invalid unicode.", ) })?, ); let state_key = utils::string_from_bytes(statekey_bytes).map_err(|_| { Error::bad_database( "Statekey in shortstatekey_statekey is invalid unicode.", ) })?; Ok((event_type, state_key)) } /// Returns `(shortstatehash, already_existed)` #[t::instrument(skip(self))] fn get_or_create_shortstatehash( &self, state_hash: &[u8], ) -> Result<(ShortStateHash, bool)> { let (short, existed) = if let Some(shortstatehash) = self.statehash_shortstatehash.get(state_hash)? { ( utils::u64_from_bytes(&shortstatehash).map_err(|_| { Error::bad_database("Invalid shortstatehash in db.") })?, true, ) } else { let shortstatehash = services().globals.next_count()?; self.statehash_shortstatehash .insert(state_hash, &shortstatehash.to_be_bytes())?; (shortstatehash, false) }; Ok((ShortStateHash::new(short), existed)) } fn get_shortroomid(&self, room_id: &RoomId) -> Result> { self.roomid_shortroomid .get(room_id.as_bytes())? .map(|bytes| { utils::u64_from_bytes(&bytes) .map_err(|_| { Error::bad_database("Invalid shortroomid in db.") }) .map(ShortRoomId::new) }) .transpose() } fn get_or_create_shortroomid( &self, room_id: &RoomId, ) -> Result { let short = if let Some(short) = self.roomid_shortroomid.get(room_id.as_bytes())? { utils::u64_from_bytes(&short).map_err(|_| { Error::bad_database("Invalid shortroomid in db.") })? } else { let short = services().globals.next_count()?; self.roomid_shortroomid .insert(room_id.as_bytes(), &short.to_be_bytes())?; short }; Ok(ShortRoomId::new(short)) } }