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

@ -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"),
)?;