Add wrapper types for short IDs

This commit is contained in:
Lambda 2024-08-27 14:27:12 +00:00
parent f1642c92d1
commit b0f33207fe
28 changed files with 427 additions and 232 deletions

View file

@ -5,12 +5,21 @@ use ruma::{events::StateEventType, EventId, RoomId};
use crate::{
database::KeyValueDatabase,
observability::{FoundIn, Lookup, METRICS},
service, services, utils, Error, Result,
service::{
self,
rooms::short::{
ShortEventId, ShortRoomId, ShortStateHash, ShortStateKey,
},
},
services, utils, Error, Result,
};
impl service::rooms::short::Data for KeyValueDatabase {
#[tracing::instrument(skip(self))]
fn get_or_create_shorteventid(&self, event_id: &EventId) -> Result<u64> {
fn get_or_create_shorteventid(
&self,
event_id: &EventId,
) -> Result<ShortEventId> {
let lookup = Lookup::CreateEventIdToShort;
if let Some(short) =
@ -39,6 +48,8 @@ impl service::rooms::short::Data for KeyValueDatabase {
shorteventid
};
let short = ShortEventId::new(short);
self.eventidshort_cache
.lock()
.unwrap()
@ -52,7 +63,7 @@ impl service::rooms::short::Data for KeyValueDatabase {
&self,
event_type: &StateEventType,
state_key: &str,
) -> Result<Option<u64>> {
) -> Result<Option<ShortStateKey>> {
let lookup = Lookup::StateKeyToShort;
if let Some(short) = self
@ -73,9 +84,11 @@ impl service::rooms::short::Data for KeyValueDatabase {
.statekey_shortstatekey
.get(&db_key)?
.map(|shortstatekey| {
utils::u64_from_bytes(&shortstatekey).map_err(|_| {
Error::bad_database("Invalid shortstatekey in db.")
})
utils::u64_from_bytes(&shortstatekey)
.map_err(|_| {
Error::bad_database("Invalid shortstatekey in db.")
})
.map(ShortStateKey::new)
})
.transpose()?;
@ -98,7 +111,7 @@ impl service::rooms::short::Data for KeyValueDatabase {
&self,
event_type: &StateEventType,
state_key: &str,
) -> Result<u64> {
) -> Result<ShortStateKey> {
let lookup = Lookup::CreateStateKeyToShort;
if let Some(short) = self
@ -134,6 +147,8 @@ impl service::rooms::short::Data for KeyValueDatabase {
shortstatekey
};
let short = ShortStateKey::new(short);
self.statekeyshort_cache
.lock()
.unwrap()
@ -145,7 +160,7 @@ impl service::rooms::short::Data for KeyValueDatabase {
#[tracing::instrument(skip(self))]
fn get_eventid_from_short(
&self,
shorteventid: u64,
shorteventid: ShortEventId,
) -> Result<Arc<EventId>> {
let lookup = Lookup::ShortToEventId;
@ -158,7 +173,7 @@ impl service::rooms::short::Data for KeyValueDatabase {
let bytes = self
.shorteventid_eventid
.get(&shorteventid.to_be_bytes())?
.get(&shorteventid.get().to_be_bytes())?
.ok_or_else(|| {
Error::bad_database("Shorteventid does not exist")
})?;
@ -187,7 +202,7 @@ impl service::rooms::short::Data for KeyValueDatabase {
#[tracing::instrument(skip(self))]
fn get_statekey_from_short(
&self,
shortstatekey: u64,
shortstatekey: ShortStateKey,
) -> Result<(StateEventType, String)> {
let lookup = Lookup::ShortToStateKey;
@ -200,7 +215,7 @@ impl service::rooms::short::Data for KeyValueDatabase {
let bytes = self
.shortstatekey_statekey
.get(&shortstatekey.to_be_bytes())?
.get(&shortstatekey.get().to_be_bytes())?
.ok_or_else(|| {
Error::bad_database("Shortstatekey does not exist")
})?;
@ -244,51 +259,56 @@ impl service::rooms::short::Data for KeyValueDatabase {
fn get_or_create_shortstatehash(
&self,
state_hash: &[u8],
) -> Result<(u64, bool)> {
Ok(
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)
},
)
) -> 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<Option<u64>> {
fn get_shortroomid(&self, room_id: &RoomId) -> Result<Option<ShortRoomId>> {
self.roomid_shortroomid
.get(room_id.as_bytes())?
.map(|bytes| {
utils::u64_from_bytes(&bytes).map_err(|_| {
Error::bad_database("Invalid shortroomid in db.")
})
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<u64> {
Ok(
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
},
)
fn get_or_create_shortroomid(
&self,
room_id: &RoomId,
) -> Result<ShortRoomId> {
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))
}
}