Skip to content

Commit

Permalink
enrich state iteration interface
Browse files Browse the repository at this point in the history 
Signed-off-by: Jason Volk <[email protected]>
  • Loading branch information
@jevolk
jevolk committed Nov 19, 2024
1 parent 7680d1b commit 411c600
Show file tree
Hide file tree
Showing 4 changed files with 115 additions and 105 deletions.
32 changes: 17 additions & 15 deletions src/service/rooms/state/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ use crate::{
globals, rooms,
rooms::{
short::{ShortEventId, ShortStateHash},
state_compressor::CompressedStateEvent,
state_compressor::{parse_compressed_state_event, CompressedStateEvent},
},
Dep,
};
Expand Down Expand Up @@ -98,12 +98,12 @@ impl Service {
_statediffremoved: Arc<HashSet<CompressedStateEvent>>,
state_lock: &RoomMutexGuard, // Take mutex guard to make sure users get the room state mutex
) -> Result {
let event_ids = statediffnew.iter().stream().filter_map(|new| {
self.services
.state_compressor
.parse_compressed_state_event(*new)
.map_ok_or_else(|_| None, |(_, event_id)| Some(event_id))
});
let event_ids = statediffnew
.iter()
.stream()
.map(|&new| parse_compressed_state_event(new).1)
.then(|shorteventid| self.services.short.get_eventid_from_short(shorteventid))
.ignore_err();

pin_mut!(event_ids);
while let Some(event_id) = event_ids.next().await {
Expand Down Expand Up @@ -428,17 +428,19 @@ impl Service {
.full_state;

let mut ret = HashMap::new();
for compressed in full_state.iter() {
let Ok((shortstatekey, event_id)) = self
.services
.state_compressor
.parse_compressed_state_event(*compressed)
.await
else {
for &compressed in full_state.iter() {
let (shortstatekey, shorteventid) = parse_compressed_state_event(compressed);

let Some((ty, state_key)) = sauthevents.remove(&shortstatekey) else {
continue;
};

let Some((ty, state_key)) = sauthevents.remove(&shortstatekey) else {
let Ok(event_id) = self
.services
.short
.get_eventid_from_short(shorteventid)
.await
else {
continue;
};

Expand Down
151 changes: 80 additions & 71 deletions src/service/rooms/state_accessor/data.rs
View file
Original file line number Diff line number Diff line change
@@ -1,11 +1,22 @@
use std::{collections::HashMap, sync::Arc};

use conduit::{err, PduEvent, Result};
use conduit::{
err,
utils::{future::TryExtExt, IterStream},
PduEvent, Result,
};
use database::{Deserialized, Map};
use futures::TryFutureExt;
use futures::{StreamExt, TryFutureExt};
use ruma::{events::StateEventType, EventId, RoomId};

use crate::{rooms, rooms::short::ShortStateHash, Dep};
use crate::{
rooms,
rooms::{
short::{ShortEventId, ShortStateHash, ShortStateKey},
state_compressor::parse_compressed_state_event,
},
Dep,
};

pub(super) struct Data {
eventid_shorteventid: Arc<Map>,
Expand Down Expand Up @@ -35,73 +46,67 @@ impl Data {
}
}

#[allow(unused_qualifications)] // async traits
pub(super) async fn state_full_ids(&self, shortstatehash: ShortStateHash) -> Result<HashMap<u64, Arc<EventId>>> {
let full_state = self
.services
.state_compressor
.load_shortstatehash_info(shortstatehash)
.await
.map_err(|e| err!(Database("Missing state IDs: {e}")))?
.pop()
.expect("there is always one layer")
.full_state;
pub(super) async fn state_full(
&self, shortstatehash: ShortStateHash,
) -> Result<HashMap<(StateEventType, String), Arc<PduEvent>>> {
Ok(self
.state_full_pdus(shortstatehash)
.await?
.into_iter()
.filter_map(|pdu| Some(((pdu.kind.to_string().into(), pdu.state_key.clone()?), pdu)))
.collect())
}

let mut result = HashMap::new();
let mut i: u8 = 0;
for compressed in full_state.iter() {
let parsed = self
.services
.state_compressor
.parse_compressed_state_event(*compressed)
.await?;

result.insert(parsed.0, parsed.1);

i = i.wrapping_add(1);
if i % 100 == 0 {
tokio::task::yield_now().await;
}
}
pub(super) async fn state_full_pdus(&self, shortstatehash: ShortStateHash) -> Result<Vec<Arc<PduEvent>>> {
Ok(self
.state_full_shortids(shortstatehash)
.await?
.iter()
.stream()
.filter_map(|(_, shorteventid)| {
self.services
.short
.get_eventid_from_short(*shorteventid)
.ok()
})
.filter_map(|eventid| async move { self.services.timeline.get_pdu(&eventid).await.ok() })
.collect()
.await)
}

Ok(result)
pub(super) async fn state_full_ids(&self, shortstatehash: ShortStateHash) -> Result<HashMap<u64, Arc<EventId>>> {
Ok(self
.state_full_shortids(shortstatehash)
.await?
.iter()
.stream()
.filter_map(|(shortstatekey, shorteventid)| {
self.services
.short
.get_eventid_from_short(*shorteventid)
.map_ok(move |eventid| (*shortstatekey, eventid))
.ok()
})
.collect()
.await)
}

#[allow(unused_qualifications)] // async traits
pub(super) async fn state_full(
pub(super) async fn state_full_shortids(
&self, shortstatehash: ShortStateHash,
) -> Result<HashMap<(StateEventType, String), Arc<PduEvent>>> {
let full_state = self
) -> Result<Vec<(ShortStateKey, ShortEventId)>> {
Ok(self
.services
.state_compressor
.load_shortstatehash_info(shortstatehash)
.await?
.await
.map_err(|e| err!(Database("Missing state IDs: {e}")))?
.pop()
.expect("there is always one layer")
.full_state;

let mut result = HashMap::new();
let mut i: u8 = 0;
for compressed in full_state.iter() {
let (_, eventid) = self
.services
.state_compressor
.parse_compressed_state_event(*compressed)
.await?;

if let Ok(pdu) = self.services.timeline.get_pdu(&eventid).await {
if let Some(state_key) = pdu.state_key.as_ref() {
result.insert((pdu.kind.to_string().into(), state_key.clone()), pdu);
}
}

i = i.wrapping_add(1);
if i % 100 == 0 {
tokio::task::yield_now().await;
}
}

Ok(result)
.full_state
.iter()
.copied()
.map(parse_compressed_state_event)
.collect())
}

/// Returns a single PDU from `room_id` with key (`event_type`,`state_key`).
Expand Down Expand Up @@ -130,18 +135,11 @@ impl Data {
.find(|bytes| bytes.starts_with(&shortstatekey.to_be_bytes()))
.ok_or(err!(Database("No shortstatekey in compressed state")))?;

let (_, shorteventid) = parse_compressed_state_event(*compressed);

self.services
.state_compressor
.parse_compressed_state_event(*compressed)
.map_ok(|(_, id)| id)
.map_err(|e| {
err!(Database(error!(
?event_type,
?state_key,
?shortstatekey,
"Failed to parse compressed: {e:?}"
)))
})
.short
.get_eventid_from_short(shorteventid)
.await
}

Expand Down Expand Up @@ -176,6 +174,17 @@ impl Data {
.await
}

/// Returns the full room state's pdus.
#[allow(unused_qualifications)] // async traits
pub(super) async fn room_state_full_pdus(&self, room_id: &RoomId) -> Result<Vec<Arc<PduEvent>>> {
self.services
.state
.get_room_shortstatehash(room_id)
.and_then(|shortstatehash| self.state_full_pdus(shortstatehash))
.map_err(|e| err!(Database("Missing state pdus for {room_id:?}: {e:?}")))
.await
}

/// Returns a single PDU from `room_id` with key (`event_type`,`state_key`).
pub(super) async fn room_state_get_id(
&self, room_id: &RoomId, event_type: &StateEventType, state_key: &str,
Expand Down
6 changes: 6 additions & 0 deletions src/service/rooms/state_accessor/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -301,6 +301,12 @@ impl Service {
self.db.room_state_full(room_id).await
}

/// Returns the full room state pdus
#[tracing::instrument(skip(self), level = "debug")]
pub async fn room_state_full_pdus(&self, room_id: &RoomId) -> Result<Vec<Arc<PduEvent>>> {
self.db.room_state_full_pdus(room_id).await
}

/// Returns a single PDU from `room_id` with key (`event_type`,
/// `state_key`).
#[tracing::instrument(skip(self), level = "debug")]
Expand Down
31 changes: 12 additions & 19 deletions src/service/rooms/state_compressor/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ use ruma::{EventId, RoomId};

use crate::{
rooms,
rooms::short::{ShortId, ShortStateHash, ShortStateKey},
rooms::short::{ShortEventId, ShortId, ShortStateHash, ShortStateKey},
Dep,
};

Expand Down Expand Up @@ -196,24 +196,6 @@ impl Service {
.expect("failed to create CompressedStateEvent")
}

/// Returns shortstatekey, event id
#[inline]
pub async fn parse_compressed_state_event(
&self, compressed_event: CompressedStateEvent,
) -> Result<(ShortStateKey, Arc<EventId>)> {
use utils::u64_from_u8;

let shortstatekey = u64_from_u8(&compressed_event[0..size_of::<ShortStateKey>()]);
let shorteventid = u64_from_u8(&compressed_event[size_of::<ShortStateKey>()..]);
let event_id = self
.services
.short
.get_eventid_from_short(shorteventid)
.await?;

Ok((shortstatekey, event_id))
}

/// Creates a new shortstatehash that often is just a diff to an already
/// existing shortstatehash and therefore very efficient.
///
Expand Down Expand Up @@ -488,6 +470,17 @@ impl Service {
}
}

#[inline]
#[must_use]
pub fn parse_compressed_state_event(compressed_event: CompressedStateEvent) -> (ShortStateKey, ShortEventId) {
use utils::u64_from_u8;

let shortstatekey = u64_from_u8(&compressed_event[0..size_of::<ShortStateKey>()]);
let shorteventid = u64_from_u8(&compressed_event[size_of::<ShortStateKey>()..]);

(shortstatekey, shorteventid)
}

#[inline]
fn compressed_state_size(compressed_state: &CompressedState) -> usize {
compressed_state
Expand Down

0 comments on commit 411c600

Please sign in to comment.