refactor: Improve code style for bundled aggregations

This commit is contained in:
Jade Ellis
2025-12-21 17:12:34 +00:00
parent 8d3e4eba99
commit 8e33f9a7d0
@@ -33,6 +33,7 @@ impl super::Service {
Direction::Backward, Direction::Backward,
) )
.await; .await;
// The relations database code still handles the basic unsigned data // The relations database code still handles the basic unsigned data
// We don't want to recursively fetch relations // We don't want to recursively fetch relations
@@ -43,65 +44,58 @@ impl super::Service {
return Ok(None); return Ok(None);
} }
// Get the original event for validation of replacement events // Partition relations by type
let original_event = self.services.timeline.get_pdu(event_id).await?; let (replace_events, reference_events): (Vec<_>, Vec<_>) = relations
.iter()
let mut replace_events = Vec::with_capacity(relations.len()); .filter_map(|relation| {
let mut reference_events = Vec::with_capacity(relations.len());
for relation in &relations {
let pdu = &relation.1; let pdu = &relation.1;
let content = pdu.get_content_as_value(); let content = pdu.get_content_as_value();
if let Some(relates_to) = content.get("m.relates_to") {
// We don't check that the event relates back, because we assume the database is content
// good. .get("m.relates_to")
if let Some(rel_type) = relates_to.get("rel_type") { .and_then(|relates_to| relates_to.get("rel_type"))
match rel_type.as_str() { .and_then(|rel_type| rel_type.as_str())
| Some("m.replace") => { .and_then(|rel_type_str| match rel_type_str {
// Only consider valid replacements | "m.replace" => Some(RelationType::Replace(relation)),
if Self::is_valid_replacement_event(&original_event, pdu).await? { | "m.reference" => Some(RelationType::Reference(relation)),
replace_events.push(relation); | _ => None, /* Ignore other relation types (threads are in DB but not
} * handled here) */
}, })
| Some("m.reference") => { })
reference_events.push(relation); .fold((Vec::new(), Vec::new()), |(mut replaces, mut references), rel_type| {
}, match rel_type {
| _ => { | RelationType::Replace(r) => replaces.push(r),
// Ignore other relation types for now | RelationType::Reference(r) => references.push(r),
// Threads are in the database but not handled here
// Other types are not specified AFAICT.
},
}
}
}
} }
(replaces, references)
});
// If no relations to bundle, return None // If no relations to bundle, return None
if replace_events.is_empty() && reference_events.is_empty() { if replace_events.is_empty() && reference_events.is_empty() {
return Ok(None); return Ok(None);
} }
let mut bundled = BundledMessageLikeRelations::new(); let mut bundled = BundledMessageLikeRelations::<Box<serde_json::value::RawValue>>::new();
// Handle m.replace relations - find the most recent one // Handle m.replace relations - find the most recent valid one (lazy load
// original event)
if !replace_events.is_empty() { if !replace_events.is_empty() {
let most_recent_replacement = Self::find_most_recent_replacement(&replace_events)?; let original_event = self.services.timeline.get_pdu(event_id).await?;
// Convert the replacement event to the bundled format if let Some(replacement) =
if let Some(replacement_pdu) = most_recent_replacement { Self::find_most_recent_valid_replacement(&original_event, &replace_events).await?
// According to the Matrix spec, we should include the full event as raw JSON {
let replacement_json = serde_json::to_string(replacement_pdu) bundled.replace = Some(Self::serialize_replacement(replacement)?);
.map_err(|e| err!(Database("Failed to serialize replacement event: {e}")))?;
let raw_value = serde_json::value::RawValue::from_string(replacement_json)
.map_err(|e| err!(Database("Failed to create RawValue: {e}")))?;
bundled.replace = Some(Box::new(raw_value));
} }
} }
// Handle m.reference relations - collect event IDs // Handle m.reference relations - collect event IDs
if !reference_events.is_empty() { if !reference_events.is_empty() {
let reference_chunk = Self::build_reference_chunk(&reference_events)?; let reference_chunk: Vec<_> = reference_events
.into_iter()
.map(|relation| BundledReference::new(relation.1.event_id().to_owned()))
.collect();
if !reference_chunk.is_empty() { if !reference_chunk.is_empty() {
bundled.reference = Some(Box::new(ReferenceChunk::new(reference_chunk))); bundled.reference = Some(Box::new(ReferenceChunk::new(reference_chunk)));
} }
@@ -112,67 +106,48 @@ impl super::Service {
Ok(Some(bundled)) Ok(Some(bundled))
} }
/// Build reference chunk for m.reference bundled aggregations /// Serialize a replacement event to the bundled format
fn build_reference_chunk( fn serialize_replacement(pdu: &PduEvent) -> Result<Box<Box<serde_json::value::RawValue>>> {
reference_events: &[&PdusIterItem], let replacement_json = serde_json::to_string(pdu)
) -> Result<Vec<BundledReference>> { .map_err(|e| err!(Database("Failed to serialize replacement event: {e}")))?;
let mut chunk = Vec::with_capacity(reference_events.len());
for relation in reference_events { let raw_value = serde_json::value::RawValue::from_string(replacement_json)
let pdu = &relation.1; .map_err(|e| err!(Database("Failed to create RawValue: {e}")))?;
let reference_entry = BundledReference::new(pdu.event_id().to_owned()); Ok(Box::new(raw_value))
chunk.push(reference_entry);
} }
// Don't sort, order is unspecified /// Find the most recent valid replacement event based on origin_server_ts
/// and lexicographic event_id ordering
Ok(chunk) async fn find_most_recent_valid_replacement<'a>(
} original_event: &PduEvent,
replacement_events: &[&'a PdusIterItem],
/// Find the most recent replacement event based on origin_server_ts and
/// lexicographic event_id ordering
fn find_most_recent_replacement<'a>(
replacement_events: &'a [&'a PdusIterItem],
) -> Result<Option<&'a PduEvent>> { ) -> Result<Option<&'a PduEvent>> {
if replacement_events.is_empty() { // Filter valid replacements and find the maximum in a single pass
return Ok(None); let mut result: Option<&PduEvent> = None;
}
let mut most_recent: Option<&PduEvent> = None;
// Jank, is there a better way to do this?
for relation in replacement_events { for relation in replacement_events {
let pdu = &relation.1; let pdu = &relation.1;
match most_recent { // Validate replacement
| None => { if !Self::is_valid_replacement_event(original_event, pdu).await? {
most_recent = Some(pdu); continue;
},
| Some(current_most_recent) => {
// Compare by origin_server_ts first
match pdu
.origin_server_ts()
.cmp(&current_most_recent.origin_server_ts())
{
| std::cmp::Ordering::Greater => {
most_recent = Some(pdu);
},
| std::cmp::Ordering::Equal => {
// If timestamps are equal, use lexicographic ordering of event_id
if pdu.event_id() > current_most_recent.event_id() {
most_recent = Some(pdu);
}
},
| std::cmp::Ordering::Less => {
// Keep current most recent
},
}
},
}
} }
Ok(most_recent) result = Some(match result {
| None => pdu,
| Some(current) => {
// Compare by origin_server_ts first, then event_id lexicographically
match pdu.origin_server_ts().cmp(&current.origin_server_ts()) {
| std::cmp::Ordering::Greater => pdu,
| std::cmp::Ordering::Equal if pdu.event_id() > current.event_id() => pdu,
| _ => current,
}
},
});
}
Ok(result)
} }
/// Adds bundled aggregations to a PDU's unsigned field /// Adds bundled aggregations to a PDU's unsigned field
@@ -200,8 +175,7 @@ impl super::Service {
Ok(()) Ok(())
} }
/// Helper method to add bundled aggregations to a PDU's unsigned /// Helper method to add bundled aggregations to a PDU's unsigned field
/// field
fn add_bundled_aggregations_to_unsigned( fn add_bundled_aggregations_to_unsigned(
pdu: &mut PduEvent, pdu: &mut PduEvent,
aggregations_json: serde_json::Value, aggregations_json: serde_json::Value,
@@ -225,9 +199,7 @@ impl super::Service {
.ok_or_else(|| err!(Database("m.relations is not an object")))?; .ok_or_else(|| err!(Database("m.relations is not an object")))?;
if let JsonValue::Object(aggregations_map) = aggregations_json { if let JsonValue::Object(aggregations_map) = aggregations_json {
for (rel_type, aggregation) in aggregations_map { relations.extend(aggregations_map);
relations.insert(rel_type, aggregation);
}
} }
pdu.unsigned = Some(to_raw_value(&unsigned)?); pdu.unsigned = Some(to_raw_value(&unsigned)?);
@@ -242,50 +214,49 @@ impl super::Service {
original_event: &PduEvent, original_event: &PduEvent,
replacement_event: &PduEvent, replacement_event: &PduEvent,
) -> Result<bool> { ) -> Result<bool> {
Ok(
// 1. Same room_id // 1. Same room_id
if original_event.room_id() != replacement_event.room_id() { original_event.room_id() == replacement_event.room_id()
return Ok(false);
}
// 2. Same sender // 2. Same sender
if original_event.sender() != replacement_event.sender() { && original_event.sender() == replacement_event.sender()
return Ok(false);
}
// 3. Same type // 3. Same type
if original_event.event_type() != replacement_event.event_type() { && original_event.event_type() == replacement_event.event_type()
return Ok(false);
}
// 4. Neither event should have a state_key property // 4. Neither event should have a state_key property
if original_event.state_key().is_some() || replacement_event.state_key().is_some() { && original_event.state_key().is_none()
return Ok(false); && replacement_event.state_key().is_none()
}
// 5. Original event must not have rel_type of m.replace // 5. Original event must not have rel_type of m.replace
let original_content = original_event.get_content_as_value(); && !Self::is_replacement_event(original_event)
if let Some(relates_to) = original_content.get("m.relates_to") { // 6. Replacement event must have m.new_content property (skip for encrypted)
if let Some(rel_type) = relates_to.get("rel_type") { && Self::has_new_content_or_encrypted(replacement_event),
if rel_type.as_str() == Some("m.replace") { )
return Ok(false);
}
}
} }
// 6. Replacement event must have m.new_content property /// Check if an event is itself a replacement
// Skip this check for encrypted events, as m.new_content would be inside the #[inline]
// encrypted payload fn is_replacement_event(event: &PduEvent) -> bool {
if replacement_event.event_type() != &ruma::events::TimelineEventType::RoomEncrypted { event
let replacement_content = replacement_event.get_content_as_value(); .get_content_as_value()
if replacement_content.get("m.new_content").is_none() { .get("m.relates_to")
return Ok(false); .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")
} }
Ok(true) /// Check if event has m.new_content or is encrypted (where m.new_content
/// would be in the encrypted payload)
#[inline]
fn has_new_content_or_encrypted(event: &PduEvent) -> bool {
event.event_type() == &ruma::events::TimelineEventType::RoomEncrypted
|| event.get_content_as_value().get("m.new_content").is_some()
} }
} }
/// Helper enum for partitioning relations
enum RelationType<'a> {
Replace(&'a PdusIterItem),
Reference(&'a PdusIterItem),
}
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use conduwuit_core::pdu::{EventHash, PduEvent}; use conduwuit_core::pdu::{EventHash, PduEvent};