diff --git a/src/service/rooms/pdu_metadata/data.rs b/src/service/rooms/pdu_metadata/data.rs index 58ee0c8fe..2f5cc38f7 100644 --- a/src/service/rooms/pdu_metadata/data.rs +++ b/src/service/rooms/pdu_metadata/data.rs @@ -60,6 +60,18 @@ impl Data { from: PduCount, dir: Direction, ) -> impl Stream + Send + '_ { + self.relation_counts(target, from, dir) + .wide_filter_map(move |count| self.get_relation_pdu(shortroomid, count)) + } + + /// The counts of the direct relations of `target`, in the order of `dir`, + /// starting after `from`. + pub(super) fn relation_counts( + &self, + target: u64, + from: PduCount, + dir: Direction, + ) -> impl Stream + Send + '_ { let start_count = match from { | PduCount::Normal(count) => count, | PduCount::Backfilled(_) => 0, @@ -77,13 +89,18 @@ impl Data { .map(|to_from| u64_from_u8(&to_from[8..16])) // The scan includes `from` itself, which should be skipped .ready_filter(move |&count| PduCount::Normal(count).is_after(from, dir)) - .wide_filter_map(move |count| async move { - let shorteventid = PduCount::Normal(count); - let pdu_id: RawPduId = PduId { shortroomid, shorteventid }.into(); - let pdu = self.services.timeline.get_pdu_from_id(&pdu_id).await.ok()?; + } - Some((count, pdu)) - }) + pub(super) async fn get_relation_pdu( + &self, + shortroomid: ShortRoomId, + count: u64, + ) -> Option<(u64, PduEvent)> { + let shorteventid = PduCount::Normal(count); + let pdu_id: RawPduId = PduId { shortroomid, shorteventid }.into(); + let pdu = self.services.timeline.get_pdu_from_id(&pdu_id).await.ok()?; + + Some((count, pdu)) } pub(super) fn mark_event_soft_failed(&self, event_id: &EventId) { diff --git a/src/service/rooms/pdu_metadata/mod.rs b/src/service/rooms/pdu_metadata/mod.rs index 2f6b11122..11d5d7568 100644 --- a/src/service/rooms/pdu_metadata/mod.rs +++ b/src/service/rooms/pdu_metadata/mod.rs @@ -8,7 +8,7 @@ use conduwuit::{ result::LogErr, utils::{IterStream, ReadyExt, stream::WidebandExt}, }; -use futures::{StreamExt, future::try_join}; +use futures::{Stream, StreamExt, future::try_join}; use ruma::{EventId, RoomId, UserId, api::Direction}; use self::data::Data; @@ -76,13 +76,9 @@ impl Service { // Direct relations are already in order if max_depth == 0 { + let counts = self.db.relation_counts(target, from, dir); return self - .db - .get_relations(room_id, target, from, dir) - .ready_filter(|(_, pdu)| matches(pdu)) - .take(limit) - .wide_then(|relation| self.finish_relation(user_id, relation)) - .collect() + .load_relations(user_id, room_id, counts, limit, &matches) .await; } @@ -95,45 +91,61 @@ impl Service { // that also replies to another reply in the thread, so it's only kept // and looked through the first time it's found let mut seen = HashSet::new(); - let mut relations = Vec::new(); + let mut counts = Vec::new(); let mut level: Vec<_> = self .db - .get_relations(room_id, target, scan_from, dir) + .relation_counts(target, scan_from, dir) .collect() .await; - seen.extend(level.iter().map(|(count, _)| *count)); + seen.extend(level.iter().copied()); for _ in 0..max_depth { let next: Vec> = level .iter() - .map(|(count, _)| *count) + .copied() .stream() .wide_then(|target| { self.db - .get_relations(room_id, target, scan_from, dir) + .relation_counts(target, scan_from, dir) .collect::>() }) .collect() .await; - relations.append(&mut level); + counts.append(&mut level); level = next .into_iter() .flatten() - .filter(|(count, _)| seen.insert(*count)) + .filter(|&count| seen.insert(count)) .collect(); } - relations.append(&mut level); + counts.append(&mut level); - relations - .retain(|(count, pdu)| PduCount::Normal(*count).is_after(from, dir) && matches(pdu)); - relations.sort_unstable_by(|(a, _), (b, _)| match dir { + counts.retain(|&count| PduCount::Normal(count).is_after(from, dir)); + counts.sort_unstable_by(|a, b| match dir { | Direction::Forward => a.cmp(b), | Direction::Backward => b.cmp(a), }); - relations.truncate(limit); - relations - .into_iter() - .stream() + self.load_relations(user_id, room_id, counts.into_iter().stream(), limit, &matches) + .await + } + + /// Loads the pdus of relations in the order of `counts` + async fn load_relations( + &self, + user_id: &UserId, + room_id: ShortRoomId, + counts: S, + limit: usize, + matches: &F, + ) -> Vec + where + S: Stream + Send, + F: Fn(&PduEvent) -> bool + Send + Sync, + { + counts + .wide_filter_map(|count| self.db.get_relation_pdu(room_id, count)) + .ready_filter(|(_, pdu)| matches(pdu)) + .take(limit) .wide_then(|relation| self.finish_relation(user_id, relation)) .collect() .await