magmaus3/continuwuity
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

add database get_batch stream wrapper

Browse source 
Signed-off-by: Jason Volk <jason@zemos.net>
This commit is contained in:
Jason Volk 2024-11-15 03:44:04 +00:00
parent 9f7a4a012b
commit 14e3b242df
3 changed files with 28 additions and 19 deletions
Download patch file Download diff file Expand all files Collapse all files

12
src/service/rooms/auth_chain/mod.rs
View file

@ -6,7 +6,7 @@ use std::{
};
use conduit::{debug, debug_error, trace, utils::IterStream, validated, warn, Err, Result};
use futures::Stream;
use futures::{Stream, StreamExt};
use ruma::{EventId, RoomId};
use self::data::Data;
@ -69,15 +69,15 @@ impl Service {
const BUCKET: BTreeSet<(u64, &EventId)> = BTreeSet::new();
let started = std::time::Instant::now();
let mut buckets = [BUCKET; NUM_BUCKETS];
for (i, &short) in self
let mut starting_ids = self
.services
.short
.multi_get_or_create_shorteventid(starting_events)
.await
.iter()
.enumerate()
{
.boxed();
let mut buckets = [BUCKET; NUM_BUCKETS];
while let Some((i, short)) = starting_ids.next().await {
let bucket: usize = short.try_into()?;
let bucket: usize = validated!(bucket % NUM_BUCKETS);
buckets[bucket].insert((short, starting_events[i]));

13
src/service/rooms/short/mod.rs
View file

@ -3,6 +3,7 @@ use std::{mem::size_of_val, sync::Arc};
pub use conduit::pdu::{ShortEventId, ShortId, ShortRoomId};
use conduit::{err, implement, utils, Result};
use database::{Deserialized, Map};
use futures::{Stream, StreamExt};
use ruma::{events::StateEventType, EventId, RoomId};
use crate::{globals, Dep};
@ -71,11 +72,12 @@ pub async fn get_or_create_shorteventid(&self, event_id: &EventId) -> ShortEvent
}
#[implement(Service)]
pub async fn multi_get_or_create_shorteventid(&self, event_ids: &[&EventId]) -> Vec<ShortEventId> {
pub fn multi_get_or_create_shorteventid<'a>(
&'a self, event_ids: &'a [&EventId],
) -> impl Stream<Item = ShortEventId> + Send + 'a {
self.db
.eventid_shorteventid
.get_batch_blocking(event_ids.iter())
.into_iter()
.get_batch(event_ids.iter())
.enumerate()
.map(|(i, result)| match result {
Ok(ref short) => utils::u64_from_u8(short),
@ -95,7 +97,6 @@ pub async fn multi_get_or_create_shorteventid(&self, event_ids: &[&EventId]) ->
short
},
})
.collect()
}
#[implement(Service)]
@ -163,10 +164,10 @@ pub async fn multi_get_eventid_from_short(&self, shorteventid: &[ShortEventId])
self.db
.shorteventid_eventid
.get_batch_blocking(keys.iter())
.into_iter()
.get_batch(keys.iter())
.map(Deserialized::deserialized)
.collect()
.await
}
#[implement(Service)]