diff --git a/src/api/client/relations.rs b/src/api/client/relations.rs index 6d1bd4ffd..9a1e017dc 100644 --- a/src/api/client/relations.rs +++ b/src/api/client/relations.rs @@ -2,7 +2,7 @@ use axum::extract::State; use conduwuit::{ Err, Result, at, debug_warn, err, matrix::{Event, event::RelationTypeEqual, pdu::PduCount}, - utils::{IterStream, ReadyExt, result::FlatOk, stream::WidebandExt}, + utils::{IterStream, result::FlatOk, stream::WidebandExt}, }; use conduwuit_service::Services; use futures::StreamExt; @@ -147,28 +147,52 @@ async fn paginate_relations_with_filter( .unwrap_or(30) .min(100); - // Spec (v1.10) recommends depth of at least 3 - let depth: u8 = if recurse { 3 } else { 1 }; + // Spec (v1.10) recommends a depth of at least 3 + const RECURSION_DEPTH: u8 = 3; + const LEVELS_BELOW: u8 = RECURSION_DEPTH - 1; + let max_depth = if recurse { LEVELS_BELOW } else { 0 }; - let events: Vec<_> = services + // One more than the limit, to know whether there's another page + let mut events = services .rooms .pdu_metadata - .get_relations(sender_user, room_id, target, start, limit, depth, dir) - .await + .get_relations( + sender_user, + room_id, + target, + start, + limit.saturating_add(1), + max_depth, + dir, + |pdu| { + filter_event_type + .as_ref() + .is_none_or(|kind| kind == pdu.kind()) + && filter_rel_type + .as_ref() + .is_none_or(|rel_type| rel_type.relation_type_equal(pdu)) + }, + ) + .await; + + let past_to = to.and_then(|to| { + events + .iter() + .position(|(count, _)| !to.is_after(*count, dir)) + }); + if let Some(past_to) = past_to { + events.truncate(past_to); + } + + let has_more = past_to.is_none() && events.len() > limit; + events.truncate(limit); + + let next_batch = + has_more.then(|| events.last().map_or(start, |(count, _)| *count).to_string()); + + let chunk: Vec<_> = events .into_iter() - .filter(|(_, pdu)| { - filter_event_type - .as_ref() - .is_none_or(|kind| kind == pdu.kind()) - }) - .filter(|(_, pdu)| { - filter_rel_type - .as_ref() - .is_none_or(|rel_type| rel_type.relation_type_equal(pdu)) - }) .stream() - .ready_take_while(|(count, _)| Some(*count) != to) - .take(limit) .wide_filter_map(|item| visibility_filter(services, sender_user, item)) .wide_filter_map(|item| ignored_filter(services, item, sender_user)) .then(async |mut pdu| { @@ -182,32 +206,15 @@ async fn paginate_relations_with_filter( } pdu }) - .collect() - .await; - - // Determine if there are more events to fetch - let has_more = events.len() >= limit; - - let next_batch = if has_more { - match dir { - | Direction::Forward => events.last(), - | Direction::Backward => events.first(), - } - .map(|(count, _)| count.to_string()) - } else { - None - }; - - let chunk: Vec<_> = events - .into_iter() .map(at!(1)) .map(Event::into_format) - .collect(); + .collect() + .await; Ok(assign!(get_relating_events::v1::Response::new(chunk), { next_batch, prev_batch: from.map(Into::into), - recursion_depth: recurse.then_some(depth.into()), + recursion_depth: recurse.then_some(RECURSION_DEPTH.into()), })) } diff --git a/src/core/matrix/event.rs b/src/core/matrix/event.rs index 1b44b4ac1..18372e171 100644 --- a/src/core/matrix/event.rs +++ b/src/core/matrix/event.rs @@ -16,7 +16,12 @@ use ruma::{ use serde::Deserialize; use serde_json::{Value as JsonValue, value::RawValue as RawJsonValue}; -pub use self::{filter::Matches, id::*, relation::RelationTypeEqual, type_ext::TypeExt}; +pub use self::{ + filter::Matches, + id::*, + relation::{RelationTypeEqual, relation_type}, + type_ext::TypeExt, +}; use super::{pdu::Pdu, state_key::StateKey}; use crate::{Result, utils}; diff --git a/src/core/matrix/event/relation.rs b/src/core/matrix/event/relation.rs index 58324e863..7766433c2 100644 --- a/src/core/matrix/event/relation.rs +++ b/src/core/matrix/event/relation.rs @@ -8,7 +8,7 @@ pub trait RelationTypeEqual { } #[derive(Clone, Debug, Deserialize)] -struct ExtractRelatesToEventId { +struct ExtractRelatesToRelType { #[serde(rename = "m.relates_to")] relates_to: ExtractRelType, } @@ -18,11 +18,16 @@ struct ExtractRelType { rel_type: RelationType, } +/// The type of an event's relation, if it has one +pub fn relation_type(event: &E) -> Option { + event + .get_content() + .map(|c: ExtractRelatesToRelType| c.relates_to.rel_type) + .ok() +} + impl RelationTypeEqual for RelationType { fn relation_type_equal(&self, event: &E) -> bool { - event - .get_content() - .map(|c: ExtractRelatesToEventId| c.relates_to.rel_type) - .is_ok_and(|r| r == *self) + relation_type(event).is_some_and(|r| r == *self) } } diff --git a/src/core/matrix/pdu/count.rs b/src/core/matrix/pdu/count.rs index 7fb125745..8429db9ac 100644 --- a/src/core/matrix/pdu/count.rs +++ b/src/core/matrix/pdu/count.rs @@ -105,6 +105,16 @@ impl Count { } } + /// Whether this comes after `other` when going in the direction of `dir` + #[inline] + #[must_use] + pub fn is_after(self, other: Self, dir: Direction) -> bool { + match dir { + | Direction::Forward => self > other, + | Direction::Backward => self < other, + } + } + #[inline] #[must_use] pub fn saturating_add(self, add: u64) -> Self { diff --git a/src/service/rooms/pdu_metadata/bundled_aggregations.rs b/src/service/rooms/pdu_metadata/bundled_aggregations.rs index b6b1cb930..bc7204e8c 100644 --- a/src/service/rooms/pdu_metadata/bundled_aggregations.rs +++ b/src/service/rooms/pdu_metadata/bundled_aggregations.rs @@ -1,14 +1,16 @@ -use conduwuit::{Event, PduEvent, Result, err}; +use conduwuit::{ + Event, PduEvent, Result, err, + matrix::{PduCount, event::relation_type}, + utils::ReadyExt, +}; use ruma::{ UserId, api::Direction, - events::relation::{BundledMessageLikeRelations, BundledReference, ReferenceChunk}, + events::relation::{ + BundledMessageLikeRelations, BundledReference, ReferenceChunk, RelationType, + }, }; -use crate::rooms::timeline::PdusIterItem; - -const MAX_BUNDLED_RELATIONS: usize = 50; - impl super::Service { /// Gets bundled aggregations for an event according to the Matrix /// specification. @@ -26,52 +28,24 @@ impl super::Service { return Ok(None); } - let relations = self - .get_relations( - user_id, - &pdu.room_id_or_hash(), - pdu.event_id(), - conduwuit::PduCount::max(), - MAX_BUNDLED_RELATIONS, - 0, - Direction::Backward, - ) + let Some((room_id, target)) = self + .relation_target(&pdu.room_id_or_hash(), pdu.event_id()) + .await + else { + return Ok(None); + }; + + let mut replace_events = Vec::new(); + let mut reference_events = Vec::new(); + self.db + .get_relations(room_id, target, PduCount::max(), Direction::Backward) + .ready_for_each(|(_, relation)| match relation_type(&relation) { + | Some(RelationType::Replacement) => replace_events.push(relation), + | Some(RelationType::Reference) => reference_events.push(relation), + | _ => {}, + }) .await; - // The relations database code still handles the basic unsigned data - // We don't want to recursively fetch relations - - if relations.is_empty() { - return Ok(None); - } - - // Partition relations by type - let (replace_events, reference_events): (Vec<_>, Vec<_>) = relations - .iter() - .filter_map(|relation| { - let pdu = &relation.1; - let content = pdu.get_content_as_value(); - - content - .get("m.relates_to") - .and_then(|relates_to| relates_to.get("rel_type")) - .and_then(|rel_type| rel_type.as_str()) - .and_then(|rel_type_str| match rel_type_str { - | "m.replace" => Some(RelationType::Replace(relation)), - | "m.reference" => Some(RelationType::Reference(relation)), - | _ => None, /* Ignore other relation types (threads are in DB but not - * handled here) */ - }) - }) - .fold((Vec::new(), Vec::new()), |(mut replaces, mut references), rel_type| { - match rel_type { - | RelationType::Replace(r) => replaces.push(r), - | RelationType::Reference(r) => references.push(r), - } - (replaces, references) - }); - - // If no relations to bundle, return None if replace_events.is_empty() && reference_events.is_empty() { return Ok(None); } @@ -85,7 +59,9 @@ impl super::Service { .find_most_recent_valid_replacement(user_id, pdu, &replace_events) .await? { - bundled.replace = Some(Self::serialize_replacement(replacement)?); + bundled.replace = Some(Self::serialize_replacement( + &self.with_unsigned(user_id, replacement.clone()).await, + )?); } } @@ -93,7 +69,7 @@ impl super::Service { if !reference_events.is_empty() { let reference_chunk: Vec<_> = reference_events .into_iter() - .map(|relation| BundledReference::new(relation.1.event_id().to_owned())) + .map(|relation| BundledReference::new(relation.event_id().to_owned())) .collect(); if !reference_chunk.is_empty() { @@ -123,14 +99,12 @@ impl super::Service { &self, user_id: &UserId, original_event: &PduEvent, - replacement_events: &[&'a PdusIterItem], + replacement_events: &'a [PduEvent], ) -> Result> { // Filter valid replacements and find the maximum in a single pass let mut result: Option<&PduEvent> = None; - for relation in replacement_events { - let pdu = &relation.1; - + for pdu in replacement_events { // Validate replacement if !Self::is_valid_replacement_event(original_event, pdu).await? { continue; @@ -245,12 +219,7 @@ impl super::Service { /// Check if an event is itself a replacement #[inline] fn is_replacement_event(event: &PduEvent) -> bool { - event - .get_content_as_value() - .get("m.relates_to") - .and_then(|relates_to| relates_to.get("rel_type")) - .and_then(|rel_type| rel_type.as_str()) - .is_some_and(|rel_type| rel_type == "m.replace") + relation_type(event) == Some(RelationType::Replacement) } /// Check if event has m.new_content or is encrypted (where m.new_content @@ -262,12 +231,6 @@ impl super::Service { } } -/// Helper enum for partitioning relations -enum RelationType<'a> { - Replace(&'a PdusIterItem), - Reference(&'a PdusIterItem), -} - #[cfg(test)] mod tests { use conduwuit_core::pdu::{EventHash, PduEvent}; diff --git a/src/service/rooms/pdu_metadata/data.rs b/src/service/rooms/pdu_metadata/data.rs index 4920df5e7..d85c226cb 100644 --- a/src/service/rooms/pdu_metadata/data.rs +++ b/src/service/rooms/pdu_metadata/data.rs @@ -2,8 +2,7 @@ use std::{mem::size_of, sync::Arc}; use conduwuit::{ arrayvec::ArrayVec, - matrix::PduCount, - result::LogErr, + matrix::{PduCount, PduEvent}, utils::{ ReadyExt, stream::{TryIgnore, WidebandExt}, @@ -12,14 +11,14 @@ use conduwuit::{ }; use database::Map; use futures::{Stream, StreamExt}; -use ruma::{EventId, UserId, api::Direction}; +use ruma::{EventId, api::Direction}; use crate::{ Dep, rooms::{ self, - short::{ShortEventId, ShortRoomId}, - timeline::{PduId, PdusIterItem, RawPduId}, + short::ShortRoomId, + timeline::{PduId, RawPduId}, }, }; @@ -56,51 +55,34 @@ impl Data { pub(super) fn get_relations<'a>( &'a self, - user_id: &'a UserId, shortroomid: ShortRoomId, - target: ShortEventId, + target: u64, from: PduCount, dir: Direction, - ) -> impl Stream + Send + 'a { - // Query from exact position then filter excludes it (saturating_inc - // could skip events at min/max boundaries) - let from_unsigned = from.into_unsigned(); - let mut current = ArrayVec::::new(); - current.extend(target.to_be_bytes()); - current.extend(from_unsigned.to_be_bytes()); - let current = current.as_slice(); + ) -> impl Stream + Send + 'a { + let start_count = match from { + | PduCount::Normal(count) => count, + | PduCount::Backfilled(_) => 0, + }; + let mut start = ArrayVec::::new(); + start.extend(target.to_be_bytes()); + start.extend(start_count.to_be_bytes()); + let start = start.as_slice(); match dir { - | Direction::Forward => self.tofrom_relation.raw_keys_from(current).boxed(), - | Direction::Backward => self.tofrom_relation.rev_raw_keys_from(current).boxed(), + | Direction::Forward => self.tofrom_relation.raw_keys_from(start).boxed(), + | Direction::Backward => self.tofrom_relation.rev_raw_keys_from(start).boxed(), } .ignore_err() .ready_take_while(move |key| key.starts_with(&target.to_be_bytes())) .map(|to_from| u64_from_u8(&to_from[8..16])) - .map(PduCount::from_unsigned) - .ready_filter(move |count| { - if from == PduCount::min() || from == PduCount::max() { - true - } else { - let count_unsigned = count.into_unsigned(); - match dir { - | Direction::Forward => count_unsigned > from_unsigned, - | Direction::Backward => count_unsigned < from_unsigned, - } - } - }) - .wide_filter_map(move |shorteventid| async move { + // 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()?; - let mut pdu = self.services.timeline.get_pdu_from_id(&pdu_id).await.ok()?; - - let ctx = self - .services - .timeline - .get_unsigned_context(&pdu, Some(user_id)) - .await; - pdu.set_unsigned(Some(ctx)).log_err().ok(); - - Some((shorteventid, pdu)) + Some((count, pdu)) }) } diff --git a/src/service/rooms/pdu_metadata/mod.rs b/src/service/rooms/pdu_metadata/mod.rs index 33c3d968b..2f6b11122 100644 --- a/src/service/rooms/pdu_metadata/mod.rs +++ b/src/service/rooms/pdu_metadata/mod.rs @@ -1,15 +1,20 @@ mod bundled_aggregations; mod data; -use std::sync::Arc; +use std::{collections::HashSet, sync::Arc}; -use conduwuit::{Result, matrix::PduCount}; +use conduwuit::{ + Result, + matrix::{PduCount, PduEvent}, + result::LogErr, + utils::{IterStream, ReadyExt, stream::WidebandExt}, +}; use futures::{StreamExt, future::try_join}; use ruma::{EventId, RoomId, UserId, api::Direction}; use self::data::Data; use crate::{ Dep, - rooms::{self, timeline::PdusIterItem}, + rooms::{self, short::ShortRoomId, timeline::PdusIterItem}, }; pub struct Service { @@ -51,7 +56,7 @@ impl Service { } #[allow(clippy::too_many_arguments)] - pub async fn get_relations<'a>( + pub async fn get_relations<'a, F>( &'a self, user_id: &'a UserId, room_id: &'a RoomId, @@ -60,59 +65,113 @@ impl Service { limit: usize, max_depth: u8, dir: Direction, - ) -> Vec { - let room_id = self.services.short.get_shortroomid(room_id); - - let target = self.services.timeline.get_pdu_count(target); - - let Ok((room_id, target)) = try_join(room_id, target).await else { + matches: F, + ) -> Vec + where + F: Fn(&PduEvent) -> bool + Send + Sync, + { + let Some((room_id, target)) = self.relation_target(room_id, target).await else { return Vec::new(); }; - let target = match target { - | PduCount::Normal(c) => c, - // TODO: Support backfilled relations - | _ => 0, // This will result in an empty iterator - }; - - let mut pdus: Vec<_> = self - .db - .get_relations(user_id, room_id, target, from, dir) - .collect() - .await; - - let mut stack: Vec<_> = pdus - .iter() - .filter(|_| max_depth > 0) - .map(|pdu| (pdu.clone(), 1)) - .collect(); - - 'limit: while let Some(stack_pdu) = stack.pop() { - let target = match stack_pdu.0.0 { - | PduCount::Normal(c) => c, - // TODO: Support backfilled relations - | PduCount::Backfilled(_) => 0, // This will result in an empty iterator - }; - - let relations: Vec<_> = self + // Direct relations are already in order + if max_depth == 0 { + return self .db - .get_relations(user_id, room_id, target, from, dir) + .get_relations(room_id, target, from, dir) + .ready_filter(|(_, pdu)| matches(pdu)) + .take(limit) + .wide_then(|relation| self.finish_relation(user_id, relation)) .collect() .await; - - for relation in relations { - if stack_pdu.1 < max_depth { - stack.push((relation.clone(), stack_pdu.1.saturating_add(1))); - } - - pdus.push(relation); - if pdus.len() >= limit { - break 'limit; - } - } } - pdus + let scan_from = match dir { + | Direction::Forward => PduCount::min(), + | Direction::Backward => from, + }; + + // An event can relate to two others in the tree, like a thread reply + // 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 level: Vec<_> = self + .db + .get_relations(room_id, target, scan_from, dir) + .collect() + .await; + seen.extend(level.iter().map(|(count, _)| *count)); + for _ in 0..max_depth { + let next: Vec> = level + .iter() + .map(|(count, _)| *count) + .stream() + .wide_then(|target| { + self.db + .get_relations(room_id, target, scan_from, dir) + .collect::>() + }) + .collect() + .await; + relations.append(&mut level); + level = next + .into_iter() + .flatten() + .filter(|(count, _)| seen.insert(*count)) + .collect(); + } + relations.append(&mut level); + + relations + .retain(|(count, pdu)| PduCount::Normal(*count).is_after(from, dir) && matches(pdu)); + relations.sort_unstable_by(|(a, _), (b, _)| match dir { + | Direction::Forward => a.cmp(b), + | Direction::Backward => b.cmp(a), + }); + relations.truncate(limit); + + relations + .into_iter() + .stream() + .wide_then(|relation| self.finish_relation(user_id, relation)) + .collect() + .await + } + + /// The short room ID and count of `target` + async fn relation_target( + &self, + room_id: &RoomId, + target: &EventId, + ) -> Option<(ShortRoomId, u64)> { + let room_id = self.services.short.get_shortroomid(room_id); + let target = self.services.timeline.get_pdu_count(target); + + match try_join(room_id, target).await { + | Ok((room_id, PduCount::Normal(target))) => Some((room_id, target)), + // TODO: Support backfilled relations + | _ => None, + } + } + + async fn finish_relation( + &self, + user_id: &UserId, + (count, pdu): (u64, PduEvent), + ) -> PdusIterItem { + (PduCount::Normal(count), self.with_unsigned(user_id, pdu).await) + } + + async fn with_unsigned(&self, user_id: &UserId, mut pdu: PduEvent) -> PduEvent { + let ctx = self + .services + .timeline + .get_unsigned_context(&pdu, Some(user_id)) + .await; + pdu.set_unsigned(Some(ctx)).log_err().ok(); + + pdu } #[inline]