From 106d7fcb6046788bf8ba8c02f2bb46aadb807321 Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Fri, 16 Oct 2026 16:36:08 +0000 Subject: [PATCH] Compute the sliding sync room lists and own the connection state in Rust In `rust_sliding_sync: on` mode, `current_sync_for_user` is now four Rust calls around the Python extensions: load the connection state, compute the interested rooms (the port of `_compute_interested_rooms_new_tables`), compute the rooms, and finish the request (unsent and sent rooms, then the persist). The connection state is a Rust pyclass. The extensions read and update it through views with the surface of Python's `PerConnectionState` and `MutablePerConnectionState`. Receipts tokens stay in their database form in Rust and become `MultiWriterStreamToken`s only when Python asks. The whole request falls back to Python when the sliding sync tables are not filled, for a filter Python does not implement either, and on any Rust error. Rooms Rust leaves to Python run `get_room_sync_data` over a Python copy of their slice of the state. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01L51DoxB8VnH8UjXJFDgpNG --- rust/src/sliding_sync/mod.rs | 441 ++++++++- rust/src/sliding_sync/result.rs | 161 +--- rust/src/sliding_sync/room_lists/filter.rs | 4 +- .../src/sliding_sync/room_lists/interested.rs | 908 ++++++++++++++++++ rust/src/sliding_sync/room_lists/lists.rs | 4 +- rust/src/sliding_sync/room_lists/mod.rs | 13 +- rust/src/sliding_sync/state.rs | 740 ++++++++++++++ rust/src/storage/sliding_sync_store.rs | 21 +- synapse/handlers/sliding_sync/__init__.py | 400 ++++++-- synapse/handlers/sliding_sync/extensions.py | 21 +- synapse/synapse_rust/sliding_sync.pyi | 238 ++++- .../client/sliding_sync/test_rust_lists.py | 514 ++++++++++ .../client/sliding_sync/test_rust_rooms.py | 104 +- 13 files changed, 3233 insertions(+), 336 deletions(-) create mode 100644 rust/src/sliding_sync/room_lists/interested.rs create mode 100644 rust/src/sliding_sync/state.rs create mode 100644 tests/rest/client/sliding_sync/test_rust_lists.py diff --git a/rust/src/sliding_sync/mod.rs b/rust/src/sliding_sync/mod.rs index b86bdaee30..7acab54e22 100644 --- a/rust/src/sliding_sync/mod.rs +++ b/rust/src/sliding_sync/mod.rs @@ -26,18 +26,41 @@ //! ([`SlidingSyncHandler::defer`]), and a `_blocking` twin runs to completion //! for tests ([`SlidingSyncHandler::block_on`]). -use std::collections::HashMap; -use std::sync::Arc; +use std::collections::{HashMap, HashSet}; +use std::sync::{Arc, Mutex}; +use anyhow::Context; use futures::future::BoxFuture; -use pyo3::{exceptions::PyRuntimeError, intern, prelude::*, types::PyModule}; -use request::SlidingSyncRequest; -use result::{RoomInput, RoomsResult}; +use pyo3::{ + create_exception, + exceptions::{PyException, PyRuntimeError}, + intern, + prelude::*, + types::PyModule, +}; +use request::{SlidingSyncRequest, SlidingSyncRequestData}; +use result::RoomsResult; +use room_entry::RoomInputData; +use room_lists::filter::NotImplementedFilter; +use room_lists::interested::{InterestedRooms, ListsSnapshot, PyInterestedRooms}; +use room_lists::lists::UnknownPosition; +use state::{ConnectionState, ConnectionStateData, InstanceMap}; +use crate::caches::lock_ignoring_poison; use crate::database::postgres::pool::{ConnectionPool, PyConnectionPool}; use crate::deferred::{create_deferred_with, OnCancel}; use crate::runtime::RustRuntime; +use crate::sliding_sync::connection::PerConnectionState; +use crate::storage::sliding_sync_store::{SlidingSyncReader, StoredConnectionState}; use crate::storage::sync_store::{anyhow_to_py, SyncDbStats, SyncStore}; +use crate::types::stream_token::RoomStreamToken; + +create_exception!( + synapse.synapse_rust.sliding_sync, + RustSlidingSyncFallback, + PyException, + "The request is one Rust leaves to Python." +); pub mod config; pub mod connection; @@ -49,6 +72,7 @@ pub mod result; pub mod room; pub mod room_entry; pub mod room_lists; +pub mod state; /// Computes sliding sync responses with the Rust implementation, reading /// the database through the native Postgres driver. @@ -109,18 +133,133 @@ impl SlidingSyncHandler { py.detach(|| handle.block_on(future)) } - /// Start computing the entries of `rooms`: read them from Python, and - /// return the future that computes them off the GIL. + /// The previous state's slice that one room's computation reads. + fn room_input( + interested: &InterestedRooms, + previous: &PerConnectionState, + room_id: &Arc, + ) -> anyhow::Result { + let membership = interested + .memberships + .get(room_id) + .with_context(|| format!("no membership for {room_id}"))? + .clone(); + let room_sync_config = interested + .relevant_room_map + .get(room_id) + .with_context(|| format!("no config for {room_id}"))? + .clone(); + Ok(RoomInputData { + membership, + room_sync_config, + prev_room_sync_config: previous.room_configs.get(&**room_id).cloned(), + room_status: previous.rooms.have_sent_room(room_id), + newly_joined: interested.newly_joined.contains(room_id), + newly_left: interested.newly_left.contains(room_id), + is_dm: interested.dm_room_ids.contains(&**room_id), + }) + } + + /// Start loading the connection state at the request's position: an + /// empty state for a first request. + fn start_load_connection_state( + &self, + py: Python<'_>, + request: &Bound<'_, SlidingSyncRequest>, + instance_ids: HashMap, + ) -> PyResult>> { + let pool = self.pool(py)?; + let request = Arc::clone(&request.get().data); + let store = Arc::clone(&self.store); + let now_ms = self.runtime.clock().now_millis() as i64; + Ok(Box::pin(async move { + let reader = store.sliding_sync_reader(pool); + let result = async { + let mut instances = reader.instance_map().await?; + instances.extend( + instance_ids + .into_iter() + .map(|(name, id)| (id, Arc::from(name))), + ); + let instances = InstanceMap::new(instances); + let stored = match &request.from { + Some(from) if from.connection_position != 0 => { + let device_id = request + .device_id + .as_deref() + .context("sliding sync needs a device")?; + let conn_id = request.body.conn_id.as_deref().unwrap_or(""); + reader + .load_connection_state( + &request.user_id, + device_id, + conn_id, + from.connection_position, + now_ms, + ) + .await? + } + _ => Arc::new(StoredConnectionState::default()), + }; + let data = ConnectionStateData::from_stored(&stored, &instances)?; + let stats = SlidingSyncStats::from_db(&reader.stats()); + Ok::<_, anyhow::Error>(ConnectionState::new(data, instances, stats)) + } + .await; + result.map_err(error_to_py) + })) + } + + /// Start computing the room lists. + fn start_interested_rooms( + &self, + py: Python<'_>, + request: &Bound<'_, SlidingSyncRequest>, + state: &Bound<'_, ConnectionState>, + snapshot: ListsSnapshot, + ) -> PyResult>> { + let pool = self.pool(py)?; + let request = Arc::clone(&request.get().data); + let previous = lock_ignoring_poison(&state.get().data).previous.clone(); + let store = Arc::clone(&self.store); + let now_ms = self.runtime.clock().now_millis() as i64; + Ok(Box::pin(async move { + let reader = store.sliding_sync_reader(pool); + let interested = room_lists::interested::compute_interested_rooms( + &reader, &request, &previous, snapshot, now_ms, + ) + .await + .map_err(error_to_py)?; + Ok(PyInterestedRooms { + data: Arc::new(interested), + stats: SlidingSyncStats::from_db(&reader.stats()), + }) + })) + } + + /// Start computing the entries of the rooms to send. The rooms' + /// changes to the connection state go into `state`. fn start_room_entries( &self, py: Python<'_>, request: &Bound<'_, SlidingSyncRequest>, - rooms: Vec>, - instance_ids: HashMap, + state: &Bound<'_, ConnectionState>, + interested: &Bound<'_, PyInterestedRooms>, ) -> PyResult>> { let pool = self.pool(py)?; let request = Arc::clone(&request.get().data); - let rooms: Vec<_> = rooms.iter().map(|room| room.get().data.clone()).collect(); + let interested = Arc::clone(&interested.get().data); + let state_data = Arc::clone(&state.get().data); + let instance_ids = state.get().instances.ids_by_name(); + let rooms = { + let data = lock_ignoring_poison(&state_data); + interested + .rooms_to_send + .iter() + .map(|room_id| Self::room_input(&interested, &data.previous, room_id)) + .collect::>>() + .map_err(anyhow_to_py)? + }; let store = Arc::clone(&self.store); let now_ms = self.runtime.clock().now_millis() as i64; Ok(Box::pin(async move { @@ -130,38 +269,289 @@ impl SlidingSyncHandler { .await .map_err(anyhow_to_py)?; let stats = SlidingSyncStats::from_db(&reader.stats()); - Ok(RoomsResult::new(outcomes, request.from.is_some(), stats)) + let result = RoomsResult::new(outcomes, request.from.is_some(), stats); + result.apply_to(&mut lock_ignoring_poison(&state_data).new); + Ok(result) })) } + + /// Start the end of the request: record what was sent and what was + /// not, and persist the state if it changed. Resolves to the new + /// connection position. + fn start_finish_request( + &self, + py: Python<'_>, + request: &Bound<'_, SlidingSyncRequest>, + state: &Bound<'_, ConnectionState>, + interested: &Bound<'_, PyInterestedRooms>, + ) -> PyResult>> { + let pool = self.pool(py)?; + let request = Arc::clone(&request.get().data); + let interested = Arc::clone(&interested.get().data); + let state_data = Arc::clone(&state.get().data); + let instances = Arc::clone(&state.get().instances); + let store = Arc::clone(&self.store); + let now_ms = self.runtime.clock().now_millis() as i64; + Ok(Box::pin(async move { + let reader = store.sliding_sync_reader(pool); + let position = finish_request( + &reader, + &request, + &state_data, + &instances, + &interested, + now_ms, + ) + .await + .map_err(error_to_py)?; + Ok((position, SlidingSyncStats::from_db(&reader.stats()))) + })) + } +} + +/// The rest of `current_sync_for_user` after the extensions. +async fn finish_request( + reader: &SlidingSyncReader<'_>, + request: &SlidingSyncRequestData, + state: &Mutex, + instances: &InstanceMap, + interested: &InterestedRooms, + now_ms: i64, +) -> anyhow::Result { + let unchanged_position = request + .from + .as_ref() + .map_or(0, |from| from.connection_position); + if !interested.has_lists_or_subscriptions { + return Ok(unchanged_position); + } + + if let Some(from) = &request.from { + // The rooms the client may care about but that are in no range nor + // subscribed to. Recording rooms without anything new for the user + // is fine. + let missing: Vec> = interested + .all_rooms + .iter() + .filter(|room_id| !interested.relevant_room_map.contains_key(*room_id)) + .cloned() + .collect(); + let unsent = reader + .rooms_with_updates_since( + &missing, + &from.room_key, + interested.events_changed_since_from.as_ref(), + ) + .await?; + lock_ignoring_poison(state) + .new + .rooms + .record_unsent_rooms(unsent.iter().map(|room_id| &**room_id), &from.room_key); + } + + let updates = { + let mut data = lock_ignoring_poison(state); + data.new + .rooms + .record_sent_rooms(interested.rooms_to_send.iter().map(|room_id| &**room_id)); + if !data.new.has_updates(now_ms) { + return Ok(unchanged_position); + } + data.updates(instances)? + }; + + // A zero position means no state was stored before. + let previous_position = request + .from + .as_ref() + .map(|from| from.connection_position) + .filter(|position| *position != 0); + let device_id = request + .device_id + .as_deref() + .context("sliding sync needs a device")?; + let conn_id = request.body.conn_id.as_deref().unwrap_or(""); + reader + .persist_connection_state( + &request.user_id, + device_id, + conn_id, + previous_position, + &updates, + now_ms, + ) + .await +} + +/// An error for Python: Python's `SlidingSyncUnknownPosition` for an +/// unknown position or a connection to reset, `RustSlidingSyncFallback` +/// for what Rust leaves to Python. +fn error_to_py(e: anyhow::Error) -> PyErr { + if e.downcast_ref::().is_some() { + return Python::attach(unknown_position_to_py); + } + if let Some(filter) = e.downcast_ref::() { + return RustSlidingSyncFallback::new_err(filter.to_string()); + } + anyhow_to_py(e) +} + +fn unknown_position_to_py(py: Python<'_>) -> PyErr { + match py + .import("synapse.api.errors") + .and_then(|errors| errors.getattr("SlidingSyncUnknownPosition")) + .and_then(|class| class.call0()) + { + Ok(exception) => PyErr::from_value(exception), + Err(e) => e, + } } #[pymethods] impl SlidingSyncHandler { - /// `get_room_sync_data` for each of `rooms`, twenty at a time: a - /// deferred resolving to a `RoomsResult`. `instance_ids` are - /// `get_id_for_instance` for the writers the rooms' tokens may name. - #[pyo3(signature = (request, rooms, *, instance_ids))] + /// `get_and_clear_connection_positions`: a deferred resolving to the + /// request's `ConnectionState`. `instance_ids` are `get_id_for_instance` + /// for the writers of the request's tokens. Fails with + /// `SlidingSyncUnknownPosition` for a position the store does not know. + #[pyo3(signature = (request, *, instance_ids))] + fn load_connection_state<'py>( + &self, + py: Python<'py>, + request: &Bound<'py, SlidingSyncRequest>, + instance_ids: HashMap, + ) -> PyResult> { + let future = self.start_load_connection_state(py, request, instance_ids)?; + self.defer(py, future) + } + + /// As [`Self::load_connection_state`], but blocks. + #[pyo3(signature = (request, *, instance_ids))] + fn load_connection_state_blocking( + &self, + py: Python<'_>, + request: &Bound<'_, SlidingSyncRequest>, + instance_ids: HashMap, + ) -> PyResult { + let future = self.start_load_connection_state(py, request, instance_ids)?; + self.block_on(py, future) + } + + /// `compute_interested_rooms`: a deferred resolving to an + /// `InterestedRooms`. The arguments after `state` are the stream change + /// caches' answers and the config the lists read; see + /// [`ListsSnapshot`]. Fails with `RustSlidingSyncFallback` for a filter + /// Python does not implement either. + #[pyo3(signature = (request, state, *, server_name, msc4155_enabled, num_rooms_threshold, membership_changed_since_from, membership_changed_since_to, events_changed_since_from, room_max_token, events_stream_cache))] + #[allow(clippy::too_many_arguments)] + fn compute_interested_rooms<'py>( + &self, + py: Python<'py>, + request: &Bound<'py, SlidingSyncRequest>, + state: &Bound<'py, ConnectionState>, + server_name: &str, + msc4155_enabled: bool, + num_rooms_threshold: usize, + membership_changed_since_from: bool, + membership_changed_since_to: bool, + events_changed_since_from: Option>, + room_max_token: RoomStreamToken, + events_stream_cache: Py, + ) -> PyResult> { + let snapshot = ListsSnapshot { + server_name: Arc::from(server_name), + msc4155_enabled, + num_rooms_threshold, + membership_changed_since_from, + membership_changed_since_to, + events_changed_since_from: events_changed_since_from + .map(|rooms| rooms.into_iter().map(Arc::from).collect()), + room_max_token, + events_stream_cache, + }; + let future = self.start_interested_rooms(py, request, state, snapshot)?; + self.defer(py, future) + } + + /// As [`Self::compute_interested_rooms`], but blocks. + #[pyo3(signature = (request, state, *, server_name, msc4155_enabled, num_rooms_threshold, membership_changed_since_from, membership_changed_since_to, events_changed_since_from, room_max_token, events_stream_cache))] + #[allow(clippy::too_many_arguments)] + fn compute_interested_rooms_blocking( + &self, + py: Python<'_>, + request: &Bound<'_, SlidingSyncRequest>, + state: &Bound<'_, ConnectionState>, + server_name: &str, + msc4155_enabled: bool, + num_rooms_threshold: usize, + membership_changed_since_from: bool, + membership_changed_since_to: bool, + events_changed_since_from: Option>, + room_max_token: RoomStreamToken, + events_stream_cache: Py, + ) -> PyResult { + let snapshot = ListsSnapshot { + server_name: Arc::from(server_name), + msc4155_enabled, + num_rooms_threshold, + membership_changed_since_from, + membership_changed_since_to, + events_changed_since_from: events_changed_since_from + .map(|rooms| rooms.into_iter().map(Arc::from).collect()), + room_max_token, + events_stream_cache, + }; + let future = self.start_interested_rooms(py, request, state, snapshot)?; + self.block_on(py, future) + } + + /// `get_room_sync_data` for each of the rooms to send, twenty at a + /// time: a deferred resolving to a `RoomsResult`. The rooms' changes to + /// the connection state go into `state`. fn generate_room_entries<'py>( &self, py: Python<'py>, request: &Bound<'py, SlidingSyncRequest>, - rooms: Vec>, - instance_ids: HashMap, + state: &Bound<'py, ConnectionState>, + interested: &Bound<'py, PyInterestedRooms>, ) -> PyResult> { - let future = self.start_room_entries(py, request, rooms, instance_ids)?; + let future = self.start_room_entries(py, request, state, interested)?; self.defer(py, future) } /// As [`Self::generate_room_entries`], but blocks. - #[pyo3(signature = (request, rooms, *, instance_ids))] fn generate_room_entries_blocking( &self, py: Python<'_>, request: &Bound<'_, SlidingSyncRequest>, - rooms: Vec>, - instance_ids: HashMap, + state: &Bound<'_, ConnectionState>, + interested: &Bound<'_, PyInterestedRooms>, ) -> PyResult { - let future = self.start_room_entries(py, request, rooms, instance_ids)?; + let future = self.start_room_entries(py, request, state, interested)?; + self.block_on(py, future) + } + + /// The end of the request, after the extensions: a deferred resolving + /// to `(connection_position, stats)`. + fn finish_request<'py>( + &self, + py: Python<'py>, + request: &Bound<'py, SlidingSyncRequest>, + state: &Bound<'py, ConnectionState>, + interested: &Bound<'py, PyInterestedRooms>, + ) -> PyResult> { + let future = self.start_finish_request(py, request, state, interested)?; + self.defer(py, future) + } + + /// As [`Self::finish_request`], but blocks. + fn finish_request_blocking( + &self, + py: Python<'_>, + request: &Bound<'_, SlidingSyncRequest>, + state: &Bound<'_, ConnectionState>, + interested: &Bound<'_, PyInterestedRooms>, + ) -> PyResult<(i64, SlidingSyncStats)> { + let future = self.start_finish_request(py, request, state, interested)?; self.block_on(py, future) } } @@ -229,9 +619,14 @@ 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_class::()?; + child_module.add_class::()?; + child_module.add( + "RustSlidingSyncFallback", + py.get_type::(), + )?; + state::register_module(&child_module)?; child_module.add_function(wrap_pyfunction!(encode::encode_rooms, &child_module)?)?; child_module.add_class::()?; config::register_test_functions(&child_module)?; diff --git a/rust/src/sliding_sync/result.rs b/rust/src/sliding_sync/result.rs index 0c88c664ba..b542f29d83 100644 --- a/rust/src/sliding_sync/result.rs +++ b/rust/src/sliding_sync/result.rs @@ -12,33 +12,25 @@ * . */ -//! What crosses to Python for the rooms of a request: [`RoomInput`] in, -//! [`RoomsResult`] out, and a [`RoomResult`] per room Rust computed. +//! What crosses to Python for the rooms of a request: [`RoomsResult`], and a +//! [`RoomResult`] per room Rust computed. //! //! A `RoomResult`'s events never become Python objects, except for module //! callbacks: Python reads only what `current_sync_for_user` and the //! extensions need, and the servlet encodes the room in Rust. -use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; +use std::collections::BTreeSet; use std::sync::{Arc, Mutex}; -use pyo3::{ - exceptions::PyValueError, - intern, - prelude::*, - types::{PyList, PyTuple}, -}; +use pyo3::{intern, prelude::*, types::PyList}; use crate::caches::lock_ignoring_poison; use crate::events::Event; use crate::sliding_sync::config::{RequiredStateMap, RoomSyncConfig, MEMBER}; -use crate::sliding_sync::connection::HaveSentRoom; -use crate::sliding_sync::room_entry::{RoomEntry, RoomInputData, RoomOutcome}; -use crate::sliding_sync::room_lists::RoomForUser; +use crate::sliding_sync::connection::{MutablePerConnectionState, RoomLazyMembershipChanges}; +use crate::sliding_sync::room_entry::{RoomEntry, RoomOutcome}; use crate::sliding_sync::SlidingSyncStats; -use crate::storage::sync_store::PersistedEventPosition; use crate::sync::relations::Bundle; -use crate::types::stream_token::RoomStreamToken; /// A `required_state_map` from Python: a mapping of types to sets of /// state keys. @@ -56,90 +48,6 @@ pub fn required_state_map_from_py(map: &Bound<'_, PyAny>) -> PyResult, - membership: &str, - event_id: Option<&str>, - instance_name: &str, - stream: i64, - room_version_id: Option<&str>, - timeline_limit: i64, - required_state_map: &Bound<'_, PyAny>, - prev_room_sync_config: Option<(i64, Bound<'_, PyAny>)>, - room_status: &str, - room_status_token: Option, - newly_joined: bool, - newly_left: bool, - is_dm: bool, - ) -> PyResult { - let room_status = match (room_status, room_status_token) { - ("never", _) => HaveSentRoom::Never, - ("live", _) => HaveSentRoom::Live, - ("previously", Some(token)) => HaveSentRoom::Previously(token), - (status, _) => { - return Err(PyValueError::new_err(format!( - "bad room status {status:?} or missing token" - ))) - } - }; - let prev_room_sync_config = prev_room_sync_config - .map(|(timeline_limit, map)| { - Ok::<_, PyErr>(RoomSyncConfig { - timeline_limit, - required_state_map: required_state_map_from_py(&map)?, - }) - }) - .transpose()?; - Ok(RoomInput { - data: RoomInputData { - membership: RoomForUser { - room_id: Arc::from(room_id), - sender: sender.map(Arc::from), - membership: Arc::from(membership), - event_id: event_id.map(Arc::from), - event_pos: PersistedEventPosition { - instance_name: Arc::from(instance_name), - stream, - }, - room_version_id: room_version_id.map(Arc::from), - // Not read by the room's computation. - has_known_state: true, - room_type: None, - is_encrypted: false, - }, - room_sync_config: RoomSyncConfig { - timeline_limit, - required_state_map: required_state_map_from_py(required_state_map)?, - }, - prev_room_sync_config, - room_status, - newly_joined, - newly_left, - is_dm, - }, - }) - } - - #[getter] - fn room_id(&self) -> &str { - &self.data.membership.room_id - } -} - /// One room's entry, computed in Rust: `SlidingSyncResult.RoomResult` as /// far as `current_sync_for_user` and the extensions read it. #[pyclass(frozen, skip_from_py_object)] @@ -227,13 +135,6 @@ impl RoomResult { } } -/// A [`RequiredStateMap`] for Python: types to sets of state keys. -type PyRequiredStateMap = BTreeMap>; - -/// A room's `RoomLazyMembershipChanges`: `(room_id, -/// returned_user_id_to_last_seen_ts_map, invalidated_user_ids)`. -type LazyMembership = (Arc, HashMap>, HashSet); - /// What [`crate::sliding_sync::room_entry::generate_room_entries`] made of /// the rooms. #[pyclass(frozen, skip_from_py_object)] @@ -245,7 +146,7 @@ pub struct RoomsResult { /// `room_configs[room_id]` for the computed rooms whose config changes. room_config_updates: Vec<(Arc, RoomSyncConfig)>, /// `room_lazy_membership[room_id]` for every computed room. - lazy_memberships: Vec, + lazy_memberships: Vec<(Arc, RoomLazyMembershipChanges)>, fallbacks: Vec<(Arc, String)>, errors: Vec<(Arc, String)>, stats: SlidingSyncStats, @@ -276,12 +177,7 @@ impl RoomsResult { if let Some(config) = mutations.room_config { room_config_updates.push((Arc::clone(&room_id), config)); } - let lazy = mutations.lazy_membership; - lazy_memberships.push(( - Arc::clone(&room_id), - lazy.returned_user_id_to_last_seen_ts_map, - lazy.invalidated_user_ids, - )); + lazy_memberships.push((Arc::clone(&room_id), mutations.lazy_membership)); if computed.entry.has_updates() || !incremental { entries.push(Arc::new(computed.entry)); } else { @@ -302,6 +198,18 @@ impl RoomsResult { stats, } } + + /// Record the computed rooms' changes in the connection state. + pub fn apply_to(&self, state: &mut MutablePerConnectionState) { + for (room_id, config) in &self.room_config_updates { + state.set_room_config(room_id, config.clone()); + } + for (room_id, changes) in &self.lazy_memberships { + state + .room_lazy_membership + .insert(room_id.to_string(), changes.clone()); + } + } } #[pymethods] @@ -322,35 +230,6 @@ impl RoomsResult { self.empty.iter().map(|room_id| &**room_id).collect() } - /// `(room_id, timeline_limit, required_state_map)` for each room whose - /// config the connection records. - #[getter] - fn room_config_updates<'py>(&self, py: Python<'py>) -> PyResult> { - let items = self - .room_config_updates - .iter() - .map(|(room_id, config)| { - let map: PyRequiredStateMap = config.required_state_map.clone(); - (&**room_id, config.timeline_limit, map).into_pyobject(py) - }) - .collect::>>>()?; - PyList::new(py, items) - } - - /// `(room_id, returned_user_id_to_last_seen_ts_map, - /// invalidated_user_ids)` for each room computed. - #[getter] - fn lazy_memberships<'py>(&self, py: Python<'py>) -> PyResult> { - let items = self - .lazy_memberships - .iter() - .map(|(room_id, returned, invalidated)| { - (&**room_id, returned, invalidated).into_pyobject(py) - }) - .collect::>>>()?; - PyList::new(py, items) - } - /// `(room_id, reason)` for each room Python must compute. #[getter] fn fallbacks(&self) -> Vec<(&str, &str)> { diff --git a/rust/src/sliding_sync/room_lists/filter.rs b/rust/src/sliding_sync/room_lists/filter.rs index 11186eac38..3f9dd444eb 100644 --- a/rust/src/sliding_sync/room_lists/filter.rs +++ b/rust/src/sliding_sync/room_lists/filter.rs @@ -52,7 +52,7 @@ pub fn filter_membership_for_sync(user_id: &str, room: &RoomForUser, newly_left: /// previous connection state). `room_tags` maps a room to the user's tags /// on it, and is only read for the tag filters. pub fn filter_rooms_using_tables<'a>( - rooms: &'a [RoomForUser], + rooms: &[&'a RoomForUser], filters: &Filters, dm_room_ids: &HashSet, room_tags: &HashMap>, @@ -117,6 +117,7 @@ pub fn filter_rooms_using_tables<'a>( Ok(rooms .iter() + .copied() .filter(|room| { passes(room) // State reset out of a room the client knows: it must hear @@ -150,6 +151,7 @@ fn filter_rooms_for_test( let filters: Filters = serde_json::from_str(filters).map_err(|e| PyValueError::new_err(e.to_string()))?; let rooms: Vec = rooms.into_iter().map(room_from_py).collect(); + let rooms: Vec<&RoomForUser> = rooms.iter().collect(); let filtered = filter_rooms_using_tables(&rooms, &filters, &dm_room_ids, &room_tags, |r| { sent_before.contains(r) }) diff --git a/rust/src/sliding_sync/room_lists/interested.rs b/rust/src/sliding_sync/room_lists/interested.rs new file mode 100644 index 0000000000..382cf860d0 --- /dev/null +++ b/rust/src/sliding_sync/room_lists/interested.rs @@ -0,0 +1,908 @@ +/* + * 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: + * . + */ + +//! `compute_interested_rooms` on the sliding sync tables +//! (`_compute_interested_rooms_new_tables`): the user's rooms at the +//! `to_token`, the lists over them, the room subscriptions, and which of +//! the rooms may have updates to send. + +use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; +use std::sync::Arc; + +use anyhow::bail; +use indexmap::IndexMap; +use pyo3::{exceptions::PyKeyError, prelude::*}; +use regex::Regex; + +use crate::push::utils::{glob_to_regex, GlobMatchType}; +use crate::sliding_sync::config::RoomSyncConfig; +use crate::sliding_sync::connection::PerConnectionState; +use crate::sliding_sync::request::SlidingSyncRequestData; +use crate::sliding_sync::room_lists::filter::{ + filter_membership_for_sync, filter_rooms_using_tables, +}; +use crate::sliding_sync::room_lists::lists::{ + add_room_subscription, needs_sort, rooms_to_send, sort_limit, sync_list_ops, RelevantRoomMap, + SyncOp, +}; +use crate::sliding_sync::room_lists::sort::SortPlan; +use crate::sliding_sync::room_lists::RoomForUser; +use crate::sliding_sync::SlidingSyncStats; +use crate::storage::sliding_sync_store::{ + MembershipChange, RoomType, SlidingSyncReader, SnapshotRooms, +}; +use crate::types::stream_token::RoomStreamToken; + +/// `AccountDataTypes.DIRECT`. +const DIRECT: &str = "m.direct"; +/// `AccountDataTypes.INVITE_PERMISSION_CONFIG` (MSC4380). +const INVITE_PERMISSION_CONFIG: &str = "m.invite_permission_config"; +/// `AccountDataTypes.MSC4155_INVITE_PERMISSION_CONFIG`. +const MSC4155_INVITE_PERMISSION_CONFIG: &str = "org.matrix.msc4155.invite_permission_config"; + +/// What Python knows that the room lists need: the in-memory stream +/// change caches at the time of the request, and server config. +pub struct ListsSnapshot { + /// Ours, for `is_mine_id`. + pub server_name: Arc, + /// `msc4155_enabled`, for the invite rules. + pub msc4155_enabled: bool, + /// `NUM_ROOMS_THRESHOLD`. + pub num_rooms_threshold: usize, + /// `_membership_stream_cache.has_entity_changed(user_id, from.stream)`, + /// true without a `from` token. + pub membership_changed_since_from: bool, + /// The same at the `to_token`'s stream. + pub membership_changed_since_to: bool, + /// `_events_stream_cache.get_all_entities_changed(from.stream)`. + pub events_changed_since_from: Option>>, + /// `get_room_max_token()`. + pub room_max_token: RoomStreamToken, + /// `_events_stream_cache`, read for the rooms a sort needs. + pub events_stream_cache: Py, +} + +/// `SlidingSyncInterestedRooms`. +#[derive(Debug, Default)] +pub struct InterestedRooms { + /// Whether the request has any lists or room subscriptions. + pub has_lists_or_subscriptions: bool, + /// Each list's name, room count and operations, in the request's order. + pub lists: Vec<(String, usize, Vec)>, + pub relevant_room_map: RelevantRoomMap, + /// The rooms of `relevant_room_map` that may have updates, in its + /// order. + pub rooms_to_send: Vec>, + /// Every room that could appear in a list, in range or not. + pub all_rooms: HashSet>, + /// The user's membership at the `to_token` of each relevant room. + pub memberships: HashMap, RoomForUser>, + pub newly_joined: HashSet>, + pub newly_left: HashSet>, + pub dm_room_ids: HashSet, + /// The snapshot's `events_changed_since_from`, for the rooms out of + /// range at the end of the request. + pub events_changed_since_from: Option>>, +} + +/// `MutableOverlayMapping` over the snapshot rooms: changes go to an +/// overlay, which iterates first, as Python's does. +struct OverlayMap { + base: Arc, + overlay: IndexMap, RoomForUser>, + deletions: HashSet>, +} + +impl OverlayMap { + fn new(base: Arc) -> Self { + OverlayMap { + base, + overlay: IndexMap::new(), + deletions: HashSet::new(), + } + } + + fn get(&self, room_id: &str) -> Option<&RoomForUser> { + if self.deletions.contains(room_id) { + return None; + } + self.overlay.get(room_id).or_else(|| self.base.get(room_id)) + } + + fn contains_key(&self, room_id: &str) -> bool { + self.get(room_id).is_some() + } + + fn insert(&mut self, room: RoomForUser) { + self.deletions.remove(&room.room_id); + self.overlay.insert(Arc::clone(&room.room_id), room); + } + + /// `pop(room_id, None)`. + fn remove(&mut self, room_id: &str) { + if !self.contains_key(room_id) { + return; + } + // `dict.pop`, so a key set again goes to the end. + self.overlay.shift_remove(room_id); + self.deletions.insert(Arc::from(room_id)); + } + + fn values(&self) -> impl Iterator { + let overlay = self + .overlay + .values() + .filter(|room| !self.deletions.contains(&room.room_id)); + let base = self.base.values().filter(|room| { + !self.deletions.contains(&room.room_id) && !self.overlay.contains_key(&room.room_id) + }); + overlay.chain(base) + } + + fn keys(&self) -> Vec> { + self.values() + .map(|room| Arc::clone(&room.room_id)) + .collect() + } +} + +/// `InviteRule`, as far as the room lists read it. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum InviteRule { + Allow, + Ignore, + Block, +} + +/// `InviteRulesConfig`: the user's choices about which invites to accept. +#[derive(Debug, Default)] +pub enum InviteRules { + /// `AllowAllInviteRulesConfig`. + #[default] + AllowAll, + /// `MSC4380InviteRulesConfig`: one rule for every invite. + Default(InviteRule), + /// `MSC4155InviteRulesConfig`: globs over user IDs, then over servers, + /// each in allow, ignore, block order. + Msc4155 { + users: Vec<(Vec, InviteRule)>, + servers: Vec<(Vec, InviteRule)>, + }, +} + +impl InviteRules { + /// `get_invite_config_for_user`. + pub async fn load( + reader: &SlidingSyncReader<'_>, + user_id: &str, + msc4155_enabled: bool, + ) -> anyhow::Result { + if let Some(data) = reader + .global_account_data_by_type(user_id, INVITE_PERMISSION_CONFIG) + .await? + { + let rule = match data.get("default_action").and_then(|v| v.as_str()) { + Some("block") => InviteRule::Block, + _ => InviteRule::Allow, + }; + return Ok(InviteRules::Default(rule)); + } + if msc4155_enabled { + if let Some(data) = reader + .global_account_data_by_type(user_id, MSC4155_INVITE_PERMISSION_CONFIG) + .await? + { + return Ok(Self::msc4155(&data)); + } + } + Ok(InviteRules::AllowAll) + } + + /// `MSC4155InviteRulesConfig(account_data)`. + pub fn msc4155(data: &serde_json::Value) -> Self { + let field = |name: &str| -> Vec { + let Some(values) = data.get(name).and_then(|v| v.as_array()) else { + return Vec::new(); + }; + values + .iter() + .filter_map(|value| value.as_str()) + // User IDs cannot exceed 255 bytes, so longer globs are + // skipped, as are ones that do not compile. + .filter(|glob| !glob.is_empty() && glob.len() <= 255) + .filter_map(|glob| glob_to_regex(glob, GlobMatchType::Whole).ok()) + .collect() + }; + InviteRules::Msc4155 { + users: vec![ + (field("allowed_users"), InviteRule::Allow), + (field("ignored_users"), InviteRule::Ignore), + (field("blocked_users"), InviteRule::Block), + ], + servers: vec![ + (field("allowed_servers"), InviteRule::Allow), + (field("ignored_servers"), InviteRule::Ignore), + (field("blocked_servers"), InviteRule::Block), + ], + } + } + + /// `get_invite_rule`. + pub fn invite_rule(&self, inviter: &str) -> InviteRule { + match self { + InviteRules::AllowAll => InviteRule::Allow, + InviteRules::Default(rule) => *rule, + InviteRules::Msc4155 { users, servers } => { + for (patterns, rule) in users { + if patterns.iter().any(|regex| regex.is_match(inviter)) { + return *rule; + } + } + let domain = inviter.split_once(':').map_or("", |(_, domain)| domain); + for (patterns, rule) in servers { + if patterns.iter().any(|regex| regex.is_match(domain)) { + return *rule; + } + } + InviteRule::Allow + } + } + } +} + +/// `_get_dm_rooms_for_user`: the rooms of the user's `m.direct`, ignoring +/// whatever is not a list of strings. +fn dm_room_ids(direct: Option<&serde_json::Value>) -> HashSet { + let Some(serde_json::Value::Object(map)) = direct else { + return HashSet::new(); + }; + map.values() + .filter_map(|room_ids| room_ids.as_array()) + .flatten() + .filter_map(|room_id| room_id.as_str()) + .map(str::to_owned) + .collect() +} + +/// A rewound membership with the room's state from `state`: Python's +/// `RoomsForUserSlidingSync(...)` "we keep the state of the room though". +fn rewound(change: &MembershipChange, state: &RoomForUser) -> RoomForUser { + RoomForUser { + room_id: Arc::clone(&change.room_id), + sender: change.sender.clone(), + membership: Arc::clone(&change.membership), + event_id: change.event_id.clone(), + event_pos: change.event_pos.clone(), + room_version_id: change.room_version_id.clone(), + has_known_state: state.has_known_state, + room_type: state.room_type.clone(), + is_encrypted: state.is_encrypted, + } +} + +/// Keep `room` if the sync covers it, else drop the room: +/// `filter_membership_for_sync` and the assignment or `pop` after it. +fn keep_if_covered(rooms: &mut OverlayMap, user_id: &str, room: RoomForUser, newly_left: bool) { + if filter_membership_for_sync(user_id, &room, newly_left) { + rooms.insert(room); + } else { + rooms.remove(&room.room_id); + } +} + +/// `_get_rewind_changes_to_current_membership_to_token`: the changes that +/// take the current memberships in `rooms` back to the `to_token`, by +/// room. `None` drops the room: the user joined it after the token, or the +/// previous membership is gone. +async fn rewind_changes( + reader: &SlidingSyncReader<'_>, + request: &SlidingSyncRequestData, + rooms: &OverlayMap, + membership_changed_since_to: bool, +) -> anyhow::Result, Option>> { + let mut changes = IndexMap::new(); + + // The token the snapshot was read at: each writer's latest membership. + let mut max_by_instance: IndexMap, i64> = IndexMap::new(); + for room in rooms.values() { + let max = max_by_instance + .entry(Arc::clone(&room.event_pos.instance_name)) + .or_insert(room.event_pos.stream); + *max = (*max).max(room.event_pos.stream); + } + let Some(min_stream) = max_by_instance.values().copied().min() else { + return Ok(changes); + }; + let snapshot_token = RoomStreamToken::with_instance_map( + min_stream, + max_by_instance + .into_iter() + .filter(|(_, stream)| *stream > min_stream) + .collect(), + )?; + if snapshot_token.is_before_or_eq(&request.to_room_key) { + return Ok(changes); + } + + let deltas = reader + .membership_delta_changes( + &request.user_id, + &request.to_room_key, + &snapshot_token, + &request.rooms_to_exclude, + membership_changed_since_to, + ) + .await?; + + // The first change after the token in each room. + let mut first_changes = IndexMap::new(); + for delta in deltas { + first_changes + .entry(Arc::clone(&delta.room_id)) + .or_insert(delta); + } + + for (room_id, first) in first_changes { + if first.prev_event_id.is_none() { + changes.insert(room_id, None); + continue; + } + // The previous event may have been culled from the database. + let (Some(prev_event_pos), Some(prev_membership), Some(prev_sender)) = ( + first.prev_event_pos, + first.prev_membership, + first.prev_sender, + ) else { + changes.insert(room_id, None); + continue; + }; + let room_version_id = match rooms.get(&room_id) { + Some(current) => current.room_version_id.clone(), + None => reader.room_version_id(&room_id).await?, + }; + changes.insert( + Arc::clone(&room_id), + Some(MembershipChange { + room_id, + sender: Some(prev_sender), + membership: prev_membership, + event_id: first.prev_event_id, + event_pos: prev_event_pos, + room_version_id, + }), + ); + } + Ok(changes) +} + +/// The positions the events stream change cache has for `room_ids`, and +/// the earliest it knows: `get_max_pos_of_last_change` and +/// `get_earliest_known_position`, read under the GIL. +fn stream_cache_positions( + cache: &Py, + room_ids: &[Arc], +) -> anyhow::Result<(i64, HashMap, i64>)> { + Python::attach(|py| -> PyResult<_> { + let cache = cache.bind(py); + let earliest: i64 = cache + .call_method0(pyo3::intern!(py, "get_earliest_known_position"))? + .extract()?; + let method = cache.getattr(pyo3::intern!(py, "get_max_pos_of_last_change"))?; + let mut positions = HashMap::with_capacity(room_ids.len()); + for room_id in room_ids { + let position: Option = method.call1((&**room_id,))?.extract()?; + if let Some(position) = position { + positions.insert(Arc::clone(room_id), position); + } + } + Ok((earliest, positions)) + }) + .map_err(|e| anyhow::anyhow!("reading the events stream cache: {e}")) +} + +/// `sort_rooms`, with a `limit`. +async fn sort_rooms<'a>( + reader: &SlidingSyncReader<'_>, + rooms: &[&'a RoomForUser], + to_token: &RoomStreamToken, + limit: usize, + snapshot: &ListsSnapshot, +) -> anyhow::Result> { + let joined = SortPlan::joined_rooms(rooms); + let (earliest, positions) = stream_cache_positions(&snapshot.events_stream_cache, &joined)?; + let mut plan = SortPlan::new(rooms, earliest, |room_id| positions.get(room_id).copied()); + for room_id in plan.rooms_to_recheck(Some(limit), to_token.stream) { + let last = reader.last_event_before(&room_id, to_token).await?; + plan.apply_recheck(&room_id, last.map(|(_, position)| position.stream)); + } + if let Some(sorted) = plan.early_result(rooms, Some(limit)) { + return Ok(sorted); + } + let joined_positions = reader + .last_event_pos_before(&joined, to_token, &snapshot.room_max_token, &positions) + .await?; + plan.sort_all(rooms, &joined_positions) +} + +/// `_compute_interested_rooms_new_tables`. Fails with +/// [`NotImplementedFilter`](crate::sliding_sync::room_lists::filter::NotImplementedFilter) +/// where Python raises `NotImplementedError`, and with +/// [`UnknownPosition`](crate::sliding_sync::room_lists::lists::UnknownPosition) +/// for a connection with too many rooms to catch up on. +pub async fn compute_interested_rooms( + reader: &SlidingSyncReader<'_>, + request: &SlidingSyncRequestData, + previous: &PerConnectionState, + snapshot: ListsSnapshot, + now_ms: i64, +) -> anyhow::Result { + let body = &request.body; + let body_lists = body.lists.as_ref().filter(|lists| !lists.is_empty()); + let subscriptions = body + .room_subscriptions + .as_ref() + .filter(|subscriptions| !subscriptions.is_empty()); + if body_lists.is_none() && subscriptions.is_none() { + return Ok(InterestedRooms::default()); + } + + let user_id = &*request.user_id; + let to_token = &request.to_room_key; + + // The current memberships, less the rooms the user left themselves, + // plus those they left after the `to_token`, which the rewind may + // bring back. + let mut rooms = OverlayMap::new(reader.rooms_from_snapshots(user_id).await?); + for (_, room) in reader.self_leave_rooms_after(user_id, to_token).await? { + rooms.insert(room); + } + + // Invites from ignored users. As in Python, the invite rules only + // apply when the user ignores someone. + let ignored_users = reader.ignored_users(user_id).await?; + if !ignored_users.is_empty() { + let invite_rules = InviteRules::load(reader, user_id, snapshot.msc4155_enabled).await?; + for room_id in rooms.keys() { + let Some(room) = rooms.get(&room_id) else { + continue; + }; + let ignored = &*room.membership == "invite" + && room.sender.as_deref().is_some_and(|sender| { + !sender.is_empty() + && (ignored_users.contains(sender) + || invite_rules.invite_rule(sender) == InviteRule::Ignore) + }); + if ignored { + rooms.remove(&room_id); + } + } + } + + let mut newly_joined = HashSet::new(); + let mut newly_left: IndexMap, MembershipChange> = IndexMap::new(); + if let Some(from) = &request.from { + let changes = reader + .snapshot_membership_changes( + user_id, + &from.room_key, + to_token, + &request.rooms_to_exclude, + snapshot.membership_changed_since_from, + ) + .await?; + for (room_id, change) in changes { + match &*change.membership { + "join" => { + newly_joined.insert(room_id); + } + "leave" => { + newly_left.insert(room_id, change); + } + _ => {} + } + } + } + + let changes = rewind_changes( + reader, + request, + &rooms, + snapshot.membership_changed_since_to, + ) + .await?; + for (room_id, change) in &changes { + let Some(change) = change else { + rooms.remove(room_id); + continue; + }; + if let Some(existing) = rooms.get(room_id) { + let room = rewound(change, existing); + keep_if_covered(&mut rooms, user_id, room, newly_left.contains_key(room_id)); + } + } + + // Add back the rooms left in the token range, which the snapshots + // leave out when the user left themselves. + let missing_newly_left: Vec> = newly_left + .keys() + .filter(|room_id| !rooms.contains_key(room_id)) + .cloned() + .collect(); + for room_id in missing_newly_left { + let newly_left_room = &newly_left[&room_id]; + if &*newly_left_room.membership != "leave" { + bail!( + "newly left room {room_id} has membership {}", + newly_left_room.membership + ); + } + match reader.snapshot_room(user_id, &room_id).await? { + // A normal leave. + Some(snapshot_room) => { + keep_if_covered(&mut rooms, user_id, snapshot_room.clone(), true); + if let Some(Some(change)) = changes.get(&room_id) { + let room = rewound(change, &snapshot_room); + keep_if_covered(&mut rooms, user_id, room, true); + } + } + // No membership: a state reset out of the room, whose state + // the tables no longer hold. + None => { + let (has_known_state, room_type) = match reader.room_type(&room_id).await? { + RoomType::Unknown => (false, None), + RoomType::Known(room_type) => (true, room_type), + }; + let at = RoomStreamToken::live(newly_left_room.event_pos.stream); + let is_encrypted = reader + .is_encrypted_at(&room_id, &at, &snapshot.room_max_token) + .await?; + let room = RoomForUser { + room_id: Arc::clone(&room_id), + sender: newly_left_room.sender.clone(), + membership: Arc::clone(&newly_left_room.membership), + event_id: newly_left_room.event_id.clone(), + event_pos: newly_left_room.event_pos.clone(), + room_version_id: newly_left_room.room_version_id.clone(), + has_known_state, + room_type, + is_encrypted, + }; + keep_if_covered(&mut rooms, user_id, room, true); + } + } + } + + for room_id in request.rooms_to_exclude.iter() { + rooms.remove(room_id); + } + + let direct = reader.global_account_data_by_type(user_id, DIRECT).await?; + let dm_room_ids = dm_room_ids(direct.as_deref()); + + let mut lists = Vec::new(); + let mut relevant_room_map = RelevantRoomMap::new(); + let mut all_rooms: HashSet> = HashSet::new(); + + if let Some(body_lists) = body_lists { + let sync_rooms: Vec<&RoomForUser> = rooms.values().collect(); + let mut room_tags: Option>> = None; + for (list_key, list) in body_lists.iter() { + let mut filtered = match &list.filters { + Some(filters) => { + if (filters.tags.is_some() || filters.not_tags.is_some()) && room_tags.is_none() + { + let tags = reader.tags_for_user(user_id).await?; + room_tags = Some( + tags.iter() + .map(|(room_id, tags)| { + (room_id.to_string(), tags.keys().cloned().collect()) + }) + .collect(), + ); + } + let no_tags = HashMap::new(); + filter_rooms_using_tables( + &sync_rooms, + filters, + &dm_room_ids, + room_tags.as_ref().unwrap_or(&no_tags), + |room_id| !previous.rooms.have_sent_room(room_id).is_never(), + )? + } + None => sync_rooms.clone(), + }; + + let partial_state_rooms = reader.partial_rooms().await?; + let room_sync_config = + RoomSyncConfig::from_room_config(list.timeline_limit, &list.required_state); + if room_sync_config.must_await_full_state(&snapshot.server_name)? { + filtered.retain(|room| !partial_state_rooms.contains(&room.room_id)); + } + + all_rooms.extend(filtered.iter().map(|room| Arc::clone(&room.room_id))); + + let mut ops = Vec::new(); + if let Some(ranges) = list.ranges.as_deref().filter(|ranges| !ranges.is_empty()) { + let sorted = if needs_sort(ranges, filtered.len()) { + sort_rooms(reader, &filtered, to_token, sort_limit(ranges), &snapshot).await? + } else { + filtered.clone() + }; + let sorted_ids: Vec> = sorted + .iter() + .map(|room| Arc::clone(&room.room_id)) + .collect(); + ops = sync_list_ops( + ranges, + &sorted_ids, + &room_sync_config, + &mut relevant_room_map, + ); + } + lists.push((list_key.to_owned(), filtered.len(), ops)); + } + } + + if let Some(subscriptions) = subscriptions { + let partial_state_rooms = reader.partial_rooms().await?; + let missing: Vec> = subscriptions + .iter() + .map(|(room_id, _)| room_id) + .filter(|room_id| !rooms.contains_key(room_id)) + .map(Arc::from) + .collect(); + for (_, room) in reader.snapshot_rooms(user_id, &missing).await? { + rooms.insert(room); + } + for (room_id, subscription) in subscriptions.iter() { + // TODO, as in Python: rooms the user is not in. + if !rooms.contains_key(room_id) { + continue; + } + let room_id: Arc = Arc::from(room_id); + all_rooms.insert(Arc::clone(&room_id)); + let room_sync_config = RoomSyncConfig::from_room_config( + subscription.timeline_limit, + &subscription.required_state, + ); + if room_sync_config.must_await_full_state(&snapshot.server_name)? + && partial_state_rooms.contains(&room_id) + { + continue; + } + add_room_subscription(&mut relevant_room_map, &room_id, room_sync_config); + } + } + + let rooms_to_send = match &request.from { + Some(from) if !relevant_room_map.is_empty() => { + let relevant: Vec> = relevant_room_map.keys().cloned().collect(); + let rooms_that_have_updates = reader + .rooms_with_updates_since( + &relevant, + &from.room_key, + snapshot.events_changed_since_from.as_ref(), + ) + .await?; + rooms_to_send( + &relevant_room_map, + previous, + &rooms_that_have_updates, + snapshot.num_rooms_threshold, + now_ms, + )? + } + _ => relevant_room_map.keys().cloned().collect(), + }; + + let memberships = relevant_room_map + .keys() + .filter_map(|room_id| rooms.get(room_id)) + .map(|room| (Arc::clone(&room.room_id), room.clone())) + .collect(); + + Ok(InterestedRooms { + has_lists_or_subscriptions: true, + lists, + relevant_room_map, + rooms_to_send, + all_rooms, + memberships, + newly_joined, + newly_left: newly_left.into_keys().collect(), + dm_room_ids, + events_changed_since_from: snapshot.events_changed_since_from, + }) +} + +/// A room's `RoomsForUserSlidingSync` for Python: `(sender, membership, +/// event_id, instance_name, stream, room_version_id, has_known_state, +/// room_type, is_encrypted)`. +type PyMembership<'a> = ( + Option<&'a str>, + &'a str, + Option<&'a str>, + &'a str, + i64, + Option<&'a str>, + bool, + Option<&'a str>, + bool, +); + +/// One room's input to Python's `get_room_sync_data`: `(room_id, +/// (timeline_limit, required_state_map), membership, newly_joined, +/// newly_left, is_dm)`. +type PyRoomInput<'a> = ( + &'a str, + (i64, BTreeMap>), + PyMembership<'a>, + bool, + bool, + bool, +); + +/// `SlidingSyncInterestedRooms`, as Python reads it. +#[pyclass(frozen, skip_from_py_object, name = "InterestedRooms")] +pub struct PyInterestedRooms { + pub(crate) data: Arc, + pub(crate) stats: SlidingSyncStats, +} + +#[pymethods] +impl PyInterestedRooms { + /// `(name, count, [((start, end), room_ids)])` for each list, in the + /// request's order. + #[getter] + #[allow(clippy::type_complexity)] + fn lists(&self) -> Vec<(&str, usize, Vec<((u64, u64), Vec<&str>)>)> { + self.data + .lists + .iter() + .map(|(name, count, ops)| { + let ops = ops + .iter() + .map(|op| (op.range, op.room_ids.iter().map(|r| &**r).collect())) + .collect(); + (name.as_str(), *count, ops) + }) + .collect() + } + + /// The keys of `relevant_room_map`. + #[getter] + fn relevant_room_ids(&self) -> Vec<&str> { + self.data.relevant_room_map.keys().map(|r| &**r).collect() + } + + /// The keys of `relevant_rooms_to_send_map`. + #[getter] + fn rooms_to_send(&self) -> Vec<&str> { + self.data.rooms_to_send.iter().map(|r| &**r).collect() + } + + #[getter] + fn all_rooms(&self) -> HashSet<&str> { + self.data.all_rooms.iter().map(|r| &**r).collect() + } + + #[getter] + fn stats(&self) -> SlidingSyncStats { + self.stats + } + + /// What Python's `get_room_sync_data` takes for each of `room_ids`, for + /// the rooms Rust leaves to Python. + fn python_room_inputs<'a>(&'a self, room_ids: Vec) -> PyResult>> { + let data = &*self.data; + room_ids + .iter() + .map(|room_id| { + let (room_id, config) = data + .relevant_room_map + .get_key_value(room_id.as_str()) + .ok_or_else(|| PyKeyError::new_err(room_id.clone()))?; + let room = data + .memberships + .get(room_id) + .ok_or_else(|| PyKeyError::new_err(room_id.to_string()))?; + let membership = ( + room.sender.as_deref(), + &*room.membership, + room.event_id.as_deref(), + &*room.event_pos.instance_name, + room.event_pos.stream, + room.room_version_id.as_deref(), + room.has_known_state, + room.room_type.as_deref(), + room.is_encrypted, + ); + Ok(( + &**room_id, + (config.timeline_limit, config.required_state_map.clone()), + membership, + data.newly_joined.contains(room_id), + data.newly_left.contains(room_id), + data.dm_room_ids.contains(&**room_id), + )) + }) + .collect() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::storage::sync_store::PersistedEventPosition; + + fn room(room_id: &str, membership: &str) -> RoomForUser { + RoomForUser { + room_id: Arc::from(room_id), + sender: Some(Arc::from("@u:test")), + membership: Arc::from(membership), + event_id: Some(Arc::from("$e")), + event_pos: PersistedEventPosition { + instance_name: Arc::from("master"), + stream: 1, + }, + room_version_id: Some(Arc::from("10")), + has_known_state: true, + room_type: None, + is_encrypted: false, + } + } + + #[test] + fn overlay_iterates_changes_first() { + let base: SnapshotRooms = ["!a", "!b", "!c"] + .into_iter() + .map(|id| (Arc::from(id), room(id, "join"))) + .collect(); + let mut rooms = OverlayMap::new(Arc::new(base)); + rooms.insert(room("!b", "leave")); + rooms.remove("!a"); + rooms.insert(room("!d", "join")); + let ids: Vec = rooms.keys().iter().map(|id| id.to_string()).collect(); + assert_eq!(ids, ["!b", "!d", "!c"]); + assert_eq!(&*rooms.get("!b").unwrap().membership, "leave"); + assert!(rooms.get("!a").is_none()); + // Set again after a pop, the room goes to the end. + rooms.remove("!b"); + rooms.insert(room("!b", "join")); + let ids: Vec = rooms.keys().iter().map(|id| id.to_string()).collect(); + assert_eq!(ids, ["!d", "!b", "!c"]); + } + + #[test] + fn msc4155_rules() { + let rules = InviteRules::msc4155(&serde_json::json!({ + "allowed_users": ["@friend:evil.example"], + "ignored_servers": ["*.example"], + "blocked_users": [5, "@bad:*"], + })); + assert_eq!(rules.invite_rule("@friend:evil.example"), InviteRule::Allow); + assert_eq!(rules.invite_rule("@other:evil.example"), InviteRule::Ignore); + assert_eq!(rules.invite_rule("@bad:good.org"), InviteRule::Block); + assert_eq!(rules.invite_rule("@ok:good.org"), InviteRule::Allow); + } + + #[test] + fn dm_rooms_ignore_bad_shapes() { + let direct = serde_json::json!({"@a:x": ["!1", 2, "!2"], "@b:x": "!3"}); + let mut rooms: Vec = dm_room_ids(Some(&direct)).into_iter().collect(); + rooms.sort(); + assert_eq!(rooms, ["!1", "!2"]); + assert!(dm_room_ids(Some(&serde_json::json!([]))).is_empty()); + } +} diff --git a/rust/src/sliding_sync/room_lists/lists.rs b/rust/src/sliding_sync/room_lists/lists.rs index ce5c37d715..691c7c60cb 100644 --- a/rust/src/sliding_sync/room_lists/lists.rs +++ b/rust/src/sliding_sync/room_lists/lists.rs @@ -30,6 +30,7 @@ pub const MINIMUM_NOT_USED_AGE_EXPIRY_MS: i64 = 60 * 60 * 1000; /// `NUM_ROOMS_THRESHOLD`: more rooms with updates than this, on a /// connection unused for [`MINIMUM_NOT_USED_AGE_EXPIRY_MS`], resets it. +/// Python passes its value per request, which tests patch. pub const NUM_ROOMS_THRESHOLD: usize = 100; /// The rooms a request covers, to their combined configs, in the order @@ -145,6 +146,7 @@ pub fn rooms_to_send( relevant_room_map: &RelevantRoomMap, previous_state: &PerConnectionState, rooms_that_have_updates: &HashSet>, + num_rooms_threshold: usize, now_ms: i64, ) -> Result>, UnknownPosition> { let mut rooms_should_send: HashSet<&Arc> = HashSet::new(); @@ -167,7 +169,7 @@ pub fn rooms_to_send( // A client with lots to catch up on does better starting over with a // small range, unless it synced recently and would loop. - if rooms_that_have_updates.len() > NUM_ROOMS_THRESHOLD { + if rooms_that_have_updates.len() > num_rooms_threshold { if let Some(last_used_ts) = previous_state.last_used_ts { if now_ms - last_used_ts > MINIMUM_NOT_USED_AGE_EXPIRY_MS { return Err(UnknownPosition); diff --git a/rust/src/sliding_sync/room_lists/mod.rs b/rust/src/sliding_sync/room_lists/mod.rs index 615bc5e25d..3948438325 100644 --- a/rust/src/sliding_sync/room_lists/mod.rs +++ b/rust/src/sliding_sync/room_lists/mod.rs @@ -31,6 +31,7 @@ use crate::storage::sync_store::PersistedEventPosition; use crate::types::stream_token::RoomStreamToken; pub mod filter; +pub mod interested; pub mod lists; pub mod sort; @@ -210,9 +211,15 @@ fn rooms_to_send_for_test( ..Default::default() }; let updates: HashSet> = rooms_that_have_updates.into_iter().map(Arc::from).collect(); - lists::rooms_to_send(&relevant, &state, &updates, now_ms) - .map(|rooms| rooms.iter().map(|r| r.to_string()).collect()) - .map_err(|_| PyValueError::new_err("unknown position")) + lists::rooms_to_send( + &relevant, + &state, + &updates, + lists::NUM_ROOMS_THRESHOLD, + now_ms, + ) + .map(|rooms| rooms.iter().map(|r| r.to_string()).collect()) + .map_err(|_| PyValueError::new_err("unknown position")) } pub fn register_test_functions(m: &Bound<'_, PyModule>) -> PyResult<()> { diff --git a/rust/src/sliding_sync/state.rs b/rust/src/sliding_sync/state.rs new file mode 100644 index 0000000000..b6614f5f31 --- /dev/null +++ b/rust/src/sliding_sync/state.rs @@ -0,0 +1,740 @@ +/* + * 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: + * . + */ + +//! A request's connection state, owned by Rust from load to persist. +//! +//! [`ConnectionState`] holds the state loaded at the request's position +//! and the request's updates to it. Python's extensions read and update it +//! through [`ConnectionStateView`] and [`RoomStatusMapView`], which have +//! the method surface of Python's `PerConnectionState` and +//! `MutablePerConnectionState` as far as the extensions use it. +//! +//! The receipts stream's tokens stay in their database form, a string: +//! Rust never compares them, and the view parses one into a Python +//! `MultiWriterStreamToken` only when Python asks for it. + +use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; +use std::sync::{Arc, Mutex}; + +use anyhow::Context; +use pyo3::{ + exceptions::{PyRuntimeError, PyValueError}, + intern, + prelude::*, + types::{PyDict, PyTuple}, +}; + +use crate::caches::lock_ignoring_poison; +use crate::sliding_sync::config::RoomSyncConfig; +use crate::sliding_sync::connection::{ + HaveSentRoom, MutablePerConnectionState, PerConnectionState, RoomLazyMembershipChanges, + RoomStatusMap, +}; +use crate::sliding_sync::result::required_state_map_from_py; +use crate::sliding_sync::SlidingSyncStats; +use crate::storage::sliding_sync_store::{ConnectionStateUpdates, StoredConnectionState}; +use crate::types::stream_token::RoomStreamToken; + +/// The `instance_map` table: each writer's instance ID and name. +#[derive(Debug, Clone, Default)] +pub struct InstanceMap { + by_id: HashMap>, + by_name: HashMap, i64>, +} + +impl InstanceMap { + pub fn new(pairs: impl IntoIterator)>) -> Self { + let mut map = InstanceMap::default(); + for (id, name) in pairs { + map.by_name.insert(Arc::clone(&name), id); + map.by_id.insert(id, name); + } + map + } + + /// `get_name_from_instance_id`. + pub fn name(&self, id: i64) -> Option> { + self.by_id.get(&id).cloned() + } + + /// `get_id_for_instance`, for a writer the table already has. + pub fn id(&self, name: &str) -> Option { + self.by_name.get(name).copied() + } + + /// Name to ID, as the room computation formats tokens with. + pub fn ids_by_name(&self) -> HashMap { + self.by_name + .iter() + .map(|(name, id)| (name.to_string(), *id)) + .collect() + } +} + +/// The streams of a connection state. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Stream { + Rooms, + Receipts, + AccountData, +} + +/// Apply `f` to a status's token. +fn map_status( + status: &HaveSentRoom, + f: impl FnOnce(&T) -> Result, +) -> Result, E> { + Ok(match status { + HaveSentRoom::Never => HaveSentRoom::Never, + HaveSentRoom::Live => HaveSentRoom::Live, + HaveSentRoom::Previously(token) => HaveSentRoom::Previously(f(token)?), + }) +} + +fn map_statuses( + statuses: &HashMap>, + f: impl Fn(&T) -> anyhow::Result, +) -> anyhow::Result>> { + statuses + .iter() + .map(|(room_id, status)| Ok((room_id.clone(), map_status(status, &f)?))) + .collect() +} + +/// A `MultiWriterStreamToken`'s database form as `(stream, [(instance ID, +/// position)])`: the bare stream, or `m{stream}~{id}.{pos}~...`. +fn parse_multi_writer(string: &str) -> anyhow::Result<(i64, Vec<(i64, i64)>)> { + let invalid = || anyhow::anyhow!("invalid stream token {string:?}"); + if string.starts_with(|c: char| c.is_ascii_digit()) { + return Ok((string.parse().map_err(|_| invalid())?, Vec::new())); + } + let rest = string.strip_prefix('m').ok_or_else(invalid)?; + let mut parts = rest.split('~'); + let stream = parts + .next() + .ok_or_else(invalid)? + .parse() + .map_err(|_| invalid())?; + let mut positions = Vec::new(); + for part in parts { + // Tokens of the form `m5~` were created by a bug. + if part.is_empty() { + continue; + } + let (id, position) = part.split_once('.').ok_or_else(invalid)?; + positions.push(( + id.parse().map_err(|_| invalid())?, + position.parse().map_err(|_| invalid())?, + )); + } + Ok((stream, positions)) +} + +/// A Python `MultiWriterStreamToken` from its database form. +fn multi_writer_to_py<'py>( + py: Python<'py>, + string: &str, + instances: &InstanceMap, +) -> PyResult> { + let (stream, positions) = + parse_multi_writer(string).map_err(|e| PyValueError::new_err(e.to_string()))?; + let instance_map = PyDict::new(py); + for (id, position) in positions { + let name = instances + .name(id) + .ok_or_else(|| PyValueError::new_err(format!("unknown instance ID {id}")))?; + instance_map.set_item(&*name, position)?; + } + let immutabledict = py + .import(intern!(py, "immutabledict"))? + .getattr(intern!(py, "immutabledict"))? + .call1((instance_map,))?; + let kwargs = PyDict::new(py); + kwargs.set_item(intern!(py, "stream"), stream)?; + kwargs.set_item(intern!(py, "instance_map"), immutabledict)?; + py.import(intern!(py, "synapse.types"))? + .getattr(intern!(py, "MultiWriterStreamToken"))? + .call((), Some(&kwargs)) +} + +/// The database form of a Python `MultiWriterStreamToken`, as its +/// `to_string` writes it: writers at or behind `stream` are left out. +fn multi_writer_from_py(token: &Bound<'_, PyAny>, instances: &InstanceMap) -> PyResult { + let py = token.py(); + let stream: i64 = token.getattr(intern!(py, "stream"))?.extract()?; + let mut entries = Vec::new(); + for item in token + .getattr(intern!(py, "instance_map"))? + .call_method0(intern!(py, "items"))? + .try_iter()? + { + let (name, position): (String, i64) = item?.extract()?; + if position <= stream { + continue; + } + let id = instances + .id(&name) + .ok_or_else(|| PyValueError::new_err(format!("unknown instance name {name}")))?; + entries.push(format!("{id}.{position}")); + } + if entries.is_empty() { + Ok(stream.to_string()) + } else { + Ok(format!("m{stream}~{}", entries.join("~"))) + } +} + +/// The loaded state and the request's updates. +pub struct ConnectionStateData { + pub previous: PerConnectionState, + pub new: MutablePerConnectionState, +} + +impl ConnectionStateData { + /// The state as stored, its tokens parsed. An empty `stored` is the + /// state of a connection's first request. + pub fn from_stored( + stored: &StoredConnectionState, + instances: &InstanceMap, + ) -> anyhow::Result { + let rooms = map_statuses(&stored.rooms, |token| { + RoomStreamToken::parse_with_instances(token, |id| instances.name(id)) + .map_err(anyhow::Error::new) + })?; + let account_data = map_statuses(&stored.account_data, |token| { + token + .parse::() + .with_context(|| format!("account data token {token:?}")) + })?; + let previous = PerConnectionState { + last_used_ts: stored.last_used_ts, + rooms: RoomStatusMap::new(rooms), + receipts: RoomStatusMap::new(stored.receipts.clone()), + account_data: RoomStatusMap::new(account_data), + room_configs: Arc::new(stored.room_configs.clone()), + }; + let new = previous.get_mutable(); + Ok(ConnectionStateData { previous, new }) + } + + /// The updates to persist, each token in its database form. + pub fn updates(&self, instances: &InstanceMap) -> anyhow::Result { + Ok(ConnectionStateUpdates { + rooms: map_statuses(self.new.rooms.get_updates(), |token| { + token + .to_string_with_instances(|name| instances.id(name)) + .map_err(anyhow::Error::new) + })?, + receipts: self.new.receipts.get_updates().clone(), + account_data: map_statuses(self.new.account_data.get_updates(), |token| { + Ok(token.to_string()) + })?, + room_configs: self.new.get_room_config_updates().clone(), + room_lazy_membership: self.new.room_lazy_membership.clone(), + }) + } +} + +/// A request's connection state. Made by +/// `SlidingSyncHandler.load_connection_state`. +#[pyclass(frozen, skip_from_py_object)] +pub struct ConnectionState { + pub(crate) data: Arc>, + pub(crate) instances: Arc, + /// What loading the state cost. + load_stats: SlidingSyncStats, +} + +impl ConnectionState { + pub fn new( + data: ConnectionStateData, + instances: InstanceMap, + load_stats: SlidingSyncStats, + ) -> Self { + ConnectionState { + data: Arc::new(Mutex::new(data)), + instances: Arc::new(instances), + load_stats, + } + } +} + +/// A Python `HaveSentRoom`. +fn have_sent_room_to_py<'py, T>( + py: Python<'py>, + status: &HaveSentRoom, + token_to_py: impl FnOnce(&T) -> PyResult>, +) -> PyResult> { + let class = py + .import(intern!(py, "synapse.types.handlers.sliding_sync"))? + .getattr(intern!(py, "HaveSentRoom"))?; + match status { + HaveSentRoom::Never => class.call_method0(intern!(py, "never")), + HaveSentRoom::Live => class.call_method0(intern!(py, "live")), + HaveSentRoom::Previously(token) => { + class.call_method1(intern!(py, "previously"), (token_to_py(token)?,)) + } + } +} + +/// A room config as Python's glue builds a `RoomSyncConfig` from it: +/// `(timeline_limit, {type: {state keys}})`. +type PyRoomConfig = (i64, BTreeMap>); + +/// A room's `RoomLazyMembershipChanges` for Python: `(room_id, +/// returned_user_id_to_last_seen_ts_map, invalidated_user_ids)`. +type PyLazyMembership = (String, HashMap>, HashSet); + +fn config_to_py(config: &RoomSyncConfig) -> PyRoomConfig { + (config.timeline_limit, config.required_state_map.clone()) +} + +#[pymethods] +impl ConnectionState { + /// What loading the state cost. + #[getter] + fn stats(&self) -> SlidingSyncStats { + self.load_stats + } + + /// The state as loaded: `PerConnectionState`. + #[getter(previous)] + fn previous_view(slf: &Bound<'_, Self>) -> ConnectionStateView { + ConnectionStateView { + state: slf.clone().unbind(), + mutable: false, + } + } + + /// The state with the request's updates: `MutablePerConnectionState`. + #[getter(new)] + fn new_view(slf: &Bound<'_, Self>) -> ConnectionStateView { + ConnectionStateView { + state: slf.clone().unbind(), + mutable: true, + } + } + + /// What Python's `get_room_sync_data` reads of the state for `room_ids`, + /// for the rooms Rust leaves to Python: `(last_used_ts, {room_id: + /// (timeline_limit, required_state_map)}, {room_id: HaveSentRoom})`. + #[allow(clippy::type_complexity)] + fn python_state_for_rooms<'py>( + &self, + py: Python<'py>, + room_ids: Vec, + ) -> PyResult<( + Option, + HashMap, + HashMap>, + )> { + let data = lock_ignoring_poison(&self.data); + let previous = &data.previous; + let mut configs = HashMap::new(); + let mut statuses = HashMap::new(); + for room_id in room_ids { + if let Some(config) = previous.room_configs.get(&room_id) { + configs.insert(room_id.clone(), config_to_py(config)); + } + let status = previous.rooms.have_sent_room(&room_id); + if !status.is_never() { + let status = have_sent_room_to_py(py, &status, |token| token.into_pyobject(py))?; + statuses.insert(room_id, status); + } + } + Ok((previous.last_used_ts, configs, statuses)) + } + + /// Record what Python's `get_room_sync_data` changed for the rooms it + /// computed: `room_configs` maps a room to `(timeline_limit, + /// required_state_map)`, and `lazy_memberships` a room to + /// `(returned_user_id_to_last_seen_ts_map, invalidated_user_ids)`. + fn apply_python_rooms( + &self, + room_configs: Vec<(String, i64, Bound<'_, PyAny>)>, + lazy_memberships: Vec, + ) -> PyResult<()> { + let mut configs = Vec::with_capacity(room_configs.len()); + for (room_id, timeline_limit, map) in room_configs { + let config = RoomSyncConfig { + timeline_limit, + required_state_map: required_state_map_from_py(&map)?, + }; + configs.push((room_id, config)); + } + let mut data = lock_ignoring_poison(&self.data); + for (room_id, config) in configs { + data.new.set_room_config(&room_id, config); + } + for (room_id, returned, invalidated) in lazy_memberships { + data.new.room_lazy_membership.insert( + room_id, + RoomLazyMembershipChanges { + returned_user_id_to_last_seen_ts_map: returned, + invalidated_user_ids: invalidated, + }, + ); + } + Ok(()) + } + + /// The request's room config and lazy membership updates, for tests: + /// `({room_id: (timeline_limit, required_state_map)}, {room_id: + /// (returned_user_id_to_last_seen_ts_map, invalidated_user_ids)})`. + #[allow(clippy::type_complexity)] + fn updates_for_test( + &self, + ) -> ( + HashMap, + HashMap>, HashSet)>, + ) { + let data = lock_ignoring_poison(&self.data); + let configs = data + .new + .get_room_config_updates() + .iter() + .map(|(room_id, config)| (room_id.clone(), config_to_py(config))) + .collect(); + let lazy = data + .new + .room_lazy_membership + .iter() + .map(|(room_id, changes)| { + ( + room_id.clone(), + ( + changes.returned_user_id_to_last_seen_ts_map.clone(), + changes.invalidated_user_ids.clone(), + ), + ) + }) + .collect(); + (configs, lazy) + } +} + +/// `PerConnectionState` (`mutable` false) or `MutablePerConnectionState` +/// over a [`ConnectionState`], as the extensions read them. +#[pyclass(frozen, skip_from_py_object)] +pub struct ConnectionStateView { + state: Py, + mutable: bool, +} + +impl ConnectionStateView { + fn stream(&self, py: Python<'_>, stream: Stream) -> RoomStatusMapView { + RoomStatusMapView { + state: self.state.clone_ref(py), + stream, + mutable: self.mutable, + } + } +} + +#[pymethods] +impl ConnectionStateView { + #[getter] + fn rooms(&self, py: Python<'_>) -> RoomStatusMapView { + self.stream(py, Stream::Rooms) + } + + #[getter] + fn receipts(&self, py: Python<'_>) -> RoomStatusMapView { + self.stream(py, Stream::Receipts) + } + + #[getter] + fn account_data(&self, py: Python<'_>) -> RoomStatusMapView { + self.stream(py, Stream::AccountData) + } + + #[getter] + fn last_used_ts(&self, py: Python<'_>) -> Option { + let state = self.state.bind(py).get(); + let data = lock_ignoring_poison(&state.data); + if self.mutable { + data.new.last_used_ts + } else { + data.previous.last_used_ts + } + } +} + +/// One stream's `RoomStatusMap` (or `MutableRoomStatusMap`) of a +/// [`ConnectionStateView`]. +#[pyclass(frozen, skip_from_py_object)] +pub struct RoomStatusMapView { + state: Py, + stream: Stream, + mutable: bool, +} + +impl RoomStatusMapView { + /// The statuses of `room_ids`, or every room's when `None`, as Python + /// `HaveSentRoom`s. + fn statuses<'py>( + &self, + py: Python<'py>, + room_ids: Option<&[String]>, + ) -> PyResult)>> { + let state = self.state.bind(py).get(); + let instances = &state.instances; + let data = lock_ignoring_poison(&state.data); + // The flattened map, for every room, or one lookup per room. + macro_rules! collect { + ($previous:expr, $new:expr, $to_py:expr) => {{ + let to_py = $to_py; + match room_ids { + Some(room_ids) => room_ids + .iter() + .map(|room_id| { + let status = if self.mutable { + $new.have_sent_room(room_id) + } else { + $previous.have_sent_room(room_id) + }; + Ok((room_id.clone(), have_sent_room_to_py(py, &status, to_py)?)) + }) + .collect::>>(), + None => { + let flattened; + let map = if self.mutable { + flattened = $new.copy(); + &flattened + } else { + &$previous + }; + map.iter() + .map(|(room_id, status)| { + Ok((room_id.to_owned(), have_sent_room_to_py(py, status, to_py)?)) + }) + .collect::>>() + } + } + }}; + } + match self.stream { + Stream::Rooms => collect!( + data.previous.rooms, + data.new.rooms, + |token: &RoomStreamToken| { token.into_pyobject(py) } + ), + Stream::Receipts => collect!( + data.previous.receipts, + data.new.receipts, + |token: &String| { multi_writer_to_py(py, token, instances) } + ), + Stream::AccountData => { + collect!( + data.previous.account_data, + data.new.account_data, + |token: &i64| { Ok(token.into_pyobject(py)?.into_any()) } + ) + } + } + } + + fn check_mutable(&self) -> PyResult<()> { + if self.mutable { + Ok(()) + } else { + Err(PyRuntimeError::new_err( + "the previous connection state is read-only", + )) + } + } +} + +/// Room IDs from any iterable of strings. +fn room_ids_from_py(room_ids: &Bound<'_, PyAny>) -> PyResult> { + room_ids + .try_iter()? + .map(|room_id| room_id?.extract()) + .collect() +} + +#[pymethods] +impl RoomStatusMapView { + fn have_sent_room<'py>(&self, py: Python<'py>, room_id: String) -> PyResult> { + let mut statuses = self.statuses(py, Some(&[room_id]))?; + Ok(statuses.pop().expect("one room asked for").1) + } + + /// Every room's status, as Python's `_statuses` mapping. + #[getter] + fn _statuses<'py>(&self, py: Python<'py>) -> PyResult> { + let dict = PyDict::new(py); + for (room_id, status) in self.statuses(py, None)? { + dict.set_item(room_id, status)?; + } + Ok(dict) + } + + fn record_sent_rooms(&self, py: Python<'_>, room_ids: &Bound<'_, PyAny>) -> PyResult<()> { + self.check_mutable()?; + let room_ids = room_ids_from_py(room_ids)?; + let state = self.state.bind(py).get(); + let mut data = lock_ignoring_poison(&state.data); + let room_ids = room_ids.iter().map(String::as_str); + match self.stream { + Stream::Rooms => data.new.rooms.record_sent_rooms(room_ids), + Stream::Receipts => data.new.receipts.record_sent_rooms(room_ids), + Stream::AccountData => data.new.account_data.record_sent_rooms(room_ids), + } + Ok(()) + } + + fn record_unsent_rooms( + &self, + py: Python<'_>, + room_ids: &Bound<'_, PyAny>, + from_token: &Bound<'_, PyAny>, + ) -> PyResult<()> { + self.check_mutable()?; + let room_ids = room_ids_from_py(room_ids)?; + let state = self.state.bind(py).get(); + let room_ids = room_ids.iter().map(String::as_str); + match self.stream { + Stream::Rooms => { + let token: RoomStreamToken = from_token.extract()?; + let mut data = lock_ignoring_poison(&state.data); + data.new.rooms.record_unsent_rooms(room_ids, &token); + } + Stream::Receipts => { + let token = multi_writer_from_py(from_token, &state.instances)?; + let mut data = lock_ignoring_poison(&state.data); + data.new.receipts.record_unsent_rooms(room_ids, &token); + } + Stream::AccountData => { + let token: i64 = from_token.extract()?; + let mut data = lock_ignoring_poison(&state.data); + data.new.account_data.record_unsent_rooms(room_ids, &token); + } + } + Ok(()) + } + + /// `get_updates`: the request's changes, as `{room_id: HaveSentRoom}`. + fn get_updates<'py>(&self, py: Python<'py>) -> PyResult> { + self.check_mutable()?; + let state = self.state.bind(py).get(); + let room_ids: Vec = { + let data = lock_ignoring_poison(&state.data); + match self.stream { + Stream::Rooms => data.new.rooms.get_updates().keys().cloned().collect(), + Stream::Receipts => data.new.receipts.get_updates().keys().cloned().collect(), + Stream::AccountData => data + .new + .account_data + .get_updates() + .keys() + .cloned() + .collect(), + } + }; + let dict = PyDict::new(py); + for (room_id, status) in self.statuses(py, Some(&room_ids))? { + dict.set_item(room_id, status)?; + } + Ok(dict) + } +} + +/// The database form of a receipts token from a Python +/// `MultiWriterStreamToken`, and back, through the Rust implementation, +/// with `instances` as `[(instance_id, name)]`. For tests. +#[pyfunction] +fn multi_writer_token_round_trip_for_test<'py>( + py: Python<'py>, + token: &Bound<'py, PyAny>, + instances: Vec<(i64, String)>, +) -> PyResult> { + let instances = InstanceMap::new( + instances + .into_iter() + .map(|(id, name)| (id, Arc::from(name))), + ); + let string = multi_writer_from_py(token, &instances)?; + let parsed = multi_writer_to_py(py, &string, &instances)?; + (string, parsed).into_pyobject(py) +} + +pub fn register_module(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_function(wrap_pyfunction!(multi_writer_token_round_trip_for_test, m)?)?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn multi_writer_tokens_parse() { + assert_eq!(parse_multi_writer("5").unwrap(), (5, vec![])); + assert_eq!( + parse_multi_writer("m5~1.7~2.9").unwrap(), + (5, vec![(1, 7), (2, 9)]) + ); + assert_eq!(parse_multi_writer("m5~").unwrap(), (5, vec![])); + assert!(parse_multi_writer("s5").is_err()); + assert!(parse_multi_writer("").is_err()); + } + + #[test] + fn stored_state_round_trips() { + let instances = InstanceMap::new([(1, Arc::from("w1"))]); + let stored = StoredConnectionState { + last_used_ts: Some(3), + rooms: HashMap::from([ + ( + "!a".to_owned(), + HaveSentRoom::Previously("m5~1.7".to_owned()), + ), + ("!b".to_owned(), HaveSentRoom::Live), + ]), + receipts: HashMap::from([( + "!a".to_owned(), + HaveSentRoom::Previously("m5~1.7".to_owned()), + )]), + account_data: HashMap::from([( + "!a".to_owned(), + HaveSentRoom::Previously("4".to_owned()), + )]), + room_configs: HashMap::new(), + }; + let mut data = ConnectionStateData::from_stored(&stored, &instances).unwrap(); + let token = match data.previous.rooms.have_sent_room("!a") { + HaveSentRoom::Previously(token) => token, + other => panic!("{other:?}"), + }; + assert_eq!(token.get_stream_pos_for_instance("w1"), 7); + data.new.rooms.record_unsent_rooms(["!b"], &token); + data.new.account_data.record_sent_rooms(["!a"]); + let updates = data.updates(&instances).unwrap(); + assert_eq!( + updates.rooms, + HashMap::from([( + "!b".to_owned(), + HaveSentRoom::Previously("m5~1.7".to_owned()) + )]) + ); + assert_eq!( + updates.account_data, + HashMap::from([("!a".to_owned(), HaveSentRoom::Live)]) + ); + } +} diff --git a/rust/src/storage/sliding_sync_store.rs b/rust/src/storage/sliding_sync_store.rs index 43de3f87b0..bbfc5e04e4 100644 --- a/rust/src/storage/sliding_sync_store.rs +++ b/rust/src/storage/sliding_sync_store.rs @@ -70,8 +70,10 @@ pub const UPDATE_INTERVAL_LAST_USED_TS_MS: i64 = 5 * 60 * 1000; /// one that fails with a serialization error or a deadlock. const MAX_TRANSACTION_ATTEMPTS: usize = 5; -/// The rooms of a user's membership snapshots, by room ID. -pub type SnapshotRooms = HashMap, RoomForUser>; +/// The rooms of a user's membership snapshots, by room ID, in the order +/// the query returns them, as Python's dict keeps it: a list that covers +/// every room is not sorted, so its rooms come in this order. +pub type SnapshotRooms = IndexMap, RoomForUser>; /// A user's tags: room ID to tag to the tag's content. pub type UserTags = HashMap, HashMap>; @@ -963,6 +965,21 @@ impl<'a> SlidingSyncReader<'a> { } } + /// The `instance_map` table: each writer's instance ID and name, as + /// stream tokens name writers by ID. + pub async fn instance_map(&self) -> anyhow::Result)>> { + let rows = self + .query_one_shot("SELECT instance_id, instance_name FROM instance_map", &[]) + .await?; + rows.iter() + .map(|row| { + let id: i32 = row.try_get(0)?; + let name: String = row.try_get(1)?; + Ok((i64::from(id), Arc::from(name))) + }) + .collect() + } + /// `get_partial_rooms`: every room the server joined with partial /// state, which is never many. pub async fn partial_rooms(&self) -> anyhow::Result>>> { diff --git a/synapse/handlers/sliding_sync/__init__.py b/synapse/handlers/sliding_sync/__init__.py index f79056d3fa..98a86777da 100644 --- a/synapse/handlers/sliding_sync/__init__.py +++ b/synapse/handlers/sliding_sync/__init__.py @@ -24,13 +24,14 @@ from typing_extensions import assert_never from twisted.internet.defer import CancelledError from synapse.api.constants import Direction, EventTypes, Membership +from synapse.api.errors import SlidingSyncUnknownPosition from synapse.events import EventBase from synapse.events.utils import FilteredEvent, strip_event from synapse.handlers.relations import BundledAggregations +from synapse.handlers.sliding_sync import room_lists from synapse.handlers.sliding_sync.extensions import SlidingSyncExtensionHandler from synapse.handlers.sliding_sync.room_lists import ( RoomsForUserType, - SlidingSyncInterestedRooms, SlidingSyncRoomLists, ) from synapse.handlers.sliding_sync.store import SlidingSyncConnectionStore @@ -48,9 +49,12 @@ from synapse.storage.databases.main.state_deltas import StateDelta from synapse.storage.databases.main.stream import PaginateFunction from synapse.storage.roommember import ( MemberSummary, + RoomsForUserSlidingSync, ) from synapse.synapse_rust.sliding_sync import ( - RoomInput as RustRoomInput, + ConnectionState as RustConnectionState, + InterestedRooms as RustInterestedRooms, + RustSlidingSyncFallback, SlidingSyncHandler as RustSlidingSyncHandler, SlidingSyncRequest as RustSlidingSyncRequest, SlidingSyncStats as RustSlidingSyncStats, @@ -74,8 +78,10 @@ from synapse.types.handlers.sliding_sync import ( HaveSentRoom, HaveSentRoomFlag, MutablePerConnectionState, + OperationType, PerConnectionState, RoomLazyMembershipChanges, + RoomStatusMap, RoomSyncConfig, SlidingSyncConfig, SlidingSyncResult, @@ -108,14 +114,15 @@ rust_sliding_sync_rooms_counter = Counter( ) # Requests served in `rust_sliding_sync: on` mode. `result` is "rust" for a -# request Rust took part in, "fallback" for one Python computed in full because -# Rust does not cover it (the sliding sync background updates have not finished, -# or a list uses a filter Python does not implement either), or "error" for one +# request Rust served (with the extensions, and any rooms Rust left to Python, +# computed in Python), "fallback" for one Python computed in full because Rust +# does not cover it (the sliding sync background updates have not finished, or a +# list uses a filter Python does not implement either), or "error" for one # Python computed in full after the Rust handler failed on it. rust_sliding_sync_requests_counter = Counter( "synapse_sliding_sync_rust_requests_total", "Sliding sync requests in `rust_sliding_sync: on` mode, by whether Rust " - "took part or Python computed the whole response", + "served them or Python computed the whole response", labelnames=["result", SERVER_NAME_LABEL], ) @@ -173,6 +180,19 @@ def observe_rust_sliding_sync_stats( ) +def _sum_stats(stats: Iterable[RustSlidingSyncStats]) -> RustSlidingSyncStats: + """One request's stats, from those of its Rust calls.""" + stats = list(stats) + return RustSlidingSyncStats( + db_time=sum(s.db_time for s in stats), + checkout_time=sum(s.checkout_time for s in stats), + transactions=sum(s.transactions for s in stats), + queries=sum(s.queries for s in stats), + rooms_computed=sum(s.rooms_computed for s in stats), + lazy_member_lookups=sum(s.lazy_member_lookups for s in stats), + ) + + # Limit the number of state_keys we should remember sending down the connection for each # (room_id, user_id). We don't want to store and pull out too much data in the database. # @@ -412,6 +432,22 @@ class SlidingSyncHandler: # See https://github.com/matrix-org/matrix-doc/issues/1144 raise NotImplementedError() + if self._rust_sliding_sync is not None: + rust_result = await self._current_sync_for_user_in_rust( + self._rust_sliding_sync, + sync_config, + to_token=to_token, + from_token=from_token, + ) + if rust_result is not None: + set_tag(SynapseTags.RESULT_PREFIX + "result", bool(rust_result)) + set_tag(SynapseTags.FUNC_ARG_PREFIX + "sync_config.user", user_id) + sync_processing_time.labels( + initial=from_token is not None, + **{SERVER_NAME_LABEL: self.server_name}, + ).observe(self.clock.time() - start_time_s) + return rust_result + # Get the per-connection state (if any). # # Raises an exception if there is a `connection_position` that we don't @@ -478,20 +514,9 @@ class SlidingSyncHandler: if relevant_rooms_to_send_map: with start_active_span("sliding_sync.generate_room_entries"): - python_room_ids: StrCollection = list(relevant_rooms_to_send_map) - if self._rust_sliding_sync is not None: - python_room_ids = await self._generate_room_entries_in_rust( - self._rust_sliding_sync, - sync_config=sync_config, - previous_connection_state=previous_connection_state, - new_connection_state=new_connection_state, - interested_rooms=interested_rooms, - from_token=from_token, - to_token=to_token, - rooms=rooms, - ) - if python_room_ids: - await concurrently_execute(handle_room, python_room_ids, 20) + await concurrently_execute( + handle_room, list(relevant_rooms_to_send_map), 20 + ) extensions = await self.extensions.get_extensions_response( sync_config=sync_config, @@ -585,109 +610,83 @@ class SlidingSyncHandler: return sliding_sync_result - async def _generate_room_entries_in_rust( + async def _current_sync_for_user_in_rust( self, rust: RustSlidingSyncHandler, - *, sync_config: SlidingSyncConfig, - previous_connection_state: PerConnectionState, - new_connection_state: MutablePerConnectionState, - interested_rooms: SlidingSyncInterestedRooms, - from_token: SlidingSyncStreamToken | None, to_token: StreamToken, - rooms: dict[str, AnySlidingSyncRoomResult], - ) -> list[str]: - """`get_room_sync_data` in Rust for every room of - `relevant_rooms_to_send_map`. Adds the entries to `rooms` and the - rooms' changes to `new_connection_state`. + from_token: SlidingSyncStreamToken | None, + ) -> SlidingSyncResult | None: + """`current_sync_for_user` with the room lists, the rooms and the + connection state in Rust, and the extensions in Python. Returns: - The rooms Python must compute: those Rust leaves to Python, and - those it failed on (outside strict mode). All of them if the Rust - call as a whole failed. + None for a request Python must compute: one Rust does not cover, + or one Rust failed on before persisting the connection state + (outside strict mode). + + Raises: + SlidingSyncUnknownPosition if the connection position is unknown, or + the connection has too much to catch up on. """ - rooms_to_send = interested_rooms.relevant_rooms_to_send_map labels = {SERVER_NAME_LABEL: self.server_name} + store = self.store + user_id = sync_config.user.to_string() + + # Rust only reads the sliding sync tables. + if not await store.have_finished_sliding_sync_background_jobs(): + rust_sliding_sync_requests_counter.labels(result="fallback", **labels).inc() + return None try: request = await self._build_rust_request(sync_config, from_token, to_token) except ValueError as e: rust_sliding_sync_requests_counter.labels(result="fallback", **labels).inc() - logger.debug("Computing the sliding sync rooms in Python: %s", e) - return list(rooms_to_send) + logger.debug("Computing the sliding sync request in Python: %s", e) + return None - # The writers that the rooms' `prev_batch` tokens may name: those of - # the tokens the timelines are bounded by. - instance_names = set(to_token.room_key.instance_map) + # The writers the tokens Rust formats and parses may name. + tokens = [to_token.room_key, to_token.receipt_key] if from_token is not None: - instance_names.update(from_token.stream_token.room_key.instance_map) - - room_inputs = [] - for room_id, room_sync_config in rooms_to_send.items(): - membership = interested_rooms.room_membership_for_user_map[room_id] - prev_room_sync_config = previous_connection_state.room_configs.get(room_id) - room_status = previous_connection_state.rooms.have_sent_room(room_id) - if room_status.last_token is not None: - instance_names.update(room_status.last_token.instance_map) - room_inputs.append( - RustRoomInput( - room_id=room_id, - sender=membership.sender, - membership=membership.membership, - event_id=membership.event_id, - instance_name=membership.event_pos.instance_name, - stream=membership.event_pos.stream, - room_version_id=membership.room_version_id, - timeline_limit=room_sync_config.timeline_limit, - required_state_map=room_sync_config.required_state_map, - prev_room_sync_config=( - ( - prev_room_sync_config.timeline_limit, - prev_room_sync_config.required_state_map, - ) - if prev_room_sync_config is not None - else None - ), - room_status=room_status.status.value, - room_status_token=room_status.last_token, - newly_joined=room_id in interested_rooms.newly_joined_rooms, - newly_left=room_id in interested_rooms.newly_left_rooms, - is_dm=room_id in interested_rooms.dm_room_ids, - ) - ) + tokens += [ + from_token.stream_token.room_key, + from_token.stream_token.receipt_key, + ] instance_ids = { - name: await self.store.get_id_for_instance(name) for name in instance_names + name: await store.get_id_for_instance(name) + for token in tokens + for name in token.instance_map } + stats: list[RustSlidingSyncStats] = [] try: - result = await rust.generate_room_entries( - request, room_inputs, instance_ids=instance_ids - ) - except CancelledError: + with start_active_span("sliding_sync.rust.load_connection_state"): + state = await rust.load_connection_state( + request, instance_ids=instance_ids + ) + stats.append(state.stats) + with start_active_span("sliding_sync.rust.compute_interested_rooms"): + interested = await self._compute_interested_rooms_in_rust( + rust, request, state, user_id, to_token, from_token + ) + stats.append(interested.stats) + rooms: dict[str, AnySlidingSyncRoomResult] = {} + with start_active_span("sliding_sync.rust.generate_room_entries"): + result = await rust.generate_room_entries(request, state, interested) + stats.append(result.stats) + except (CancelledError, SlidingSyncUnknownPosition): raise + except RustSlidingSyncFallback as e: + rust_sliding_sync_requests_counter.labels(result="fallback", **labels).inc() + logger.debug("Computing the sliding sync request in Python: %s", e) + return None except Exception: rust_sliding_sync_requests_counter.labels(result="error", **labels).inc() if self._rust_sliding_sync_strict: raise logger.exception( - "Rust sliding sync failed for %s. Computing the rooms in Python", - sync_config.user.to_string(), - ) - return list(rooms_to_send) - - rust_sliding_sync_requests_counter.labels(result="rust", **labels).inc() - observe_rust_sliding_sync_stats(result.stats, self.server_name) - - for room_id, timeline_limit, required_state_map in result.room_config_updates: - new_connection_state.room_configs[room_id] = RoomSyncConfig( - timeline_limit=timeline_limit, required_state_map=required_state_map - ) - for room_id, returned, invalidated in result.lazy_memberships: - new_connection_state.room_lazy_membership[room_id] = ( - RoomLazyMembershipChanges( - returned_user_id_to_last_seen_ts_map=returned, - invalidated_user_ids=invalidated, - ) + "Rust sliding sync failed for %s. Computing it in Python", user_id ) + return None entries = result.take_entries() for entry in entries: @@ -714,7 +713,210 @@ class SlidingSyncHandler: error, ) python_room_ids.append(room_id) - return python_room_ids + if python_room_ids: + await self._generate_room_entries_in_python( + sync_config, + state, + interested, + python_room_ids, + from_token=from_token, + to_token=to_token, + rooms=rooms, + ) + + lists = { + name: SlidingSyncResult.SlidingWindowList( + count=count, + ops=[ + SlidingSyncResult.SlidingWindowList.Operation( + op=OperationType.SYNC, range=range, room_ids=room_ids + ) + for range, room_ids in ops + ], + ) + for name, count, ops in interested.lists + } + extensions = await self.extensions.get_extensions_response( + sync_config=sync_config, + actual_lists=lists, + previous_connection_state=state.previous, + new_connection_state=state.new, + all_interested_room_ids=interested.all_rooms, + # All the rooms we could send, with event updates or not: see the + # Python path. + actual_room_ids=set(interested.relevant_room_ids), + actual_room_response_map=rooms, + from_token=from_token, + to_token=to_token, + ) + + try: + with start_active_span("sliding_sync.rust.finish_request"): + connection_position, finish_stats = await rust.finish_request( + request, state, interested + ) + except CancelledError: + raise + except Exception: + # The persist is one transaction, so nothing was stored, and + # Python can compute the request from the start. + rust_sliding_sync_requests_counter.labels(result="error", **labels).inc() + if self._rust_sliding_sync_strict: + raise + logger.exception( + "Rust sliding sync failed for %s. Computing it in Python", user_id + ) + return None + stats.append(finish_stats) + + rust_sliding_sync_requests_counter.labels(result="rust", **labels).inc() + observe_rust_sliding_sync_stats(_sum_stats(stats), self.server_name) + return SlidingSyncResult( + next_pos=SlidingSyncStreamToken(to_token, connection_position), + lists=lists, + rooms=rooms, + extensions=extensions, + ) + + async def _compute_interested_rooms_in_rust( + self, + rust: RustSlidingSyncHandler, + request: RustSlidingSyncRequest, + state: RustConnectionState, + user_id: str, + to_token: StreamToken, + from_token: SlidingSyncStreamToken | None, + ) -> RustInterestedRooms: + """`compute_interested_rooms` in Rust, with the stream change caches' + answers for the request.""" + store = self.store + membership_changed_since_from = True + events_changed_since_from: set[str] | None = None + if from_token is not None: + from_stream = from_token.stream_token.room_key.stream + membership_changed_since_from = ( + store._membership_stream_cache.has_entity_changed(user_id, from_stream) + ) + changed = store._events_stream_cache.get_all_entities_changed(from_stream) + if changed.hit: + events_changed_since_from = set(changed.entities) + return await rust.compute_interested_rooms( + request, + state, + server_name=self.server_name, + msc4155_enabled=self._hs_config.experimental.msc4155_enabled, + # Read per request, as tests patch it. + num_rooms_threshold=room_lists.NUM_ROOMS_THRESHOLD, + membership_changed_since_from=membership_changed_since_from, + membership_changed_since_to=store._membership_stream_cache.has_entity_changed( + user_id, to_token.room_key.stream + ), + events_changed_since_from=events_changed_since_from, + room_max_token=store.get_room_max_token(), + events_stream_cache=store._events_stream_cache, + ) + + async def _generate_room_entries_in_python( + self, + sync_config: SlidingSyncConfig, + state: RustConnectionState, + interested: RustInterestedRooms, + room_ids: list[str], + *, + from_token: SlidingSyncStreamToken | None, + to_token: StreamToken, + rooms: dict[str, AnySlidingSyncRoomResult], + ) -> None: + """`get_room_sync_data` in Python for the rooms Rust left to it, over + a Python copy of their slice of the connection state. Their changes go + back into `state`.""" + last_used_ts, room_configs, statuses = state.python_state_for_rooms(room_ids) + previous_connection_state = PerConnectionState( + last_used_ts=last_used_ts, + rooms=RoomStatusMap(statuses), + room_configs={ + room_id: RoomSyncConfig( + timeline_limit=timeline_limit, required_state_map=required_state_map + ) + for room_id, ( + timeline_limit, + required_state_map, + ) in room_configs.items() + }, + ) + new_connection_state = previous_connection_state.get_mutable() + + inputs = { + room_input[0]: room_input + for room_input in interested.python_room_inputs(room_ids) + } + + async def handle_room(room_id: str) -> None: + ( + _, + (timeline_limit, required_state_map), + membership, + newly_joined, + newly_left, + is_dm, + ) = inputs[room_id] + ( + sender, + membership_value, + event_id, + instance_name, + stream, + room_version_id, + has_known_state, + room_type, + is_encrypted, + ) = membership + # The tables only hold rooms of known versions. + assert room_version_id is not None, f"{room_id} has no room version" + room_sync_result = await self.get_room_sync_data( + sync_config=sync_config, + previous_connection_state=previous_connection_state, + new_connection_state=new_connection_state, + room_id=room_id, + room_sync_config=RoomSyncConfig( + timeline_limit=timeline_limit, required_state_map=required_state_map + ), + room_membership_for_user_at_to_token=RoomsForUserSlidingSync( + room_id=room_id, + sender=sender, + membership=membership_value, + event_id=event_id, + event_pos=PersistedEventPosition(instance_name, stream), + room_version_id=room_version_id, + has_known_state=has_known_state, + room_type=room_type, + is_encrypted=is_encrypted, + ), + from_token=from_token, + to_token=to_token, + newly_joined=newly_joined, + newly_left=newly_left, + is_dm=is_dm, + ) + if room_sync_result or not from_token: + rooms[room_id] = room_sync_result + + await concurrently_execute(handle_room, room_ids, 20) + + state.apply_python_rooms( + [ + (room_id, config.timeline_limit, config.required_state_map) + for room_id, config in new_connection_state.room_configs.maps[0].items() + ], + [ + ( + room_id, + changes.returned_user_id_to_last_seen_ts_map, + changes.invalidated_user_ids, + ) + for room_id, changes in new_connection_state.room_lazy_membership.items() + ], + ) @trace async def get_current_state_ids_at( diff --git a/synapse/handlers/sliding_sync/extensions.py b/synapse/handlers/sliding_sync/extensions.py index f009b0d330..7176dbef6e 100644 --- a/synapse/handlers/sliding_sync/extensions.py +++ b/synapse/handlers/sliding_sync/extensions.py @@ -37,6 +37,7 @@ from synapse.events.utils import FilteredEvent from synapse.handlers.receipts import ReceiptEventSource from synapse.logging.opentracing import trace from synapse.storage.databases.main.receipts import ReceiptInRoom +from synapse.synapse_rust.sliding_sync import ConnectionStateView from synapse.types import ( Absent, AbsentType, @@ -78,6 +79,14 @@ _ThreadUnsubscription: TypeAlias = ( if TYPE_CHECKING: from synapse.server import HomeServer +# The connection state of a request: Python's, or the Rust-owned one in +# `rust_sliding_sync: on` mode, which has the same surface as far as the +# extensions read and update it. +AnyPerConnectionState: TypeAlias = PerConnectionState | ConnectionStateView +AnyMutablePerConnectionState: TypeAlias = ( + MutablePerConnectionState | ConnectionStateView +) + logger = logging.getLogger(__name__) @@ -99,8 +108,8 @@ class SlidingSyncExtensionHandler: async def get_extensions_response( self, sync_config: SlidingSyncConfig, - previous_connection_state: "PerConnectionState", - new_connection_state: "MutablePerConnectionState", + previous_connection_state: AnyPerConnectionState, + new_connection_state: AnyMutablePerConnectionState, all_interested_room_ids: set[str], actual_lists: Mapping[str, SlidingSyncResult.SlidingWindowList], actual_room_ids: set[str], @@ -455,8 +464,8 @@ class SlidingSyncExtensionHandler: async def get_account_data_extension_response( self, sync_config: SlidingSyncConfig, - previous_connection_state: "PerConnectionState", - new_connection_state: "MutablePerConnectionState", + previous_connection_state: AnyPerConnectionState, + new_connection_state: AnyMutablePerConnectionState, actual_lists: Mapping[str, SlidingSyncResult.SlidingWindowList], actual_room_ids: set[str], account_data_request: SlidingSyncConfig.Extensions.AccountDataExtension, @@ -687,8 +696,8 @@ class SlidingSyncExtensionHandler: async def get_receipts_extension_response( self, sync_config: SlidingSyncConfig, - previous_connection_state: "PerConnectionState", - new_connection_state: "MutablePerConnectionState", + previous_connection_state: AnyPerConnectionState, + new_connection_state: AnyMutablePerConnectionState, actual_lists: Mapping[str, SlidingSyncResult.SlidingWindowList], actual_room_ids: set[str], actual_room_response_map: Mapping[str, AnySlidingSyncRoomResult], diff --git a/synapse/synapse_rust/sliding_sync.pyi b/synapse/synapse_rust/sliding_sync.pyi index 537b0efba0..b0a5a72d5b 100644 --- a/synapse/synapse_rust/sliding_sync.pyi +++ b/synapse/synapse_rust/sliding_sync.pyi @@ -13,64 +13,214 @@ """Simplified sliding sync (MSC4186) in Rust. See `rust/src/sliding_sync/`.""" from collections.abc import Mapping -from typing import AbstractSet +from typing import AbstractSet, Generic, TypeVar from twisted.internet.defer import Deferred from synapse.events import EventBase from synapse.synapse_rust.database.postgres import ConnectionPool from synapse.synapse_rust.events import SerializeEventConfig -from synapse.types import JsonDict, JsonMapping, RoomStreamToken +from synapse.types import ( + JsonDict, + JsonMapping, + MultiWriterStreamToken, + RoomStreamToken, + StrCollection, +) +from synapse.types.handlers.sliding_sync import HaveSentRoom from synapse.types.state import StateFilter +from synapse.util.caches.stream_change_cache import StreamChangeCache + +T = TypeVar("T", str, RoomStreamToken, MultiWriterStreamToken, int) class SlidingSyncHandler: """Computes sliding sync responses in Rust, reading the database through the native Postgres driver. Built by `RustHandlers`.""" + def load_connection_state( + self, request: SlidingSyncRequest, *, instance_ids: Mapping[str, int] + ) -> Deferred[ConnectionState]: + """`get_and_clear_connection_positions`. `instance_ids` are + `get_id_for_instance` for the writers of the request's tokens. + + Raises: + SlidingSyncUnknownPosition for a position the store does not know. + """ + def load_connection_state_blocking( + self, request: SlidingSyncRequest, *, instance_ids: Mapping[str, int] + ) -> ConnectionState: ... + def compute_interested_rooms( + self, + request: SlidingSyncRequest, + state: ConnectionState, + *, + server_name: str, + msc4155_enabled: bool, + num_rooms_threshold: int, + membership_changed_since_from: bool, + membership_changed_since_to: bool, + events_changed_since_from: AbstractSet[str] | None, + room_max_token: RoomStreamToken, + events_stream_cache: StreamChangeCache, + ) -> Deferred[InterestedRooms]: + """`compute_interested_rooms` on the sliding sync tables. + `membership_changed_since_*` are `_membership_stream_cache`'s + `has_entity_changed` at the from and to tokens' streams (true without a + from token), `events_changed_since_from` is `_events_stream_cache`'s + `get_all_entities_changed` at the from token's stream, and + `events_stream_cache` is read for the rooms a sort needs. + + Raises: + RustSlidingSyncFallback for a filter Python does not implement. + SlidingSyncUnknownPosition for a connection to reset. + """ + def compute_interested_rooms_blocking( + self, + request: SlidingSyncRequest, + state: ConnectionState, + *, + server_name: str, + msc4155_enabled: bool, + num_rooms_threshold: int, + membership_changed_since_from: bool, + membership_changed_since_to: bool, + events_changed_since_from: AbstractSet[str] | None, + room_max_token: RoomStreamToken, + events_stream_cache: StreamChangeCache, + ) -> InterestedRooms: ... def generate_room_entries( self, request: SlidingSyncRequest, - rooms: list[RoomInput], - *, - instance_ids: Mapping[str, int], + state: ConnectionState, + interested: InterestedRooms, ) -> Deferred[RoomsResult]: - """`get_room_sync_data` for each of `rooms`, twenty at a time. - `instance_ids` are `get_id_for_instance` for the writers the rooms' - tokens may name.""" + """`get_room_sync_data` for each of the rooms to send, twenty at a + time. The rooms' changes to the connection state go into `state`.""" def generate_room_entries_blocking( self, request: SlidingSyncRequest, - rooms: list[RoomInput], - *, - instance_ids: Mapping[str, int], + state: ConnectionState, + interested: InterestedRooms, ) -> RoomsResult: ... - -class RoomInput: - """The input of one room's `get_room_sync_data`.""" - - def __init__( + def finish_request( self, - *, - room_id: str, - sender: str | None, - membership: str, - event_id: str | None, - instance_name: str, - stream: int, - room_version_id: str | None, - timeline_limit: int, - required_state_map: Mapping[str, AbstractSet[str]], - prev_room_sync_config: tuple[int, Mapping[str, AbstractSet[str]]] | None, - room_status: str, - room_status_token: RoomStreamToken | None, - newly_joined: bool, - newly_left: bool, - is_dm: bool, - ) -> None: - """`room_status` is a `HaveSentRoomFlag` value, with its token for - "previously".""" + request: SlidingSyncRequest, + state: ConnectionState, + interested: InterestedRooms, + ) -> Deferred[tuple[int, SlidingSyncStats]]: + """The rest of `current_sync_for_user` after the extensions: record + the sent and unsent rooms, and persist the state if it changed. + Resolves to the connection position for `pos`.""" + def finish_request_blocking( + self, + request: SlidingSyncRequest, + state: ConnectionState, + interested: InterestedRooms, + ) -> tuple[int, SlidingSyncStats]: ... + +class RustSlidingSyncFallback(Exception): + """The request is one Rust leaves to Python.""" + +class ConnectionState: + """A request's connection state, owned by Rust from load to persist.""" + @property - def room_id(self) -> str: ... + def stats(self) -> SlidingSyncStats: + """What loading the state cost.""" + @property + def previous(self) -> ConnectionStateView: + """The state as loaded: `PerConnectionState`.""" + @property + def new(self) -> ConnectionStateView: + """The state with the request's updates: `MutablePerConnectionState`.""" + def python_state_for_rooms( + self, room_ids: list[str] + ) -> tuple[ + int | None, + dict[str, tuple[int, dict[str, set[str]]]], + dict[str, HaveSentRoom[RoomStreamToken]], + ]: + """`(last_used_ts, room_configs, room_statuses)` of the previous state + for `room_ids`, for the rooms Python computes.""" + def apply_python_rooms( + self, + room_configs: list[tuple[str, int, Mapping[str, AbstractSet[str]]]], + lazy_memberships: list[tuple[str, Mapping[str, int | None], AbstractSet[str]]], + ) -> None: + """Record the changes of the rooms Python computed.""" + def updates_for_test( + self, + ) -> tuple[ + dict[str, tuple[int, dict[str, set[str]]]], + dict[str, tuple[dict[str, int | None], set[str]]], + ]: ... + +class ConnectionStateView: + """`PerConnectionState` or `MutablePerConnectionState` over a + `ConnectionState`, as the extensions read them.""" + + @property + def rooms(self) -> RoomStatusMapView[RoomStreamToken]: ... + @property + def receipts(self) -> RoomStatusMapView[MultiWriterStreamToken]: ... + @property + def account_data(self) -> RoomStatusMapView[int]: ... + @property + def last_used_ts(self) -> int | None: ... + +class RoomStatusMapView(Generic[T]): + """One stream's `RoomStatusMap` of a `ConnectionStateView`. The record + methods raise on the previous state's.""" + + def have_sent_room(self, room_id: str) -> HaveSentRoom[T]: ... + @property + def _statuses(self) -> dict[str, HaveSentRoom[T]]: ... + def record_sent_rooms(self, room_ids: StrCollection) -> None: ... + def record_unsent_rooms(self, room_ids: StrCollection, from_token: T) -> None: ... + def get_updates(self) -> dict[str, HaveSentRoom[T]]: ... + +class InterestedRooms: + """`SlidingSyncInterestedRooms`, as Python reads it.""" + + @property + def lists( + self, + ) -> list[tuple[str, int, list[tuple[tuple[int, int], list[str]]]]]: + """`(name, count, [((start, end), room_ids)])` for each list.""" + @property + def relevant_room_ids(self) -> list[str]: ... + @property + def rooms_to_send(self) -> list[str]: ... + @property + def all_rooms(self) -> set[str]: ... + @property + def stats(self) -> SlidingSyncStats: ... + def python_room_inputs( + self, room_ids: list[str] + ) -> list[ + tuple[ + str, + tuple[int, dict[str, set[str]]], + tuple[ + str | None, + str, + str | None, + str, + int, + str | None, + bool, + str | None, + bool, + ], + bool, + bool, + bool, + ] + ]: + """`(room_id, (timeline_limit, required_state_map), (sender, + membership, event_id, instance_name, stream, room_version_id, + has_known_state, room_type, is_encrypted), newly_joined, newly_left, + is_dm)` for each room, for the rooms Python computes.""" class RoomResult: """One room's entry, computed in Rust: what `current_sync_for_user` and the @@ -103,16 +253,6 @@ class RoomsResult: def empty(self) -> list[str]: """The rooms computed with nothing to send.""" @property - def room_config_updates(self) -> list[tuple[str, int, dict[str, set[str]]]]: - """`(room_id, timeline_limit, required_state_map)` for each room whose - config the connection records.""" - @property - def lazy_memberships( - self, - ) -> list[tuple[str, dict[str, int | None], set[str]]]: - """`(room_id, returned_user_id_to_last_seen_ts_map, - invalidated_user_ids)` for each room computed.""" - @property def fallbacks(self) -> list[tuple[str, str]]: ... @property def errors(self) -> list[tuple[str, str]]: ... @@ -472,3 +612,9 @@ class SlidingSyncStoreForTest: room_lazy_membership: dict[str, tuple[dict[str, int | None], set[str]]], now_ms: int, ) -> int: ... + +def multi_writer_token_round_trip_for_test( + token: MultiWriterStreamToken, instances: list[tuple[int, str]] +) -> tuple[str, MultiWriterStreamToken]: + """A receipts token's database form, as Rust writes it, and the token Rust + parses back from it.""" diff --git a/tests/rest/client/sliding_sync/test_rust_lists.py b/tests/rest/client/sliding_sync/test_rust_lists.py new file mode 100644 index 0000000000..64fdbf464b --- /dev/null +++ b/tests/rest/client/sliding_sync/test_rust_lists.py @@ -0,0 +1,514 @@ +# +# 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 room lists, room subscriptions and connection state of a sliding sync +request in Rust (`rust/src/sliding_sync/room_lists/interested.rs`, +`state.rs`) against Python's. + +Each request goes to two connections of the same user, one served by Python +and one by Rust. The two responses must be equal but for `pos`, and so must +the changes each side makes to its connection's state: the room configs, the +lazy-loaded members, and the status of each room in the rooms, receipts and +account data streams. +""" + +from typing import Any, Mapping +from unittest import SkipTest + +from immutabledict import immutabledict + +from twisted.internet.defer import Deferred +from twisted.internet.testing import MemoryReactor + +import synapse.rest.admin +from synapse.api.constants import AccountDataTypes, EventTypes, RoomTypes +from synapse.handlers.sliding_sync import rust_sliding_sync_requests_counter +from synapse.metrics import SERVER_NAME_LABEL +from synapse.rest.client import account_data, knock, login, receipts, room, sync, tags +from synapse.server import HomeServer +from synapse.synapse_rust.sliding_sync import ( + ConnectionState as RustConnectionState, + InterestedRooms as RustInterestedRooms, + SlidingSyncHandler as RustSlidingSyncHandler, + SlidingSyncRequest as RustSlidingSyncRequest, + SlidingSyncStats as RustSlidingSyncStats, + multi_writer_token_round_trip_for_test, +) +from synapse.types import JsonDict, MultiWriterStreamToken, SlidingSyncStreamToken +from synapse.types.handlers.sliding_sync import ( + MutablePerConnectionState, + SlidingSyncConfig, +) +from synapse.util.clock import Clock + +from tests.rest.client.sliding_sync.test_rust_rooms import _normalized +from tests.rest.client.sliding_sync.test_sliding_sync import SlidingSyncBase +from tests.server import _BlockingRustCalls +from tests.unittest import TestCase + + +def _changes( + configs: Mapping[str, Any], + lazy: Mapping[str, Any], + streams: Mapping[str, Mapping[str, Any]], +) -> JsonDict: + """A request's changes to its connection state, comparable across the two + sides. A status is its flag and whether it has a token, as each connection + has its own tokens. A lazy member's timestamp is only whether there is + one, as the two connections send members a moment apart.""" + return { + "room_configs": { + room_id: (limit, {t: set(keys) for t, keys in state_map.items()}) + for room_id, (limit, state_map) in configs.items() + }, + "lazy_membership": { + room_id: ( + {user_id: ts is not None for user_id, ts in returned.items()}, + set(invalidated), + ) + for room_id, (returned, invalidated) in lazy.items() + }, + "streams": { + stream: { + room_id: (status.status, status.last_token is not None) + for room_id, status in updates.items() + } + for stream, updates in streams.items() + }, + } + + +def _python_changes(state: MutablePerConnectionState) -> JsonDict: + return _changes( + { + room_id: (config.timeline_limit, config.required_state_map) + for room_id, config in state.get_room_config_updates().items() + }, + { + room_id: ( + changes.returned_user_id_to_last_seen_ts_map, + changes.invalidated_user_ids, + ) + for room_id, changes in state.room_lazy_membership.items() + }, + { + "rooms": state.rooms.get_updates(), + "receipts": state.receipts.get_updates(), + "account_data": state.account_data.get_updates(), + }, + ) + + +def _rust_changes(state: RustConnectionState) -> JsonDict: + configs, lazy = state.updates_for_test() + new = state.new + return _changes( + configs, + lazy, + { + "rooms": new.rooms.get_updates(), + "receipts": new.receipts.get_updates(), + "account_data": new.account_data.get_updates(), + }, + ) + + +class _RecordingRustCalls(_BlockingRustCalls): + """Blocking Rust calls that record each request's changes to its + connection state, by `conn_id`, once the request is finished.""" + + def __init__( + self, handler: RustSlidingSyncHandler, recorded: dict[str, JsonDict] + ) -> None: + super().__init__(handler) + self._recorded = recorded + + def finish_request( + self, + request: RustSlidingSyncRequest, + state: RustConnectionState, + interested: RustInterestedRooms, + ) -> "Deferred[tuple[int, RustSlidingSyncStats]]": + result = super().__getattr__("finish_request")(request, state, interested) + conn_id = request.body_for_test().get("conn_id") + assert isinstance(conn_id, str) + self._recorded[conn_id] = _rust_changes(state) + return result + + +class SlidingSyncRustListsTestCase(SlidingSyncBase): + maxDiff = None + + servlets = [ + synapse.rest.admin.register_servlets, + account_data.register_servlets, + knock.register_servlets, + login.register_servlets, + receipts.register_servlets, + room.register_servlets, + sync.register_servlets, + tags.register_servlets, + ] + + def default_config(self) -> JsonDict: + config = super().default_config() + config.setdefault("experimental_features", {})["msc4155_enabled"] = True + return config + + def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None: + super().prepare(reactor, clock, hs) + rust = hs.get_rust_handlers().sliding_sync + if rust is None: + raise SkipTest("needs the native Rust Postgres driver") + self.recorded: dict[str, JsonDict] = {} + self.rust = _RecordingRustCalls(rust, self.recorded) + self.handler = hs.get_sliding_sync_handler() + self.handler._rust_sliding_sync_strict = True + self.pos: dict[str, str | None] = {"python": None, "rust": None} + + connection_store = self.handler.connection_store + record_new_state = connection_store.record_new_state + + async def recording_record_new_state( + sync_config: SlidingSyncConfig, + from_token: SlidingSyncStreamToken | None, + new_connection_state: MutablePerConnectionState, + ) -> int: + assert sync_config.conn_id is not None + self.recorded[sync_config.conn_id] = _python_changes(new_connection_state) + return await record_new_state(sync_config, from_token, new_connection_state) + + connection_store.record_new_state = recording_record_new_state # type: ignore[method-assign] + + self.user1_id = self.register_user("user1", "pass") + self.user1_tok = self.login(self.user1_id, "pass") + self.user2_id = self.register_user("user2", "pass") + self.user2_tok = self.login(self.user2_id, "pass") + + def _rust_requests(self) -> float: + return rust_sliding_sync_requests_counter.labels( + result="rust", **{SERVER_NAME_LABEL: self.hs.hostname} + )._value.get() + + def _sync_both( + self, sync_body: JsonDict, unordered_lists: bool = False + ) -> JsonDict: + """Sync `sync_body` on the Python and the Rust connection, assert the + two responses and the two sides' changes to the connection state are + equal, and return Rust's response. + + With `unordered_lists`, each operation's rooms compare as a set: Python + filters a list into a set, so a list it does not sort comes out in + hash order.""" + responses = {} + self.recorded.clear() + for side in ("python", "rust"): + self.handler._rust_sliding_sync = ( + self.rust if side == "rust" else None # type: ignore[assignment] + ) + before = self._rust_requests() + response, self.pos[side] = self.do_sync( + {**sync_body, "conn_id": side}, + since=self.pos[side], + tok=self.user1_tok, + ) + self.assertEqual(self._rust_requests() - before, int(side == "rust")) + response.pop("pos") + response["rooms"] = _normalized(response.get("rooms", {})) + if unordered_lists: + for list_result in response.get("lists", {}).values(): + for op in list_result.get("ops", []): + op["room_ids"] = set(op["room_ids"]) + responses[side] = response + self.assertEqual(responses["python"], responses["rust"]) + self.assertEqual(self.recorded.get("python"), self.recorded.get("rust")) + return responses["rust"] + + def _send_message(self, room_id: str) -> str: + return self.helper.send(room_id, "hi", tok=self.user2_tok)["event_id"] + + def _put_account_data(self, data_type: str, content: JsonDict) -> None: + channel = self.make_request( + "PUT", + f"/_matrix/client/v3/user/{self.user1_id}/account_data/{data_type}", + content=content, + access_token=self.user1_tok, + ) + self.assertEqual(channel.code, 200, channel.json_body) + + def _tag(self, room_id: str, tag: str) -> None: + channel = self.make_request( + "PUT", + f"/_matrix/client/v3/user/{self.user1_id}/rooms/{room_id}/tags/{tag}", + content={}, + access_token=self.user1_tok, + ) + self.assertEqual(channel.code, 200, channel.json_body) + + def _joined_rooms(self, count: int) -> list[str]: + rooms = [] + for _ in range(count): + room_id = self.helper.create_room_as(self.user2_id, tok=self.user2_tok) + self.helper.join(room_id, self.user1_id, tok=self.user1_tok) + rooms.append(room_id) + return rooms + + def test_sorted_ranges_and_unsent_rooms(self) -> None: + """A list cut to a range sorts by the latest event, rooms out of range + with new events are recorded as unsent, and widening the range sends + them.""" + rooms = self._joined_rooms(6) + for room_id in reversed(rooms): + self._send_message(room_id) + sync_body: JsonDict = { + "lists": { + "top": {"ranges": [[0, 1]], "required_state": [], "timeline_limit": 2}, + "all": { + "ranges": [[0, 99]], + "required_state": [["m.room.name", ""]], + "timeline_limit": 0, + }, + } + } + self._sync_both(sync_body) + + # Activity in rooms out of the top range, and in one in it. + self._send_message(rooms[4]) + self._send_message(rooms[5]) + self._sync_both({"lists": {"top": sync_body["lists"]["top"]}}) + + # Rooms with updates the connection did not send come down now. + self._send_message(rooms[2]) + self._sync_both( + { + "lists": { + "top": {**sync_body["lists"]["top"], "ranges": [[0, 4]]}, + } + } + ) + # Nothing new. + self._sync_both( + {"lists": {"top": {**sync_body["lists"]["top"], "ranges": [[0, 4]]}}} + ) + + def test_filters(self) -> None: + """The `is_dm`, `is_invite`, `is_encrypted`, room type and tag filters, + over rooms with known and unknown state.""" + joined = self._joined_rooms(2) + dm = self._create_dm_room( + inviter_user_id=self.user2_id, + inviter_tok=self.user2_tok, + invitee_user_id=self.user1_id, + invitee_tok=self.user1_tok, + ) + space = self.helper.create_room_as( + self.user2_id, + tok=self.user2_tok, + extra_content={"creation_content": {"type": RoomTypes.SPACE}}, + ) + self.helper.join(space, self.user1_id, tok=self.user1_tok) + encrypted = self.helper.create_room_as(self.user2_id, tok=self.user2_tok) + self.helper.send_state( + encrypted, + EventTypes.RoomEncryption, + {"algorithm": "m.megolm.v1.aes-sha2"}, + tok=self.user2_tok, + ) + self.helper.join(encrypted, self.user1_id, tok=self.user1_tok) + invite = self.helper.create_room_as(self.user2_id, tok=self.user2_tok) + self.helper.invite( + invite, src=self.user2_id, targ=self.user1_id, tok=self.user2_tok + ) + self._create_remote_invite_room_for_user(self.user1_id, None) + self._tag(joined[0], "m.favourite") + self._tag(joined[1], "u.work") + + def list_with(filters: JsonDict) -> JsonDict: + return { + "ranges": [[0, 99]], + "required_state": [], + "timeline_limit": 1, + "filters": filters, + } + + sync_body = { + "lists": { + "dms": list_with({"is_dm": True}), + "not_dms": list_with({"is_dm": False}), + "invites": list_with({"is_invite": True}), + "encrypted": list_with({"is_encrypted": True}), + "unencrypted": list_with({"is_encrypted": False}), + "spaces": list_with({"room_types": [RoomTypes.SPACE]}), + "rooms": list_with({"room_types": [None]}), + "not_spaces": list_with({"not_room_types": [RoomTypes.SPACE]}), + "favourites": list_with({"tags": ["m.favourite"]}), + "not_work": list_with({"not_tags": ["u.work"]}), + } + } + response = self._sync_both(sync_body, unordered_lists=True) + self.assertEqual(response["lists"]["dms"]["count"], 1) + self.assertEqual(response["lists"]["spaces"]["count"], 1) + + self._send_message(dm) + self._send_message(space) + self._sync_both(sync_body, unordered_lists=True) + + def test_ignored_invites(self) -> None: + """Invites from ignored users, and from servers the MSC4155 rules + ignore, which only apply when the user ignores someone.""" + user3_id = self.register_user("user3", "pass") + user3_tok = self.login(user3_id, "pass") + from_user2 = self.helper.create_room_as(self.user2_id, tok=self.user2_tok) + self.helper.invite( + from_user2, src=self.user2_id, targ=self.user1_id, tok=self.user2_tok + ) + from_user3 = self.helper.create_room_as(user3_id, tok=user3_tok) + self.helper.invite(from_user3, src=user3_id, targ=self.user1_id, tok=user3_tok) + self._put_account_data( + AccountDataTypes.MSC4155_INVITE_PERMISSION_CONFIG, + {"ignored_users": [user3_id]}, + ) + sync_body = { + "lists": { + "all": {"ranges": [[0, 99]], "required_state": [], "timeline_limit": 1} + } + } + # The rules alone ignore nothing. + response = self._sync_both(sync_body) + self.assertEqual(response["lists"]["all"]["count"], 2) + + self._put_account_data( + AccountDataTypes.IGNORED_USER_LIST, {"ignored_users": {"@nobody:test": {}}} + ) + response = self._sync_both(sync_body) + self.assertEqual(response["lists"]["all"]["count"], 1) + + self._put_account_data( + AccountDataTypes.IGNORED_USER_LIST, {"ignored_users": {self.user2_id: {}}} + ) + response = self._sync_both(sync_body) + self.assertEqual(response["lists"]["all"]["count"], 0) + + def test_membership_changes(self) -> None: + """Rooms left, kicked from and banned from in the token range, joined + in it, and a room subscription to a room the user left.""" + rooms = self._joined_rooms(4) + sync_body: JsonDict = { + "lists": { + "all": {"ranges": [[0, 99]], "required_state": [], "timeline_limit": 3} + } + } + self._sync_both(sync_body) + + self.helper.leave(rooms[0], self.user1_id, tok=self.user1_tok) + self.helper.change_membership( + rooms[1], self.user2_id, self.user1_id, "leave", tok=self.user2_tok + ) + self.helper.ban(rooms[2], self.user2_id, self.user1_id, tok=self.user2_tok) + (new_room,) = self._joined_rooms(1) + self._send_message(new_room) + self._sync_both(sync_body) + + # The rooms left before the from token drop out, but for the kick + # and the ban. + self._send_message(rooms[3]) + self._sync_both(sync_body) + + # A subscription to the room left, which the lists leave out. + self._sync_both( + { + **sync_body, + "room_subscriptions": { + rooms[0]: {"required_state": [["*", "*"]], "timeline_limit": 1}, + new_room: {"required_state": [], "timeline_limit": 5}, + }, + } + ) + + def test_extensions_state(self) -> None: + """The receipts and account data extensions read and update the Rust + connection state through its views: rooms sent before, rooms with + updates out of range, and the range widening onto them.""" + rooms = self._joined_rooms(4) + for room_id in rooms: + event_id = self._send_message(room_id) + self.helper.send_read_receipt(room_id, event_id, tok=self.user2_tok) + self.get_success( + self.hs.get_datastores().main.add_account_data_to_room( + self.user1_id, room_id, "org.example.data", {"room": room_id} + ) + ) + extensions = { + "receipts": {"enabled": True}, + "account_data": {"enabled": True}, + } + sync_body: JsonDict = { + "lists": { + "top": {"ranges": [[0, 1]], "required_state": [], "timeline_limit": 1} + }, + "extensions": extensions, + } + self._sync_both(sync_body) + + # Receipts and account data in every room, two of them out of range. + for room_id in rooms: + event_id = self._send_message(room_id) + self.helper.send_read_receipt(room_id, event_id, tok=self.user2_tok) + self._sync_both( + { + **sync_body, + "lists": {"top": {**sync_body["lists"]["top"], "ranges": [[0, 0]]}}, + } + ) + for room_id in rooms: + self.get_success( + self.hs.get_datastores().main.add_account_data_to_room( + self.user1_id, room_id, "org.example.data", {"again": room_id} + ) + ) + self._sync_both( + { + **sync_body, + "lists": {"top": {**sync_body["lists"]["top"], "ranges": [[0, 3]]}}, + } + ) + self._sync_both( + { + **sync_body, + "lists": {"top": {**sync_body["lists"]["top"], "ranges": [[0, 3]]}}, + } + ) + + +class MultiWriterTokenTestCase(TestCase): + """A receipts token's database form, which Rust writes and parses without + the store.""" + + def test_round_trip(self) -> None: + instances = [(1, "w1"), (2, "w2")] + for token, string in [ + (MultiWriterStreamToken(stream=5), "5"), + ( + MultiWriterStreamToken( + stream=5, + instance_map=immutabledict({"w1": 7, "w2": 9}), + ), + "m5~1.7~2.9", + ), + ]: + rust_string, parsed = multi_writer_token_round_trip_for_test( + token, instances + ) + self.assertEqual(rust_string, string) + self.assertEqual(parsed, token) diff --git a/tests/rest/client/sliding_sync/test_rust_rooms.py b/tests/rest/client/sliding_sync/test_rust_rooms.py index 7e08233543..f015d24b34 100644 --- a/tests/rest/client/sliding_sync/test_rust_rooms.py +++ b/tests/rest/client/sliding_sync/test_rust_rooms.py @@ -20,11 +20,12 @@ and one by Rust. The `rooms` of the two responses must be equal, and so must the changes each side makes to its connection's state. """ -from typing import Any +from typing import AbstractSet, Any, Mapping from unittest import SkipTest from parameterized import parameterized_class +from twisted.internet.defer import Deferred from twisted.internet.testing import MemoryReactor import synapse.rest.admin @@ -34,6 +35,13 @@ from synapse.handlers.sliding_sync import rust_sliding_sync_rooms_counter from synapse.metrics import SERVER_NAME_LABEL from synapse.rest.client import knock, login, room, sync from synapse.server import HomeServer +from synapse.synapse_rust.sliding_sync import ( + ConnectionState as RustConnectionState, + InterestedRooms as RustInterestedRooms, + SlidingSyncHandler as RustSlidingSyncHandler, + SlidingSyncRequest as RustSlidingSyncRequest, + SlidingSyncStats as RustSlidingSyncStats, +) from synapse.types import JsonDict, SlidingSyncStreamToken from synapse.types.handlers.sliding_sync import ( MutablePerConnectionState, @@ -70,37 +78,102 @@ def _normalized(rooms: JsonDict) -> JsonDict: return rooms -def _connection_changes(state: MutablePerConnectionState) -> JsonDict: +def _unordered_lists(response: JsonDict) -> JsonDict: + """The response's lists with each operation's rooms sorted. The lists of + these tests cover every room, so neither side sorts them, and Python adds + rooms left in the token range back in set order. `test_rust_lists` + compares the order of sorted lists.""" + return { + name: { + **list_result, + "ops": [ + {**op, "room_ids": sorted(op["room_ids"])} + for op in list_result.get("ops", []) + ], + } + for name, list_result in response.get("lists", {}).items() + } + + +# What Rust's `ConnectionState.updates_for_test` returns. +RustUpdates = tuple[ + Mapping[str, tuple[int, Mapping[str, AbstractSet[str]]]], + Mapping[str, tuple[Mapping[str, int | None], AbstractSet[str]]], +] + + +def _connection_changes(state: MutablePerConnectionState | RustUpdates) -> JsonDict: """What a request changed in its connection's state that the room - computation writes: the room configs and the lazy-loaded members. + computation writes: the room configs and the lazy-loaded members. `state` + is Python's state, or what Rust's `updates_for_test` returns. A lazy member's timestamp is only whether there is one: it is when its own connection last sent the member, which the two connections did a moment apart.""" + configs: Mapping[str, tuple[int, Mapping[str, AbstractSet[str]]]] + lazy: Mapping[str, tuple[Mapping[str, int | None], AbstractSet[str]]] + if isinstance(state, MutablePerConnectionState): + configs = { + room_id: (config.timeline_limit, config.required_state_map) + for room_id, config in state.get_room_config_updates().items() + } + lazy = { + room_id: ( + changes.returned_user_id_to_last_seen_ts_map, + changes.invalidated_user_ids, + ) + for room_id, changes in state.room_lazy_membership.items() + } + else: + configs, lazy = state return { "room_configs": { room_id: ( - config.timeline_limit, + timeline_limit, { state_type: set(state_keys) - for state_type, state_keys in config.required_state_map.items() + for state_type, state_keys in required_state_map.items() }, ) - for room_id, config in state.get_room_config_updates().items() + for room_id, (timeline_limit, required_state_map) in configs.items() }, "lazy_membership": { room_id: ( { user_id: last_seen_ts is not None - for user_id, last_seen_ts in changes.returned_user_id_to_last_seen_ts_map.items() + for user_id, last_seen_ts in returned.items() }, - set(changes.invalidated_user_ids), + set(invalidated), ) - for room_id, changes in state.room_lazy_membership.items() + for room_id, (returned, invalidated) in lazy.items() }, } +class _RecordingRustCalls(_BlockingRustCalls): + """Blocking Rust calls that record the connection state's updates of each + request before it is persisted, by `conn_id`.""" + + def __init__( + self, + handler: RustSlidingSyncHandler, + recorded: dict[str, MutablePerConnectionState | RustUpdates], + ) -> None: + super().__init__(handler) + self._recorded = recorded + + def finish_request( + self, + request: RustSlidingSyncRequest, + state: RustConnectionState, + interested: RustInterestedRooms, + ) -> "Deferred[tuple[int, RustSlidingSyncStats]]": + conn_id = request.body_for_test().get("conn_id") + assert isinstance(conn_id, str) + self._recorded[conn_id] = state.updates_for_test() + return super().__getattr__("finish_request")(request, state, interested) + + @parameterized_class( ("use_new_tables",), [ @@ -127,13 +200,13 @@ class SlidingSyncRustRoomsTestCase(SlidingSyncBase): rust = hs.get_rust_handlers().sliding_sync if rust is None: raise SkipTest("needs the native Rust Postgres driver") - self.rust = _BlockingRustCalls(rust) + # The connection state each side records, by `conn_id`. + self.recorded: dict[str, MutablePerConnectionState | RustUpdates] = {} + self.rust = _RecordingRustCalls(rust, self.recorded) self.handler = hs.get_sliding_sync_handler() self.handler._rust_sliding_sync_strict = True self.pos: dict[str, str | None] = {"python": None, "rust": None} - # The connection state each side records, by `conn_id`. - self.recorded: dict[str, MutablePerConnectionState] = {} connection_store = self.handler.connection_store record_new_state = connection_store.record_new_state @@ -177,9 +250,11 @@ class SlidingSyncRustRoomsTestCase(SlidingSyncBase): tok=self.user1_tok, ) computed = self._rust_rooms_computed() - before - if side == "rust": + if side == "rust" and self.use_new_tables: self.assertGreaterEqual(computed, len(response["rooms"])) else: + # Rust leaves a request to Python until the sliding sync + # tables are filled. self.assertEqual(computed, 0) responses[side] = response python_rooms = _normalized(responses["python"]["rooms"]) @@ -188,7 +263,8 @@ class SlidingSyncRustRoomsTestCase(SlidingSyncBase): for room_id, python_room in python_rooms.items(): self.assertEqual(python_room, rust_rooms[room_id], room_id) self.assertEqual( - responses["python"].get("lists"), responses["rust"].get("lists") + _unordered_lists(responses["python"]), + _unordered_lists(responses["rust"]), ) self.assertEqual( _connection_changes(self.recorded["python"]),