Add PduId wrapper struct

Death to Vec<u8>
This commit is contained in:
Lambda 2024-08-26 16:47:50 +00:00
parent 341f4213d0
commit 26322d5a95
15 changed files with 110 additions and 71 deletions

View file

@ -4,7 +4,10 @@ use ruma::{EventId, RoomId, UserId};
use crate::{
database::KeyValueDatabase,
service::{self, rooms::timeline::PduCount},
service::{
self,
rooms::timeline::{PduCount, PduId},
},
services, utils, Error, PduEvent, Result,
};
@ -50,6 +53,7 @@ impl service::rooms::pdu_metadata::Data for KeyValueDatabase {
let mut pduid = shortroomid.to_be_bytes().to_vec();
pduid.extend_from_slice(&from.to_be_bytes());
let pduid = PduId::new(pduid);
let mut pdu = services()
.rooms

View file

@ -1,6 +1,10 @@
use ruma::RoomId;
use crate::{database::KeyValueDatabase, service, services, utils, Result};
use crate::{
database::KeyValueDatabase,
service::{self, rooms::timeline::PduId},
services, utils, Result,
};
/// Splits a string into tokens used as keys in the search inverted index
///
@ -18,7 +22,7 @@ impl service::rooms::search::Data for KeyValueDatabase {
fn index_pdu(
&self,
shortroomid: u64,
pdu_id: &[u8],
pdu_id: &PduId,
message_body: &str,
) -> Result<()> {
let mut batch = tokenize(message_body).map(|word| {
@ -26,7 +30,7 @@ impl service::rooms::search::Data for KeyValueDatabase {
key.extend_from_slice(word.as_bytes());
key.push(0xFF);
// TODO: currently we save the room id a second time here
key.extend_from_slice(pdu_id);
key.extend_from_slice(pdu_id.as_bytes());
(key, Vec::new())
});
@ -37,7 +41,7 @@ impl service::rooms::search::Data for KeyValueDatabase {
fn deindex_pdu(
&self,
shortroomid: u64,
pdu_id: &[u8],
pdu_id: &PduId,
message_body: &str,
) -> Result<()> {
let batch = tokenize(message_body).map(|word| {
@ -45,7 +49,7 @@ impl service::rooms::search::Data for KeyValueDatabase {
key.extend_from_slice(word.as_bytes());
key.push(0xFF);
// TODO: currently we save the room id a second time here
key.extend_from_slice(pdu_id);
key.extend_from_slice(pdu_id.as_bytes());
key
});
@ -62,7 +66,7 @@ impl service::rooms::search::Data for KeyValueDatabase {
&'a self,
room_id: &RoomId,
search_string: &str,
) -> Result<Option<(Box<dyn Iterator<Item = Vec<u8>> + 'a>, Vec<String>)>>
) -> Result<Option<(Box<dyn Iterator<Item = PduId> + 'a>, Vec<String>)>>
{
let prefix = services()
.rooms
@ -87,12 +91,14 @@ impl service::rooms::search::Data for KeyValueDatabase {
// Newest pdus first
.iter_from(&last_possible_id, true)
.take_while(move |(k, _)| k.starts_with(&prefix2))
.map(move |(key, _)| key[prefix3.len()..].to_vec())
.map(move |(key, _)| PduId::new(key[prefix3.len()..].to_vec()))
});
// We compare b with a because we reversed the iterator earlier
let Some(common_elements) =
utils::common_elements(iterators, |a, b| b.cmp(a))
utils::common_elements(iterators, |a, b| {
b.as_bytes().cmp(a.as_bytes())
})
else {
return Ok(None);
};

View file

@ -6,8 +6,9 @@ use ruma::{
};
use crate::{
database::KeyValueDatabase, service, services, utils, Error, PduEvent,
Result,
database::KeyValueDatabase,
service::{self, rooms::timeline::PduId},
services, utils, Error, PduEvent, Result,
};
impl service::rooms::threads::Data for KeyValueDatabase {
@ -42,6 +43,9 @@ impl service::rooms::threads::Data for KeyValueDatabase {
"Invalid pduid in threadid_userids.",
)
})?;
let pduid = PduId::new(pduid);
let mut pdu = services()
.rooms
.timeline
@ -61,7 +65,7 @@ impl service::rooms::threads::Data for KeyValueDatabase {
fn update_participants(
&self,
root_id: &[u8],
root_id: &PduId,
participants: &[OwnedUserId],
) -> Result<()> {
let users = participants
@ -70,16 +74,16 @@ impl service::rooms::threads::Data for KeyValueDatabase {
.collect::<Vec<_>>()
.join(&[0xFF][..]);
self.threadid_userids.insert(root_id, &users)?;
self.threadid_userids.insert(root_id.as_bytes(), &users)?;
Ok(())
}
fn get_participants(
&self,
root_id: &[u8],
root_id: &PduId,
) -> Result<Option<Vec<OwnedUserId>>> {
if let Some(users) = self.threadid_userids.get(root_id)? {
if let Some(users) = self.threadid_userids.get(root_id.as_bytes())? {
Ok(Some(
users
.split(|b| *b == 0xFF)

View file

@ -10,7 +10,8 @@ use tracing::error;
use crate::{
database::KeyValueDatabase,
observability::{FoundIn, Lookup, METRICS},
service, services, utils, Error, PduEvent, Result,
service::{self, rooms::timeline::PduId},
services, utils, Error, PduEvent, Result,
};
impl service::rooms::timeline::Data for KeyValueDatabase {
@ -102,8 +103,8 @@ impl service::rooms::timeline::Data for KeyValueDatabase {
}
/// Returns the pdu's id.
fn get_pdu_id(&self, event_id: &EventId) -> Result<Option<Vec<u8>>> {
self.eventid_pduid.get(event_id.as_bytes())
fn get_pdu_id(&self, event_id: &EventId) -> Result<Option<PduId>> {
self.eventid_pduid.get(event_id.as_bytes()).map(|x| x.map(PduId::new))
}
/// Returns the pdu.
@ -170,8 +171,8 @@ impl service::rooms::timeline::Data for KeyValueDatabase {
/// Returns the pdu.
///
/// This does __NOT__ check the outliers `Tree`.
fn get_pdu_from_id(&self, pdu_id: &[u8]) -> Result<Option<PduEvent>> {
self.pduid_pdu.get(pdu_id)?.map_or(Ok(None), |pdu| {
fn get_pdu_from_id(&self, pdu_id: &PduId) -> Result<Option<PduEvent>> {
self.pduid_pdu.get(pdu_id.as_bytes())?.map_or(Ok(None), |pdu| {
Ok(Some(
serde_json::from_slice(&pdu)
.map_err(|_| Error::bad_database("Invalid PDU in db."))?,
@ -182,9 +183,9 @@ impl service::rooms::timeline::Data for KeyValueDatabase {
/// Returns the pdu as a `BTreeMap<String, CanonicalJsonValue>`.
fn get_pdu_json_from_id(
&self,
pdu_id: &[u8],
pdu_id: &PduId,
) -> Result<Option<CanonicalJsonObject>> {
self.pduid_pdu.get(pdu_id)?.map_or(Ok(None), |pdu| {
self.pduid_pdu.get(pdu_id.as_bytes())?.map_or(Ok(None), |pdu| {
Ok(Some(
serde_json::from_slice(&pdu)
.map_err(|_| Error::bad_database("Invalid PDU in db."))?,
@ -194,13 +195,13 @@ impl service::rooms::timeline::Data for KeyValueDatabase {
fn append_pdu(
&self,
pdu_id: &[u8],
pdu_id: &PduId,
pdu: &PduEvent,
json: &CanonicalJsonObject,
count: u64,
) -> Result<()> {
self.pduid_pdu.insert(
pdu_id,
pdu_id.as_bytes(),
&serde_json::to_vec(json)
.expect("CanonicalJsonObject is always a valid"),
)?;
@ -210,7 +211,8 @@ impl service::rooms::timeline::Data for KeyValueDatabase {
.unwrap()
.insert(pdu.room_id.clone(), PduCount::Normal(count));
self.eventid_pduid.insert(pdu.event_id.as_bytes(), pdu_id)?;
self.eventid_pduid
.insert(pdu.event_id.as_bytes(), pdu_id.as_bytes())?;
self.eventid_outlierpdu.remove(pdu.event_id.as_bytes())?;
Ok(())
@ -218,17 +220,17 @@ impl service::rooms::timeline::Data for KeyValueDatabase {
fn prepend_backfill_pdu(
&self,
pdu_id: &[u8],
pdu_id: &PduId,
event_id: &EventId,
json: &CanonicalJsonObject,
) -> Result<()> {
self.pduid_pdu.insert(
pdu_id,
pdu_id.as_bytes(),
&serde_json::to_vec(json)
.expect("CanonicalJsonObject is always a valid"),
)?;
self.eventid_pduid.insert(event_id.as_bytes(), pdu_id)?;
self.eventid_pduid.insert(event_id.as_bytes(), pdu_id.as_bytes())?;
self.eventid_outlierpdu.remove(event_id.as_bytes())?;
Ok(())
@ -237,13 +239,13 @@ impl service::rooms::timeline::Data for KeyValueDatabase {
/// Removes a pdu and creates a new one with the same id.
fn replace_pdu(
&self,
pdu_id: &[u8],
pdu_id: &PduId,
pdu_json: &CanonicalJsonObject,
pdu: &PduEvent,
) -> Result<()> {
if self.pduid_pdu.get(pdu_id)?.is_some() {
if self.pduid_pdu.get(pdu_id.as_bytes())?.is_some() {
self.pduid_pdu.insert(
pdu_id,
pdu_id.as_bytes(),
&serde_json::to_vec(pdu_json)
.expect("CanonicalJsonObject is always a valid"),
)?;

View file

@ -4,6 +4,7 @@ use crate::{
database::KeyValueDatabase,
service::{
self,
rooms::timeline::PduId,
sending::{Destination, RequestKey, SendingEventType},
},
services, utils, Error, Result,
@ -61,7 +62,7 @@ impl service::sending::Data for KeyValueDatabase {
for (destination, event) in requests {
let mut key = destination.get_prefix();
if let SendingEventType::Pdu(value) = &event {
key.extend_from_slice(value);
key.extend_from_slice(value.as_bytes());
} else {
key.extend_from_slice(
&services().globals.next_count()?.to_be_bytes(),
@ -202,7 +203,7 @@ fn parse_servercurrentevent(
Ok((
destination,
if value.is_empty() {
SendingEventType::Pdu(event.to_vec())
SendingEventType::Pdu(PduId::new(event.to_vec()))
} else {
SendingEventType::Edu(value)
},