From 4b0fd01aefa0e86d7b89febb712b5d4dc17e495e Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Mon, 15 Jun 2026 14:47:51 +0100 Subject: [PATCH] Port the BundledAggregations class to Rust Move `BundledAggregations` and `_ThreadAggregation` from the Python `RelationsHandler` into frozen Rust pyclasses exposed from `synapse.synapse_rust.events`. They are re-exported from `synapse.handlers.relations` so existing call sites are unaffected. The referenced events are stored by value (the `Event` pyclass now derives a cheap, Arc-sharing `Clone`). The `references` field uses the `JsonObject` type, which gains a `FromPyObject` impl so it can be built from any Python mapping in place of a plain dict. As the pyclass is immutable, `get_bundled_aggregations` is reworked to collect each kind of aggregation and build one instance per event. Co-Authored-By: Claude Opus 4.8 --- rust/src/events/formats/mod.rs | 12 ++-- rust/src/events/json_object.rs | 82 +++++++++++++++++++------- rust/src/events/mod.rs | 9 ++- rust/src/events/relations.rs | 101 ++++++++++++++++++++++++++++++++ synapse/handlers/relations.py | 72 +++++++++++------------ synapse/synapse_rust/events.pyi | 41 +++++++++++++ 6 files changed, 254 insertions(+), 63 deletions(-) create mode 100644 rust/src/events/relations.rs diff --git a/rust/src/events/formats/mod.rs b/rust/src/events/formats/mod.rs index 86023add17..8f24ffd757 100644 --- a/rust/src/events/formats/mod.rs +++ b/rust/src/events/formats/mod.rs @@ -95,9 +95,13 @@ pub use vmsc4242::EventFormatVMSC4242; /// pyclass. /// /// The `signatures` and `unsigned` fields are kept separate from the other -/// fields as they are mutable (and must be deep-copied if the event is cloned). -/// `common_fields` and `specific_fields` are both `#[serde(flatten)]`ed so that -/// the serialised JSON is a single flat object matching the Matrix spec. +/// fields as they are mutable. Note the derived [`Clone`] is *shallow*: it +/// shares the mutable `signatures`/`unsigned`/internal state behind their +/// `Arc`s (cheap, and fine for read-only uses such as bundled aggregations). +/// Use [`FormattedEvent::deep_copy`] when an independently-mutable copy is +/// required. `common_fields` and `specific_fields` are both +/// `#[serde(flatten)]`ed so that the serialised JSON is a single flat object +/// matching the Matrix spec. /// /// Note, deserialization of this struct must not be done from /// [`serde_json::Value`] nor [`pythonize::depythonize`], due to a bug with @@ -105,7 +109,7 @@ pub use vmsc4242::EventFormatVMSC4242; /// Instead, deserialize directly from a JSON string with /// `serde_json::from_str`. See https://github.com/serde-rs/serde/issues/2230 /// for details. -#[derive(Serialize, Deserialize)] +#[derive(Clone, Serialize, Deserialize)] pub struct FormattedEvent> { /// The event's signatures. /// diff --git a/rust/src/events/json_object.rs b/rust/src/events/json_object.rs index bb4877d482..bfb3396792 100644 --- a/rust/src/events/json_object.rs +++ b/rust/src/events/json_object.rs @@ -17,14 +17,16 @@ use std::{collections::BTreeMap, sync::Arc}; use pyo3::{ exceptions::{PyKeyError, PyTypeError}, + prelude::Borrowed, pyclass, pymethods, types::{ PyAnyMethods, PyIterator, PyList, PyListMethods, PyMapping, PySet, PySetMethods, PyTuple, }, - Bound, IntoPyObject, IntoPyObjectExt, Py, PyAny, PyResult, Python, + Bound, FromPyObject, IntoPyObject, IntoPyObjectExt, Py, PyAny, PyErr, PyResult, Python, }; use pythonize::{depythonize, pythonize}; use serde::{Deserialize, Serialize}; +use serde_json::Value; /// A generic class for representing immutable JSON objects. /// @@ -40,33 +42,45 @@ pub struct JsonObject { object: Arc, serde_json::Value>>, } +// We implement `FromPyObject` to allow `JsonObject` to be used as function +// arguments. +impl<'py> FromPyObject<'_, 'py> for JsonObject { + type Error = PyErr; + + fn extract(ob: Borrowed<'_, 'py, PyAny>) -> Result { + // Fast path: already a JsonObject, so just share the underlying map + // (cheap, as it's immutable and behind an `Arc`). + if let Ok(obj) = ob.cast::() { + return Ok(JsonObject { + object: obj.get().object.clone(), + }); + } + + // Otherwise accept any mapping and convert it via pythonize. Unlike the + // `#[new]` constructor we don't accept `None` here: an absent value is + // represented as `Option` at the field/argument level. + let mapping = ob + .cast::() + .map_err(|_| PyTypeError::new_err("expected a mapping"))?; + let object: BTreeMap, Value> = depythonize(&mapping)?; + Ok(JsonObject { + object: Arc::new(object), + }) + } +} + #[pymethods] impl JsonObject { #[new] #[pyo3(signature = (content = None))] - fn new<'a, 'py>(content: Option<&'a Bound<'py, PyAny>>) -> PyResult { - let Some(content) = content else { + fn new(content: Option<&Bound<'_, PyAny>>) -> PyResult { + match content { // If no content is provided, default to an empty object. - return Ok(Self::default()); - }; - - if let Ok(content) = content.cast::() { - // If the content is already a JsonObject, we can just clone the - // underlying map (this is safe as the object is immutable). - return Ok(JsonObject { - object: content.get().object.clone(), - }); + None => Ok(Self::default()), + // Otherwise reuse the `FromPyObject` path, which accepts an + // existing `JsonObject` or any Python mapping. + Some(content) => JsonObject::extract(content.as_borrowed()), } - - let Ok(content) = content.cast::() else { - return Err(PyTypeError::new_err("'content' must be a mapping")); - }; - - // Use pythonize to try and convert from a mapping. - let content = depythonize(content)?; - Ok(Self { - object: Arc::new(content), - }) } fn __len__(&self) -> usize { @@ -197,6 +211,30 @@ impl JsonObject { pub fn get_field(&self, key: &str) -> Option<&serde_json::Value> { self.object.get(key) } + + /// Returns a `serde_json::Value::Object` containing a clone of this + /// object's entries. + pub fn as_map(&self) -> &BTreeMap, Value> { + &self.object + } + + /// Whether the object has no entries. + pub fn is_empty(&self) -> bool { + self.object.is_empty() + } + + pub fn iter(&self) -> impl Iterator, &Value)> { + self.object.iter() + } +} + +impl<'a> IntoIterator for &'a JsonObject { + type Item = (&'a Box, &'a serde_json::Value); + type IntoIter = std::collections::btree_map::Iter<'a, Box, serde_json::Value>; + + fn into_iter(self) -> Self::IntoIter { + self.object.as_ref().iter() + } } /// Helper class returned by `JsonObject.keys()` to act as a view into the keys diff --git a/rust/src/events/mod.rs b/rust/src/events/mod.rs index 83900e14ba..d1f4a01597 100644 --- a/rust/src/events/mod.rs +++ b/rust/src/events/mod.rs @@ -87,6 +87,7 @@ pub mod filter; pub mod formats; pub mod internal_metadata; pub mod json_object; +pub mod relations; pub mod signatures; pub mod unsigned; pub mod utils; @@ -107,6 +108,8 @@ pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> child_module.add_class::()?; child_module.add_class::()?; child_module.add_class::()?; + child_module.add_class::()?; + child_module.add_class::()?; child_module.add_function(wrap_pyfunction!(filter::event_visible_to_server_py, m)?)?; child_module.add_function(wrap_pyfunction!(redact_event_py, m)?)?; child_module.add_function(wrap_pyfunction!(redact_event_dict, m)?)?; @@ -129,7 +132,11 @@ pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> /// metadata, rejection reason, and a reference to the room version that /// produced this event). See the module-level docs for the high-level /// design. -#[pyclass(frozen, weakref)] +/// +/// `Clone` is shallow (see [`FormattedEvent`]) and lets an `Event` be held by +/// value, e.g. inside [`BundledAggregations`](crate::events::relations::BundledAggregations). +#[pyclass(frozen, weakref, skip_from_py_object)] +#[derive(Clone)] pub struct Event { /// The parsed event JSON. parsed_event: FormattedEvent, diff --git a/rust/src/events/relations.rs b/rust/src/events/relations.rs new file mode 100644 index 0000000000..175a3a4cf6 --- /dev/null +++ b/rust/src/events/relations.rs @@ -0,0 +1,101 @@ +/* + * This file is licensed under the Affero General Public License (AGPL) version 3. + * + * Copyright (C) 2026 Element Creations Ltd + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the GNU Affero General Public License as + * published by the Free Software Foundation, either version 3 of the + * License, or (at your option) any later version. + * + * See the GNU Affero General Public License for more details: + * . + */ + +//! The bundled aggregations attached to an event for client serialization. +//! +//! These mirror the Matrix "server-side aggregation" data (references, edits +//! and thread summaries) that is folded into an event's `unsigned.m.relations` +//! section when serializing for clients. They are built by the Python +//! `RelationsHandler` and consumed by [`serialize_events`](crate::events::serialize::serialize_events). +//! +//! The events they reference ([`Event`]) are stored by value rather than as +//! Python handles; cloning an `Event` is cheap (it shares the underlying data +//! behind `Arc`s) and the events are only ever read here. + +use pyo3::{pyclass, pymethods}; + +use crate::events::{json_object::JsonObject, Event}; + +/// A thread's bundled summary: its latest event, the number of events in the +/// thread, and whether the requesting user has participated. +#[pyclass(frozen, name = "ThreadAggregation", skip_from_py_object)] +#[derive(Clone)] +pub struct ThreadAggregation { + /// The latest event in the thread. + #[pyo3(get)] + pub latest_event: Event, + /// The total number of events in the thread. + #[pyo3(get)] + pub count: i64, + /// Whether the requesting user has sent an event to the thread. + #[pyo3(get)] + pub current_user_participated: bool, +} + +#[pymethods] +impl ThreadAggregation { + #[new] + fn new(latest_event: &Event, count: i64, current_user_participated: bool) -> Self { + Self { + latest_event: latest_event.clone(), + count, + current_user_participated, + } + } +} + +/// The bundled aggregations for a single event. +/// +/// Some values require additional processing during serialization (the edit +/// and the thread's latest event are themselves serialized). +#[pyclass(frozen, skip_from_py_object)] +#[derive(Clone)] +pub struct BundledAggregations { + /// The `m.reference` aggregation (e.g. `{"chunk": [{"event_id": ...}]}`). + #[pyo3(get)] + pub references: Option, + /// The edit (`m.replace`) event that applies to this event. + #[pyo3(get)] + pub replace: Option, + /// The thread (`m.thread`) summary for this event. + #[pyo3(get)] + pub thread: Option, +} + +#[pymethods] +impl BundledAggregations { + #[new] + #[pyo3(signature = (references = None, replace = None, thread = None))] + fn new( + references: Option, + replace: Option<&Event>, + thread: Option<&ThreadAggregation>, + ) -> Self { + Self { + references, + replace: replace.cloned(), + thread: thread.cloned(), + } + } + + /// Whether there are any aggregations to bundle. + /// + /// Matches the Python `bool(self.references or self.replace or self.thread)`: + /// an empty `references` mapping counts as falsey. + fn __bool__(&self) -> bool { + self.references.as_ref().is_some_and(|r| !r.is_empty()) + || self.replace.is_some() + || self.thread.is_some() + } +} diff --git a/synapse/handlers/relations.py b/synapse/handlers/relations.py index ee4f8d672e..4fd301a18e 100644 --- a/synapse/handlers/relations.py +++ b/synapse/handlers/relations.py @@ -28,8 +28,6 @@ from typing import ( Sequence, ) -import attr - from synapse.api.constants import Direction, EventTypes, RelationTypes from synapse.api.errors import SynapseError from synapse.events import EventBase, relation_from_event @@ -38,6 +36,14 @@ from synapse.logging.context import make_deferred_yieldable, run_in_background from synapse.logging.opentracing import trace from synapse.storage.databases.main.relations import ThreadsNextBatch, _RelatedEvent from synapse.streams.config import PaginationConfig + +# `BundledAggregations` and `ThreadAggregation` are implemented in Rust; they +# are re-exported here so existing call sites can keep importing them from +# `synapse.handlers.relations`. +from synapse.synapse_rust.events import ( # noqa: F401 + BundledAggregations, + ThreadAggregation, +) from synapse.types import JsonDict, Requester, UserID from synapse.util.async_helpers import gather_results from synapse.visibility import filter_and_transform_events_for_client @@ -56,32 +62,6 @@ class ThreadsListInclude(str, enum.Enum): participated = "participated" -@attr.s(slots=True, frozen=True, auto_attribs=True) -class _ThreadAggregation: - # The latest event in the thread. - latest_event: EventBase - # The total number of events in the thread. - count: int - # True if the current user has sent an event to the thread. - current_user_participated: bool - - -@attr.s(slots=True, auto_attribs=True) -class BundledAggregations: - """ - The bundled aggregations for an event. - - Some values require additional processing during serialization. - """ - - references: JsonDict | None = None - replace: EventBase | None = None - thread: _ThreadAggregation | None = None - - def __bool__(self) -> bool: - return bool(self.references or self.replace or self.thread) - - class RelationsHandler: def __init__(self, hs: "HomeServer"): self._main_store = hs.get_datastores().main @@ -310,7 +290,7 @@ class RelationsHandler: relations_by_id: dict[str, str], user_id: str, ignored_users: frozenset[str], - ) -> dict[str, _ThreadAggregation]: + ) -> dict[str, ThreadAggregation]: """Get the bundled aggregations for threads for the requested events. Args: @@ -421,7 +401,7 @@ class RelationsHandler: continue latest_thread_event = event.event - results[event_id] = _ThreadAggregation( + results[event_id] = ThreadAggregation( latest_event=latest_thread_event, count=thread_count, # If there's a thread summary it must also exist in the @@ -478,8 +458,12 @@ class RelationsHandler: # The event should get bundled aggregations. events_by_id[event.event_id] = event - # event ID -> bundled aggregation in non-serialized form. - results: dict[str, BundledAggregations] = {} + # `BundledAggregations` is immutable, so we collect each kind of + # aggregation into its own map keyed by event ID and assemble the + # results once everything has been fetched. + thread_by_id: dict[str, ThreadAggregation] = {} + references_by_id: dict[str, JsonDict] = {} + replace_by_id: dict[str, EventBase] = {} # Fetch any ignored users of the requesting user. ignored_users = await self._main_store.ignored_users(user_id) @@ -495,7 +479,7 @@ class RelationsHandler: ignored_users, ) for event_id, thread in threads.items(): - results.setdefault(event_id, BundledAggregations()).thread = thread + thread_by_id[event_id] = thread # If the latest event in a thread is not already being fetched, # add it. This ensures that the bundled aggregations for the @@ -516,7 +500,7 @@ class RelationsHandler: ) for event_id, references in references_by_event_id.items(): if references: - results.setdefault(event_id, BundledAggregations()).references = { + references_by_id[event_id] = { "chunk": [{"event_id": ev.event_id} for ev in references] } @@ -535,7 +519,13 @@ class RelationsHandler: ] ) for event_id, edit in edits.items(): - results.setdefault(event_id, BundledAggregations()).replace = edit + # `get_applicable_edits` returns `None` for events with no + # applicable edit. Skip those rather than recording an entry: a + # `None` replace contributes nothing during serialization, so + # the old code's empty `BundledAggregations` for such events was + # inert anyway. + if edit is not None: + replace_by_id[event_id] = edit # Parallelize the calls for annotations, references, and edits since they # are unrelated. @@ -548,7 +538,17 @@ class RelationsHandler: ) ) - return results + # Assemble one (immutable) bundled aggregation per event that has any. + return { + event_id: BundledAggregations( + references=references_by_id.get(event_id), + replace=replace_by_id.get(event_id), + thread=thread_by_id.get(event_id), + ) + for event_id in thread_by_id.keys() + | references_by_id.keys() + | replace_by_id.keys() + } async def get_threads( self, diff --git a/synapse/synapse_rust/events.pyi b/synapse/synapse_rust/events.pyi index f84eeb55d6..28f4857508 100644 --- a/synapse/synapse_rust/events.pyi +++ b/synapse/synapse_rust/events.pyi @@ -309,6 +309,47 @@ class Event: ``SynapseDuration`` representing the sticky duration. Otherwise returns ``None``.""" +class ThreadAggregation: + """The bundled thread summary for an event.""" + + def __init__( + self, + latest_event: Event, + count: int, + current_user_participated: bool, + ) -> None: ... + @property + def latest_event(self) -> Event: + """The latest event in the thread.""" + + @property + def count(self) -> int: + """The total number of events in the thread.""" + + @property + def current_user_participated(self) -> bool: + """Whether the requesting user has sent an event to the thread.""" + +class BundledAggregations: + """The bundled aggregations for an event. + + Some values require additional processing during serialization. + """ + + def __init__( + self, + references: JsonMapping | None = None, + replace: Event | None = None, + thread: ThreadAggregation | None = None, + ) -> None: ... + @property + def references(self) -> JsonMapping | None: ... + @property + def replace(self) -> Event | None: ... + @property + def thread(self) -> ThreadAggregation | None: ... + def __bool__(self) -> bool: ... + def redact_event(event: Event) -> Event: """Returns a pruned version of the given event, which removes all keys we don't know about or think could potentially be dodgy.