backups/grapevine
1
0
Fork 0
mirror of https://gitlab.computer.surgery/matrix/grapevine.git synced 2025-12-18 08:11:24 +01:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add wrapper types for short IDs

Browse source
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
Download patch file Download diff file Expand all files Collapse all files

1
src/database/key_value/globals.rs
View file

@ -74,6 +74,7 @@ impl service::globals::Data for KeyValueDatabase {
.ok()
.flatten()
.expect("room exists")
.get()
.to_be_bytes()
.to_vec();

23
src/database/key_value/rooms/auth_chain.rs
View file

@ -3,15 +3,16 @@ use std::{collections::HashSet, mem::size_of, sync::Arc};
use crate::{
database::KeyValueDatabase,
observability::{FoundIn, Lookup, METRICS},
service, utils, Result,
service::{self, rooms::short::ShortEventId},
utils, Result,
};
impl service::rooms::auth_chain::Data for KeyValueDatabase {
#[tracing::instrument(skip(self, key))]
fn get_cached_eventid_authchain(
&self,
key: &[u64],
) -> Result<Option<Arc<HashSet<u64>>>> {
key: &[ShortEventId],
) -> Result<Option<Arc<HashSet<ShortEventId>>>> {
let lookup = Lookup::AuthChain;
// Check RAM cache
@ -26,13 +27,15 @@ impl service::rooms::auth_chain::Data for KeyValueDatabase {
// Check DB cache
let chain = self
.shorteventid_authchain
.get(&key[0].to_be_bytes())?
.get(&key[0].get().to_be_bytes())?
.map(|chain| {
chain
.chunks_exact(size_of::<u64>())
.map(|chunk| {
utils::u64_from_bytes(chunk)
.expect("byte length is correct")
ShortEventId::new(
utils::u64_from_bytes(chunk)
.expect("byte length is correct"),
)
})
.collect()
});
@ -57,16 +60,16 @@ impl service::rooms::auth_chain::Data for KeyValueDatabase {
fn cache_auth_chain(
&self,
key: Vec<u64>,
auth_chain: Arc<HashSet<u64>>,
key: Vec<ShortEventId>,
auth_chain: Arc<HashSet<ShortEventId>>,
) -> Result<()> {
// Only persist single events in db
if key.len() == 1 {
self.shorteventid_authchain.insert(
&key[0].to_be_bytes(),
&key[0].get().to_be_bytes(),
&auth_chain
.iter()
.flat_map(|s| s.to_be_bytes().to_vec())
.flat_map(|s| s.get().to_be_bytes().to_vec())
.collect::<Vec<u8>>(),
)?;
}

2
src/database/key_value/rooms/metadata.rs
View file

@ -8,7 +8,7 @@ impl service::rooms::metadata::Data for KeyValueDatabase {
#[tracing::instrument(skip(self))]
fn exists(&self, room_id: &RoomId) -> Result<bool> {
let prefix = match services().rooms.short.get_shortroomid(room_id)? {
Some(b) => b.to_be_bytes().to_vec(),
Some(b) => b.get().to_be_bytes().to_vec(),
None => return Ok(false),
};

9
src/database/key_value/rooms/pdu_metadata.rs
View file

@ -6,7 +6,10 @@ use crate::{
database::KeyValueDatabase,
service::{
self,
rooms::timeline::{PduCount, PduId},
rooms::{
short::ShortRoomId,
timeline::{PduCount, PduId},
},
},
services, utils, Error, PduEvent, Result,
};
@ -22,7 +25,7 @@ impl service::rooms::pdu_metadata::Data for KeyValueDatabase {
fn relations_until<'a>(
&'a self,
user_id: &'a UserId,
shortroomid: u64,
shortroomid: ShortRoomId,
target: u64,
until: PduCount,
) -> Result<Box<dyn Iterator<Item = Result<(PduCount, PduEvent)>> + 'a>>
@ -51,7 +54,7 @@ impl service::rooms::pdu_metadata::Data for KeyValueDatabase {
Error::bad_database("Invalid count in tofrom_relation.")
})?;
let mut pduid = shortroomid.to_be_bytes().to_vec();
let mut pduid = shortroomid.get().to_be_bytes().to_vec();
pduid.extend_from_slice(&from.to_be_bytes());
let pduid = PduId::new(pduid);

14
src/database/key_value/rooms/search.rs
View file

@ -2,7 +2,10 @@ use ruma::RoomId;
use crate::{
database::KeyValueDatabase,
service::{self, rooms::timeline::PduId},
service::{
self,
rooms::{short::ShortRoomId, timeline::PduId},
},
services, utils, Result,
};
@ -21,12 +24,12 @@ impl service::rooms::search::Data for KeyValueDatabase {
#[tracing::instrument(skip(self))]
fn index_pdu(
&self,
shortroomid: u64,
shortroomid: ShortRoomId,
pdu_id: &PduId,
message_body: &str,
) -> Result<()> {
let mut batch = tokenize(message_body).map(|word| {
let mut key = shortroomid.to_be_bytes().to_vec();
let mut key = shortroomid.get().to_be_bytes().to_vec();
key.extend_from_slice(word.as_bytes());
key.push(0xFF);
// TODO: currently we save the room id a second time here
@ -40,12 +43,12 @@ impl service::rooms::search::Data for KeyValueDatabase {
#[tracing::instrument(skip(self))]
fn deindex_pdu(
&self,
shortroomid: u64,
shortroomid: ShortRoomId,
pdu_id: &PduId,
message_body: &str,
) -> Result<()> {
let batch = tokenize(message_body).map(|word| {
let mut key = shortroomid.to_be_bytes().to_vec();
let mut key = shortroomid.get().to_be_bytes().to_vec();
key.extend_from_slice(word.as_bytes());
key.push(0xFF);
// TODO: currently we save the room id a second time here
@ -73,6 +76,7 @@ impl service::rooms::search::Data for KeyValueDatabase {
.short
.get_shortroomid(room_id)?
.expect("room exists")
.get()
.to_be_bytes()
.to_vec();

116
src/database/key_value/rooms/short.rs
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))
}
}

39
src/database/key_value/rooms/state.rs
View file

@ -4,21 +4,30 @@ use ruma::{EventId, OwnedEventId, OwnedRoomId, RoomId};
use crate::{
database::KeyValueDatabase,
service::{self, globals::marker},
service::{
self,
globals::marker,
rooms::short::{ShortEventId, ShortStateHash},
},
utils::{self, on_demand_hashmap::KeyToken},
Error, Result,
};
impl service::rooms::state::Data for KeyValueDatabase {
fn get_room_shortstatehash(&self, room_id: &RoomId) -> Result<Option<u64>> {
fn get_room_shortstatehash(
&self,
room_id: &RoomId,
) -> Result<Option<ShortStateHash>> {
self.roomid_shortstatehash.get(room_id.as_bytes())?.map_or(
Ok(None),
|bytes| {
Ok(Some(utils::u64_from_bytes(&bytes).map_err(|_| {
Error::bad_database(
"Invalid shortstatehash in roomid_shortstatehash",
)
})?))
Ok(Some(ShortStateHash::new(
utils::u64_from_bytes(&bytes).map_err(|_| {
Error::bad_database(
"Invalid shortstatehash in roomid_shortstatehash",
)
})?,
)))
},
)
}
@ -26,21 +35,23 @@ impl service::rooms::state::Data for KeyValueDatabase {
fn set_room_state(
&self,
room_id: &KeyToken<OwnedRoomId, marker::State>,
new_shortstatehash: u64,
new_shortstatehash: ShortStateHash,
) -> Result<()> {
self.roomid_shortstatehash
.insert(room_id.as_bytes(), &new_shortstatehash.to_be_bytes())?;
self.roomid_shortstatehash.insert(
room_id.as_bytes(),
&new_shortstatehash.get().to_be_bytes(),
)?;
Ok(())
}
fn set_event_state(
&self,
shorteventid: u64,
shortstatehash: u64,
shorteventid: ShortEventId,
shortstatehash: ShortStateHash,
) -> Result<()> {
self.shorteventid_shortstatehash.insert(
&shorteventid.to_be_bytes(),
&shortstatehash.to_be_bytes(),
&shorteventid.get().to_be_bytes(),
&shortstatehash.get().to_be_bytes(),
)?;
Ok(())
}

39
src/database/key_value/rooms/state_accessor.rs
View file

@ -4,16 +4,20 @@ use async_trait::async_trait;
use ruma::{events::StateEventType, EventId, RoomId};
use crate::{
database::KeyValueDatabase, service, services, utils, Error, PduEvent,
Result,
database::KeyValueDatabase,
service::{
self,
rooms::short::{ShortStateHash, ShortStateKey},
},
services, utils, Error, PduEvent, Result,
};
#[async_trait]
impl service::rooms::state_accessor::Data for KeyValueDatabase {
async fn state_full_ids(
&self,
shortstatehash: u64,
) -> Result<HashMap<u64, Arc<EventId>>> {
shortstatehash: ShortStateHash,
) -> Result<HashMap<ShortStateKey, Arc<EventId>>> {
let full_state = services()
.rooms
.state_compressor
@ -40,7 +44,7 @@ impl service::rooms::state_accessor::Data for KeyValueDatabase {
async fn state_full(
&self,
shortstatehash: u64,
shortstatehash: ShortStateHash,
) -> Result<HashMap<(StateEventType, String), Arc<PduEvent>>> {
let full_state = services()
.rooms
@ -87,7 +91,7 @@ impl service::rooms::state_accessor::Data for KeyValueDatabase {
/// `state_key`).
fn state_get_id(
&self,
shortstatehash: u64,
shortstatehash: ShortStateHash,
event_type: &StateEventType,
state_key: &str,
) -> Result<Option<Arc<EventId>>> {
@ -105,7 +109,7 @@ impl service::rooms::state_accessor::Data for KeyValueDatabase {
.full_state;
Ok(full_state
.iter()
.find(|bytes| bytes.starts_with(&shortstatekey.to_be_bytes()))
.find(|compressed| compressed.state == shortstatekey)
.and_then(|compressed| {
services()
.rooms
@ -120,7 +124,7 @@ impl service::rooms::state_accessor::Data for KeyValueDatabase {
/// `state_key`).
fn state_get(
&self,
shortstatehash: u64,
shortstatehash: ShortStateHash,
event_type: &StateEventType,
state_key: &str,
) -> Result<Option<Arc<PduEvent>>> {
@ -131,19 +135,24 @@ impl service::rooms::state_accessor::Data for KeyValueDatabase {
}
/// Returns the state hash for this pdu.
fn pdu_shortstatehash(&self, event_id: &EventId) -> Result<Option<u64>> {
fn pdu_shortstatehash(
&self,
event_id: &EventId,
) -> Result<Option<ShortStateHash>> {
self.eventid_shorteventid.get(event_id.as_bytes())?.map_or(
Ok(None),
|shorteventid| {
self.shorteventid_shortstatehash
.get(&shorteventid)?
.map(|bytes| {
utils::u64_from_bytes(&bytes).map_err(|_| {
Error::bad_database(
"Invalid shortstatehash bytes in \
shorteventid_shortstatehash",
)
})
utils::u64_from_bytes(&bytes)
.map_err(|_| {
Error::bad_database(
"Invalid shortstatehash bytes in \
shorteventid_shortstatehash",
)
})
.map(ShortStateHash::new)
})
.transpose()
},

37
src/database/key_value/rooms/state_compressor.rs
View file

@ -2,19 +2,28 @@ use std::{collections::HashSet, mem::size_of, sync::Arc};
use crate::{
database::KeyValueDatabase,
service::{self, rooms::state_compressor::data::StateDiff},
service::{
self,
rooms::{
short::ShortStateHash,
state_compressor::{data::StateDiff, CompressedStateEvent},
},
},
utils, Error, Result,
};
impl service::rooms::state_compressor::Data for KeyValueDatabase {
fn get_statediff(&self, shortstatehash: u64) -> Result<StateDiff> {
fn get_statediff(
&self,
shortstatehash: ShortStateHash,
) -> Result<StateDiff> {
let value = self
.shortstatehash_statediff
.get(&shortstatehash.to_be_bytes())?
.get(&shortstatehash.get().to_be_bytes())?
.ok_or_else(|| Error::bad_database("State hash does not exist"))?;
let parent = utils::u64_from_bytes(&value[0..size_of::<u64>()])
.expect("bytes have right length");
let parent = (parent != 0).then_some(parent);
let parent = (parent != 0).then_some(ShortStateHash::new(parent));
let mut add_mode = true;
let mut added = HashSet::new();
@ -28,10 +37,13 @@ impl service::rooms::state_compressor::Data for KeyValueDatabase {
continue;
}
if add_mode {
added.insert(v.try_into().expect("we checked the size above"));
added.insert(CompressedStateEvent::from_bytes(
v.try_into().expect("we checked the size above"),
));
} else {
removed
.insert(v.try_into().expect("we checked the size above"));
removed.insert(CompressedStateEvent::from_bytes(
v.try_into().expect("we checked the size above"),
));
}
i += 2 * size_of::<u64>();
}
@ -45,22 +57,23 @@ impl service::rooms::state_compressor::Data for KeyValueDatabase {
fn save_statediff(
&self,
shortstatehash: u64,
shortstatehash: ShortStateHash,
diff: StateDiff,
) -> Result<()> {
let mut value = diff.parent.unwrap_or(0).to_be_bytes().to_vec();
let mut value =
diff.parent.map_or(0, |h| h.get()).to_be_bytes().to_vec();
for new in diff.added.iter() {
value.extend_from_slice(&new[..]);
value.extend_from_slice(&new.as_bytes());
}
if !diff.removed.is_empty() {
value.extend_from_slice(&0_u64.to_be_bytes());
for removed in diff.removed.iter() {
value.extend_from_slice(&removed[..]);
value.extend_from_slice(&removed.as_bytes());
}
}
self.shortstatehash_statediff
.insert(&shortstatehash.to_be_bytes(), &value)
.insert(&shortstatehash.get().to_be_bytes(), &value)
}
}

1
src/database/key_value/rooms/threads.rs
View file

@ -24,6 +24,7 @@ impl service::rooms::threads::Data for KeyValueDatabase {
.short
.get_shortroomid(room_id)?
.expect("room exists")
.get()
.to_be_bytes()
.to_vec();

1
src/database/key_value/rooms/timeline.rs
View file

@ -383,6 +383,7 @@ fn count_to_id(
.ok_or_else(|| {
Error::bad_database("Looked for bad shortroomid in timeline")
})?
.get()
.to_be_bytes()
.to_vec();
let mut pdu_id = prefix.clone();

28
src/database/key_value/rooms/user.rs
View file

@ -1,7 +1,9 @@
use ruma::{OwnedRoomId, OwnedUserId, RoomId, UserId};
use crate::{
database::KeyValueDatabase, service, services, utils, Error, Result,
database::KeyValueDatabase,
service::{self, rooms::short::ShortStateHash},
services, utils, Error, Result,
};
impl service::rooms::user::Data for KeyValueDatabase {
@ -95,7 +97,7 @@ impl service::rooms::user::Data for KeyValueDatabase {
&self,
room_id: &RoomId,
token: u64,
shortstatehash: u64,
shortstatehash: ShortStateHash,
) -> Result<()> {
let shortroomid = services()
.rooms
@ -103,36 +105,38 @@ impl service::rooms::user::Data for KeyValueDatabase {
.get_shortroomid(room_id)?
.expect("room exists");
let mut key = shortroomid.to_be_bytes().to_vec();
let mut key = shortroomid.get().to_be_bytes().to_vec();
key.extend_from_slice(&token.to_be_bytes());
self.roomsynctoken_shortstatehash
.insert(&key, &shortstatehash.to_be_bytes())
.insert(&key, &shortstatehash.get().to_be_bytes())
}
fn get_token_shortstatehash(
&self,
room_id: &RoomId,
token: u64,
) -> Result<Option<u64>> {
) -> Result<Option<ShortStateHash>> {
let shortroomid = services()
.rooms
.short
.get_shortroomid(room_id)?
.expect("room exists");
let mut key = shortroomid.to_be_bytes().to_vec();
let mut key = shortroomid.get().to_be_bytes().to_vec();
key.extend_from_slice(&token.to_be_bytes());
self.roomsynctoken_shortstatehash
.get(&key)?
.map(|bytes| {
utils::u64_from_bytes(&bytes).map_err(|_| {
Error::bad_database(
"Invalid shortstatehash in \
roomsynctoken_shortstatehash",
)
})
utils::u64_from_bytes(&bytes)
.map_err(|_| {
Error::bad_database(
"Invalid shortstatehash in \
roomsynctoken_shortstatehash",
)
})
.map(ShortStateHash::new)
})
.transpose()
}