diff --git a/src/service/rooms/short/mod.rs b/src/service/rooms/short/mod.rs index 20685442b..53d70fe85 100644 --- a/src/service/rooms/short/mod.rs +++ b/src/service/rooms/short/mod.rs @@ -2,7 +2,7 @@ pub use conduwuit::matrix::pdu::{ShortEventId, ShortId, ShortRoomId, ShortStateKey}; use conduwuit::{ - Result, err, implement, + Result, err, matrix::StateKey, pair_of, utils::{self, IterStream, ReadyExt}, @@ -57,232 +57,240 @@ fn build(args: crate::Args<'_>) -> Result> { fn name(&self) -> &str { crate::service::make_name(std::module_path!()) } } -#[implement(Service)] -pub async fn get_or_create_shorteventid(&self, event_id: &EventId) -> ShortEventId { - if let Ok(shorteventid) = self.get_shorteventid(event_id).await { - return shorteventid; +impl Service { + /// Gets or creates a short event ID + pub async fn get_or_create_shorteventid(&self, event_id: &EventId) -> ShortEventId { + if let Ok(shorteventid) = self.get_shorteventid(event_id).await { + return shorteventid; + } + + self.create_shorteventid(event_id) } - self.create_shorteventid(event_id) -} - -#[implement(Service)] -pub fn multi_get_or_create_shorteventid<'a, I>( - &'a self, - event_ids: I, -) -> impl Stream + Send + 'a -where - I: Iterator + Clone + Debug + Send + 'a, -{ - event_ids - .clone() - .stream() - .get(&self.db.eventid_shorteventid) - .zip(event_ids.into_iter().stream()) - .map(|(result, event_id)| match result { - | Ok(ref short) => utils::u64_from_u8(short), - | Err(_) => self.create_shorteventid(event_id), - }) -} - -#[implement(Service)] -fn create_shorteventid(&self, event_id: &EventId) -> ShortEventId { - const BUFSIZE: usize = size_of::(); - - let short = self.services.globals.next_count().unwrap(); - debug_assert!(size_of_val(&short) == BUFSIZE, "buffer requirement changed"); - - self.db - .eventid_shorteventid - .raw_aput::(event_id, short); - - self.db - .shorteventid_eventid - .aput_raw::(short, event_id); - - short -} - -#[implement(Service)] -pub async fn get_shorteventid(&self, event_id: &EventId) -> Result { - self.db - .eventid_shorteventid - .get(event_id) - .await - .deserialized() -} - -#[implement(Service)] -pub async fn get_or_create_shortstatekey( - &self, - event_type: &StateEventType, - state_key: &str, -) -> ShortStateKey { - const BUFSIZE: usize = size_of::(); - - if let Ok(shortstatekey) = self.get_shortstatekey(event_type, state_key).await { - return shortstatekey; - } - - let key = (event_type, state_key); - let shortstatekey = self.services.globals.next_count().unwrap(); - debug_assert!(size_of_val(&shortstatekey) == BUFSIZE, "buffer requirement changed"); - - self.db - .statekey_shortstatekey - .put_aput::(key, shortstatekey); - - self.db - .shortstatekey_statekey - .aput_put::(shortstatekey, key); - - shortstatekey -} - -#[implement(Service)] -pub async fn get_shortstatekey( - &self, - event_type: &StateEventType, - state_key: &str, -) -> Result { - let key = (event_type, state_key); - self.db - .statekey_shortstatekey - .qry(&key) - .await - .deserialized() -} - -#[implement(Service)] -pub async fn get_eventid_from_short(&self, shorteventid: ShortEventId) -> Result -where - Id: for<'de> Deserialize<'de> + Sized + ToOwned, - ::Owned: Borrow, -{ - const BUFSIZE: usize = size_of::(); - - self.db - .shorteventid_eventid - .aqry::(&shorteventid) - .await - .deserialized() - .map_err(|e| err!(Database("Failed to find EventId from short {shorteventid:?}: {e:?}"))) -} - -#[implement(Service)] -pub fn multi_get_eventid_from_short<'a, Id, S>( - &'a self, - shorteventid: S, -) -> impl Stream> + Send + 'a -where - S: Stream + Send + 'a, - Id: for<'de> Deserialize<'de> + Sized + ToOwned + 'a, - ::Owned: Borrow, -{ - shorteventid - .qry(&self.db.shorteventid_eventid) - .map(Deserialized::deserialized) -} - -#[implement(Service)] -pub async fn get_statekey_from_short( - &self, - shortstatekey: ShortStateKey, -) -> Result<(StateEventType, StateKey)> { - const BUFSIZE: usize = size_of::(); - - self.db - .shortstatekey_statekey - .aqry::(&shortstatekey) - .await - .deserialized() - .map_err(|e| { - err!(Database( - "Failed to find (StateEventType, state_key) from short {shortstatekey:?}: {e:?}" - )) - }) -} - -#[implement(Service)] -pub fn multi_get_statekey_from_short<'a, S>( - &'a self, - shortstatekey: S, -) -> impl Stream> + Send + 'a -where - S: Stream + Send + 'a, -{ - shortstatekey - .qry(&self.db.shortstatekey_statekey) - .map(Deserialized::deserialized) -} - -/// Returns (shortstatehash, already_existed) -#[implement(Service)] -pub async fn get_or_create_shortstatehash(&self, state_hash: &[u8]) -> (ShortStateHash, bool) { - const BUFSIZE: usize = size_of::(); - - if let Ok(shortstatehash) = self - .db - .statehash_shortstatehash - .get(state_hash) - .await - .deserialized() + /// Gets or creates multiple short event IDs. + pub fn multi_get_or_create_shorteventid<'a, I>( + &'a self, + event_ids: I, + ) -> impl Stream + Send + 'a + where + I: Iterator + Clone + Debug + Send + 'a, { - return (shortstatehash, true); + event_ids + .clone() + .stream() + .get(&self.db.eventid_shorteventid) + .zip(event_ids.into_iter().stream()) + .map(|(result, event_id)| match result { + | Ok(ref short) => utils::u64_from_u8(short), + | Err(_) => self.create_shorteventid(event_id), + }) } - let shortstatehash = self.services.globals.next_count().unwrap(); - debug_assert!(size_of_val(&shortstatehash) == BUFSIZE, "buffer requirement changed"); + /// Creates a short event ID + fn create_shorteventid(&self, event_id: &EventId) -> ShortEventId { + const BUFSIZE: usize = size_of::(); - self.db - .statehash_shortstatehash - .raw_aput::(state_hash, shortstatehash); + let short = self.services.globals.next_count().unwrap(); + debug_assert!(size_of_val(&short) == BUFSIZE, "buffer requirement changed"); - (shortstatehash, false) -} + self.db + .eventid_shorteventid + .raw_aput::(event_id, short); -#[implement(Service)] -pub async fn get_shortroomid(&self, room_id: &RoomId) -> Result { - self.db.roomid_shortroomid.get(room_id).await.deserialized() -} + self.db + .shorteventid_eventid + .aput_raw::(short, event_id); -#[implement(Service)] -pub async fn get_or_create_shortroomid(&self, room_id: &RoomId) -> ShortRoomId { - self.db - .roomid_shortroomid - .get(room_id) - .await - .deserialized() - .unwrap_or_else(|_| { - const BUFSIZE: usize = size_of::(); + short + } - let short = self.services.globals.next_count().unwrap(); - debug_assert!(size_of_val(&short) == BUFSIZE, "buffer requirement changed"); + /// Gets a short event ID. + pub async fn get_shorteventid(&self, event_id: &EventId) -> Result { + self.db + .eventid_shorteventid + .get(event_id) + .await + .deserialized() + } - self.db - .roomid_shortroomid - .raw_aput::(room_id, short); + /// Gets or creates a short ID for a state key pair. + pub async fn get_or_create_shortstatekey( + &self, + event_type: &StateEventType, + state_key: &str, + ) -> ShortStateKey { + const BUFSIZE: usize = size_of::(); - short + if let Ok(shortstatekey) = self.get_shortstatekey(event_type, state_key).await { + return shortstatekey; + } + + let key = (event_type, state_key); + let shortstatekey = self.services.globals.next_count().unwrap(); + debug_assert!(size_of_val(&shortstatekey) == BUFSIZE, "buffer requirement changed"); + + self.db + .statekey_shortstatekey + .put_aput::(key, shortstatekey); + + self.db + .shortstatekey_statekey + .aput_put::(shortstatekey, key); + + shortstatekey + } + + /// Gets a short ID for a state key pair. + pub async fn get_shortstatekey( + &self, + event_type: &StateEventType, + state_key: &str, + ) -> Result { + let key = (event_type, state_key); + self.db + .statekey_shortstatekey + .qry(&key) + .await + .deserialized() + } + + /// Gets a full event ID from a short event ID. + pub async fn get_eventid_from_short(&self, shorteventid: ShortEventId) -> Result + where + Id: for<'de> Deserialize<'de> + Sized + ToOwned, + ::Owned: Borrow, + { + const BUFSIZE: usize = size_of::(); + + self.db + .shorteventid_eventid + .aqry::(&shorteventid) + .await + .deserialized() + .map_err(|e| { + err!(Database("Failed to find EventId from short {shorteventid:?}: {e:?}")) + }) + } + + /// Gets multiple full event IDs from a short event ID. + pub fn multi_get_eventid_from_short<'a, Id, S>( + &'a self, + shorteventid: S, + ) -> impl Stream> + Send + 'a + where + S: Stream + Send + 'a, + Id: for<'de> Deserialize<'de> + Sized + ToOwned + 'a, + ::Owned: Borrow, + { + shorteventid + .qry(&self.db.shorteventid_eventid) + .map(Deserialized::deserialized) + } + + /// Gets a state key pair from a short state key ID. + pub async fn get_statekey_from_short( + &self, + shortstatekey: ShortStateKey, + ) -> Result<(StateEventType, StateKey)> { + const BUFSIZE: usize = size_of::(); + + self.db + .shortstatekey_statekey + .aqry::(&shortstatekey) + .await + .deserialized() + .map_err(|e| { + err!(Database( + "Failed to find (StateEventType, state_key) from short {shortstatekey:?}: \ + {e:?}" + )) + }) + } + + /// Gets multiple state key pairs from their short IDs. + pub fn multi_get_statekey_from_short<'a, S>( + &'a self, + shortstatekey: S, + ) -> impl Stream> + Send + 'a + where + S: Stream + Send + 'a, + { + shortstatekey + .qry(&self.db.shortstatekey_statekey) + .map(Deserialized::deserialized) + } + + /// Gets or creates a short state hash ID. The boolean indicates whether a + /// new short ID was created. + pub async fn get_or_create_shortstatehash( + &self, + state_hash: &[u8], + ) -> (ShortStateHash, bool) { + const BUFSIZE: usize = size_of::(); + + if let Ok(shortstatehash) = self + .db + .statehash_shortstatehash + .get(state_hash) + .await + .deserialized() + { + return (shortstatehash, true); + } + + let shortstatehash = self.services.globals.next_count().unwrap(); + debug_assert!(size_of_val(&shortstatehash) == BUFSIZE, "buffer requirement changed"); + + self.db + .statehash_shortstatehash + .raw_aput::(state_hash, shortstatehash); + + (shortstatehash, false) + } + + /// Gets a short room ID. + pub async fn get_shortroomid(&self, room_id: &RoomId) -> Result { + self.db.roomid_shortroomid.get(room_id).await.deserialized() + } + + /// Gets or creates a short room ID. + pub async fn get_or_create_shortroomid(&self, room_id: &RoomId) -> ShortRoomId { + self.db + .roomid_shortroomid + .get(room_id) + .await + .deserialized() + .unwrap_or_else(|_| { + const BUFSIZE: usize = size_of::(); + + let short = self.services.globals.next_count().unwrap(); + debug_assert!(size_of_val(&short) == BUFSIZE, "buffer requirement changed"); + + self.db + .roomid_shortroomid + .raw_aput::(room_id, short); + + short + }) + } + + /// Gets the state map associated with a short state hash. + pub async fn multi_get_state_from_short<'a, S>( + &'a self, + short_state: S, + ) -> impl Stream> + Send + 'a + where + S: Stream + Send + 'a, + { + let (short_state_keys, short_event_ids): pair_of!(Vec<_>) = short_state.unzip().await; + + StreamExt::zip( + self.multi_get_statekey_from_short(stream::iter(short_state_keys)), + self.multi_get_eventid_from_short(stream::iter(short_event_ids)), + ) + .ready_filter_map(|state_event| match state_event { + | (Ok(state_key), Ok(event_id)) => Some(Ok((state_key, event_id))), + | (Err(e), _) | (_, Err(e)) => Some(Err(e)), }) -} - -#[implement(Service)] -pub async fn multi_get_state_from_short<'a, S>( - &'a self, - short_state: S, -) -> impl Stream> + Send + 'a -where - S: Stream + Send + 'a, -{ - let (short_state_keys, short_event_ids): pair_of!(Vec<_>) = short_state.unzip().await; - - StreamExt::zip( - self.multi_get_statekey_from_short(stream::iter(short_state_keys)), - self.multi_get_eventid_from_short(stream::iter(short_event_ids)), - ) - .ready_filter_map(|state_event| match state_event { - | (Ok(state_key), Ok(event_id)) => Some(Ok((state_key, event_id))), - | (Err(e), _) | (_, Err(e)) => Some(Err(e)), - }) + } }