diff --git a/Cargo.lock b/Cargo.lock index c238030f4f..c137810acf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2028,6 +2028,7 @@ dependencies = [ "hyper", "hyper-util", "icu_segmenter", + "indexmap", "ipnet", "itertools", "lazy_static", diff --git a/changelog.d/rust-sss-bench.misc b/changelog.d/rust-sss-bench.misc new file mode 100644 index 0000000000..2f742afcfd --- /dev/null +++ b/changelog.d/rust-sss-bench.misc @@ -0,0 +1 @@ +Add a benchmark of sliding sync in Rust against Python. diff --git a/changelog.d/rust-sss-encode-bytes.misc b/changelog.d/rust-sss-encode-bytes.misc new file mode 100644 index 0000000000..6869c08371 --- /dev/null +++ b/changelog.d/rust-sss-encode-bytes.misc @@ -0,0 +1 @@ +Encode the sliding sync rooms section in Rust as JSON bytes, without the GIL. diff --git a/changelog.d/rust-sss-lists.misc b/changelog.d/rust-sss-lists.misc new file mode 100644 index 0000000000..3317cef9b1 --- /dev/null +++ b/changelog.d/rust-sss-lists.misc @@ -0,0 +1 @@ +Compute the room lists of simplified sliding sync responses, and keep the per-connection state, in Rust when the experimental `rust_sliding_sync` option is on. diff --git a/changelog.d/rust-sss-prereqs.misc b/changelog.d/rust-sss-prereqs.misc new file mode 100644 index 0000000000..6e61edd5fb --- /dev/null +++ b/changelog.d/rust-sss-prereqs.misc @@ -0,0 +1 @@ +Add the prerequisites for computing simplified sliding sync responses in Rust, behind `experimental_features.rust_sliding_sync`. diff --git a/changelog.d/rust-sss-pure.misc b/changelog.d/rust-sss-pure.misc new file mode 100644 index 0000000000..634a108f91 --- /dev/null +++ b/changelog.d/rust-sss-pure.misc @@ -0,0 +1 @@ +Port the pure logic of simplified sliding sync to Rust, held equal to the Python by differential tests. diff --git a/changelog.d/rust-sss-queries.misc b/changelog.d/rust-sss-queries.misc new file mode 100644 index 0000000000..e1d0205bde --- /dev/null +++ b/changelog.d/rust-sss-queries.misc @@ -0,0 +1 @@ +Cut the database queries of sliding sync in Rust. diff --git a/changelog.d/rust-sss-rooms.misc b/changelog.d/rust-sss-rooms.misc new file mode 100644 index 0000000000..663947aefd --- /dev/null +++ b/changelog.d/rust-sss-rooms.misc @@ -0,0 +1 @@ +Compute and encode the rooms of simplified sliding sync responses in Rust when the experimental `rust_sliding_sync` option is on. diff --git a/changelog.d/rust-sss-store.misc b/changelog.d/rust-sss-store.misc new file mode 100644 index 0000000000..9c94d2e863 --- /dev/null +++ b/changelog.d/rust-sss-store.misc @@ -0,0 +1 @@ +Port the database reads and connection-state transactions of simplified sliding sync to Rust, held equal to the Python by differential tests. diff --git a/docker/complement/conf/start_for_complement.sh b/docker/complement/conf/start_for_complement.sh index 247b05149f..396fd8ca55 100755 --- a/docker/complement/conf/start_for_complement.sh +++ b/docker/complement/conf/start_for_complement.sh @@ -6,7 +6,7 @@ set -e echo "Complement Synapse launcher" echo " Args: $*" -echo " Env: SYNAPSE_COMPLEMENT_DATABASE=$SYNAPSE_COMPLEMENT_DATABASE SYNAPSE_COMPLEMENT_USE_WORKERS=$SYNAPSE_COMPLEMENT_USE_WORKERS SYNAPSE_COMPLEMENT_USE_ASYNCIO_REACTOR=$SYNAPSE_COMPLEMENT_USE_ASYNCIO_REACTOR SYNAPSE_USE_RUST_DRIVER=$SYNAPSE_USE_RUST_DRIVER SYNAPSE_RUST_STATE_RES=$SYNAPSE_RUST_STATE_RES SYNAPSE_RUST_SYNC=$SYNAPSE_RUST_SYNC SYNAPSE_RUST_PERSIST=$SYNAPSE_RUST_PERSIST" +echo " Env: SYNAPSE_COMPLEMENT_DATABASE=$SYNAPSE_COMPLEMENT_DATABASE SYNAPSE_COMPLEMENT_USE_WORKERS=$SYNAPSE_COMPLEMENT_USE_WORKERS SYNAPSE_COMPLEMENT_USE_ASYNCIO_REACTOR=$SYNAPSE_COMPLEMENT_USE_ASYNCIO_REACTOR SYNAPSE_USE_RUST_DRIVER=$SYNAPSE_USE_RUST_DRIVER SYNAPSE_RUST_STATE_RES=$SYNAPSE_RUST_STATE_RES SYNAPSE_RUST_SYNC=$SYNAPSE_RUST_SYNC SYNAPSE_RUST_PERSIST=$SYNAPSE_RUST_PERSIST SYNAPSE_RUST_SLIDING_SYNC=$SYNAPSE_RUST_SLIDING_SYNC" function log { d=$(printf '%(%Y-%m-%d %H:%M:%S)T,%.3s\n' ${EPOCHREALTIME/./ }) @@ -146,6 +146,11 @@ if [[ -n "$SYNAPSE_RUST_PERSIST" ]]; then sed -i "s/^experimental_features:$/experimental_features:\n rust_event_persistence: \"$SYNAPSE_RUST_PERSIST\"/" /conf/shared.yaml.j2 fi +# Compute sliding sync responses in Rust (`experimental_features.rust_sliding_sync`). +if [[ -n "$SYNAPSE_RUST_SLIDING_SYNC" ]]; then + sed -i "s/^experimental_features:$/experimental_features:\n rust_sliding_sync: \"$SYNAPSE_RUST_SLIDING_SYNC\"/" /conf/shared.yaml.j2 +fi + # Add a directory for tests to add config overrides if they want mkdir --parents /conf/homeserver.d export _SYNAPSE_COMPLEMENT_EXTRA_CONFIG_DIR=/conf/homeserver.d diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 98a4c35f8b..463cac4b20 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -84,6 +84,8 @@ once_cell = "1.18.0" # Fast hashing for the integer-keyed maps in state resolution. rustc-hash = "2.1.1" itertools = "0.15.0" +# Maps that keep insertion order, as the Python dicts they port do. +indexmap = "2.10.0" # Used by the `libpq` module to load the system libpq at runtime, so that # libpq is an optional runtime dependency rather than a build-time one. libloading = "0.8" diff --git a/rust/src/caches/python.rs b/rust/src/caches/python.rs index e132a33f85..132ae73f31 100644 --- a/rust/src/caches/python.rs +++ b/rust/src/caches/python.rs @@ -35,7 +35,7 @@ use std::{hash::Hash, sync::Arc, time::SystemTime}; use pyo3::{ pyclass, pymethods, - types::{PyAnyMethods, PyModule, PyModuleMethods, PyTuple}, + types::{PyAnyMethods, PyModule, PyModuleMethods, PyTuple, PyTupleMethods}, Bound, PyResult, Python, }; @@ -94,6 +94,13 @@ impl PyCacheKey for Arc { } } +/// The key of a cache of a function without arguments: the empty tuple. +impl PyCacheKey for () { + fn from_py_key(key: &Bound<'_, PyTuple>) -> Option { + key.is_empty().then_some(()) + } +} + /// A key that is one integer, such as a state group ID. Over the invalidation /// stream it arrives as a string. impl PyCacheKey for i64 { diff --git a/rust/src/events/internal_metadata.rs b/rust/src/events/internal_metadata.rs index 8c4209bd50..3e4d31a598 100644 --- a/rust/src/events/internal_metadata.rs +++ b/rust/src/events/internal_metadata.rs @@ -623,6 +623,13 @@ impl EventInternalMetadata { .ok() .and_then(|inner| inner.stream_ordering) } + + /// The writer that persisted the event, if known. + pub fn instance_name(&self) -> Option { + self.read_inner() + .ok() + .and_then(|inner| inner.instance_name.clone()) + } } /// Helper to convert `None` to an `AttributeError` for a property getter. diff --git a/rust/src/handlers/mod.rs b/rust/src/handlers/mod.rs index 659bd1732c..2db712518b 100644 --- a/rust/src/handlers/mod.rs +++ b/rust/src/handlers/mod.rs @@ -24,6 +24,7 @@ use pyo3::{ use crate::homeserver::HomeServer; use crate::persist::python::Persister; +use crate::sliding_sync::SlidingSyncHandler; use crate::state::StateResolutionHandler; use crate::sync::SyncHandler; @@ -39,6 +40,9 @@ struct RustHandlers { /// Only when `rust_event_persistence` is not `off`, the main database /// uses the native Rust driver and this instance writes events. persister: Option>, + /// Only when the main database uses the native Rust driver. Shares + /// `sync`'s store. + sliding_sync: Option>, } #[pymethods] @@ -78,45 +82,53 @@ impl RustHandlers { let rust_connection_pool_class = py .import("synapse.storage.rust_pool")? .getattr("RustConnectionPool")?; - let (state_resolution, sync, persister) = - if connection_pool.is_instance(&rust_connection_pool_class)? { - let connection_pool = connection_pool.unbind(); - let state_resolution = StateResolutionHandler::new( + let (state_resolution, sync, sliding_sync, persister) = if connection_pool + .is_instance(&rust_connection_pool_class)? + { + let connection_pool = connection_pool.unbind(); + let state_resolution = StateResolutionHandler::new( + connection_pool.clone_ref(py), + runtime.clone(), + Arc::clone(&store), + ); + let persister = if config.experimental.rust_event_persistence != "off" { + Persister::for_homeserver( + py, + &homeserver, connection_pool.clone_ref(py), runtime.clone(), Arc::clone(&store), - ); - let persister = if config.experimental.rust_event_persistence != "off" { - Persister::for_homeserver( - py, - &homeserver, - connection_pool.clone_ref(py), - runtime.clone(), - Arc::clone(&store), - state_resolution.conflict_cache(), - )? - .map(|persister| Py::new(py, persister)) - .transpose()? - } else { - None - }; - ( - Some(Py::new(py, state_resolution)?), - Some(Py::new( - py, - SyncHandler::new(connection_pool, runtime.clone(), store, &config), - )?), - persister, - ) + state_resolution.conflict_cache(), + )? + .map(|persister| Py::new(py, persister)) + .transpose()? } else { - (None, None, None) + None }; + let sync = SyncHandler::new( + connection_pool.clone_ref(py), + runtime.clone(), + store, + &config, + ); + let sliding_sync = + SlidingSyncHandler::new(connection_pool, runtime.clone(), Arc::clone(sync.store())); + ( + Some(Py::new(py, state_resolution)?), + Some(Py::new(py, sync)?), + Some(Py::new(py, sliding_sync)?), + persister, + ) + } else { + (None, None, None, None) + }; Ok(RustHandlers { versions, state_resolution, sync, persister, + sliding_sync, }) } @@ -146,6 +158,13 @@ impl RustHandlers { fn persister(&self, py: Python<'_>) -> Option> { self.persister.as_ref().map(|h| h.clone_ref(py)) } + + /// The Rust sliding sync handler, or `None` if the main database does + /// not use the native Rust driver. + #[getter] + fn sliding_sync(&self, py: Python<'_>) -> Option> { + self.sliding_sync.as_ref().map(|h| h.clone_ref(py)) + } } /// Called when registering modules with python. diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 100526bd74..58ef62932e 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -37,6 +37,7 @@ pub mod rendezvous; pub mod room_versions; pub mod runtime; pub mod segmenter; +pub mod sliding_sync; pub mod state; pub mod storage; pub mod sync; @@ -113,6 +114,7 @@ fn synapse_rust(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { msc4388_rendezvous::register_module(py, m)?; persist::register_module(py, m)?; segmenter::register_module(py, m)?; + sliding_sync::register_module(py, m)?; state::register_module(py, m)?; sync::register_module(py, m)?; room_versions::register_module(py, m)?; diff --git a/rust/src/room_versions.rs b/rust/src/room_versions.rs index 1f8570a03e..b279176feb 100644 --- a/rust/src/room_versions.rs +++ b/rust/src/room_versions.rs @@ -550,6 +550,16 @@ static KNOWN_ROOM_VERSIONS: LazyLock = LazyLock::new(| } }); +/// Whether `identifier` is in `KNOWN_ROOM_VERSIONS`, which experimental +/// versions can join at runtime. +pub fn is_known_room_version(identifier: &str) -> bool { + KNOWN_ROOM_VERSIONS + .versions + .read() + .map(|versions| versions.iter().any(|v| v.identifier == identifier)) + .unwrap_or(false) +} + /// Container class for room version constants. /// /// This should contain all room versions that we know about. diff --git a/rust/src/sliding_sync/config.rs b/rust/src/sliding_sync/config.rs new file mode 100644 index 0000000000..5322625651 --- /dev/null +++ b/rust/src/sliding_sync/config.rs @@ -0,0 +1,338 @@ +/* + * 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: + * . + */ + +//! `RoomSyncConfig`: what to fetch for a room, from the lists and room +//! subscriptions that include it. + +use std::collections::{BTreeMap, BTreeSet}; + +use anyhow::bail; +use pyo3::{prelude::*, types::PyModule}; + +/// `StateValues.WILDCARD`: every state key, or every type. +pub const WILDCARD: &str = "*"; +/// `StateValues.LAZY`: the members the timeline needs. Only a state key of +/// `m.room.member`. +pub const LAZY: &str = "$LAZY"; +/// `StateValues.ME`: the requester's user ID. +pub const ME: &str = "$ME"; + +pub const MEMBER: &str = "m.room.member"; + +/// A type to its state keys, with the `*`, `$LAZY` and `$ME` values still +/// in them. Python's is a dict of sets. Neither order is observable. +pub type RequiredStateMap = BTreeMap>; + +/// `RoomSyncConfig`. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct RoomSyncConfig { + pub timeline_limit: i64, + pub required_state_map: RequiredStateMap, +} + +/// Remove `state_key` from every type's set, and drop the sets it empties: +/// a wildcard type is about to cover it. +fn remove_state_key_everywhere(map: &mut RequiredStateMap, state_key: &str) { + map.retain(|_, state_keys| { + state_keys.remove(state_key); + !state_keys.is_empty() + }); +} + +fn has_key(map: &RequiredStateMap, event_type: &str, state_key: &str) -> bool { + map.get(event_type) + .is_some_and(|state_keys| state_keys.contains(state_key)) +} + +impl RoomSyncConfig { + /// `from_room_config`: the config of a list or a room subscription, + /// with the wildcards collapsed. + pub fn from_room_config(timeline_limit: i64, required_state: &[(String, String)]) -> Self { + let mut map = RequiredStateMap::new(); + for (state_type, state_key) in required_state { + // A wildcard type already covers this state key. + if has_key(&map, WILDCARD, state_key) { + continue; + } + // A wildcard state key already covers this type. + if has_key(&map, state_type, WILDCARD) { + continue; + } + if state_type == WILDCARD && state_key == WILDCARD { + map = RequiredStateMap::from([( + WILDCARD.to_owned(), + BTreeSet::from([WILDCARD.to_owned()]), + )]); + break; + } else if state_type == WILDCARD { + remove_state_key_everywhere(&mut map, state_key); + } + + if state_key == WILDCARD { + map.insert(state_type.clone(), BTreeSet::from([state_key.clone()])); + } else { + map.entry(state_type.clone()) + .or_default() + .insert(state_key.clone()); + } + } + RoomSyncConfig { + timeline_limit, + required_state_map: map, + } + } + + /// `combine_room_sync_config`: the union of the two. + pub fn combine(&self, other: &RoomSyncConfig) -> RoomSyncConfig { + let timeline_limit = self.timeline_limit.max(other.timeline_limit); + let mut map = self.required_state_map.clone(); + + for (state_type, state_key_set) in &other.required_state_map { + // Everything is already wanted. + if has_key(&map, WILDCARD, WILDCARD) { + break; + } + // Every state key of this type is already wanted. + if has_key(&map, state_type, WILDCARD) { + continue; + } + if state_type == WILDCARD && state_key_set.contains(WILDCARD) { + map = RequiredStateMap::from([( + state_type.clone(), + BTreeSet::from([WILDCARD.to_owned()]), + )]); + break; + } + + for state_key in state_key_set { + // A wildcard type already covers this state key. + if has_key(&map, WILDCARD, state_key) { + continue; + } + if state_type == WILDCARD { + remove_state_key_everywhere(&mut map, state_key); + } + if state_key == WILDCARD { + map.insert(state_type.clone(), BTreeSet::from([state_key.clone()])); + break; + } + map.entry(state_type.clone()) + .or_default() + .insert(state_key.clone()); + } + } + + RoomSyncConfig { + timeline_limit, + required_state_map: map, + } + } + + /// `must_await_full_state`: whether the room's partial state cannot + /// satisfy the config, so a partially-stated room must be left out. + /// `server_name` is ours, for `is_mine_id`. + /// + /// Fails where Python raises: an empty state key under the wildcard + /// type, which Python indexes into. Python only gets there if no + /// remote user ID comes before it in the set's order. Rust checks the + /// user IDs first, which is one of Python's orders. + pub fn must_await_full_state(&self, server_name: &str) -> anyhow::Result { + let is_mine_id = |user_id: &str| { + user_id + .split_once(':') + .is_some_and(|(_, host)| host == server_name) + }; + + let wildcard_state_keys = self.required_state_map.get(WILDCARD); + if let Some(state_keys) = wildcard_state_keys { + if state_keys.contains(WILDCARD) { + return Ok(true); + } + for possible_user_id in state_keys { + if !possible_user_id.starts_with('@') { + continue; + } + if !possible_user_id.contains(':') { + continue; + } + if !is_mine_id(possible_user_id) { + return Ok(true); + } + } + if state_keys.contains("") { + bail!("an empty state key under the wildcard type"); + } + } + + let Some(membership_state_keys) = self.required_state_map.get(MEMBER) else { + return Ok(false); + }; + for user_id in membership_state_keys { + match user_id.as_str() { + ME | LAZY => continue, + WILDCARD => return Ok(false), + _ if !is_mine_id(user_id) => return Ok(true), + _ => {} + } + } + Ok(false) + } +} + +/// A `RoomSyncConfig` from `(timeline_limit, {type: state keys})`. +pub fn config_from_py(timeline_limit: i64, map: BTreeMap>) -> RoomSyncConfig { + RoomSyncConfig { + timeline_limit, + required_state_map: map + .into_iter() + .map(|(k, v)| (k, v.into_iter().collect())) + .collect(), + } +} + +/// `RoomSyncConfig.from_room_config` through the Rust implementation, as +/// `(timeline_limit, {type: state keys})`. +#[pyfunction] +fn room_sync_config_from_room_config_for_test( + timeline_limit: i64, + required_state: Vec<(String, String)>, +) -> (i64, BTreeMap>) { + to_test_tuple(&RoomSyncConfig::from_room_config( + timeline_limit, + &required_state, + )) +} + +/// `combine_room_sync_config` through the Rust implementation. +#[pyfunction] +fn combine_room_sync_config_for_test( + a: (i64, BTreeMap>), + b: (i64, BTreeMap>), +) -> (i64, BTreeMap>) { + let a = config_from_py(a.0, a.1); + let b = config_from_py(b.0, b.1); + to_test_tuple(&a.combine(&b)) +} + +/// `must_await_full_state` through the Rust implementation. Raises +/// `ValueError` where Python raises. +#[pyfunction] +fn must_await_full_state_for_test( + config: (i64, BTreeMap>), + server_name: &str, +) -> PyResult { + config_from_py(config.0, config.1) + .must_await_full_state(server_name) + .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string())) +} + +pub fn to_test_tuple(config: &RoomSyncConfig) -> (i64, BTreeMap>) { + ( + config.timeline_limit, + config + .required_state_map + .iter() + .map(|(k, v)| (k.clone(), v.iter().cloned().collect())) + .collect(), + ) +} + +pub fn register_test_functions(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_function(wrap_pyfunction!( + room_sync_config_from_room_config_for_test, + m + )?)?; + m.add_function(wrap_pyfunction!(combine_room_sync_config_for_test, m)?)?; + m.add_function(wrap_pyfunction!(must_await_full_state_for_test, m)?)?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn rs(pairs: &[(&str, &str)]) -> Vec<(String, String)> { + pairs + .iter() + .map(|(t, k)| (t.to_string(), k.to_string())) + .collect() + } + + fn map(entries: &[(&str, &[&str])]) -> RequiredStateMap { + entries + .iter() + .map(|(t, keys)| (t.to_string(), keys.iter().map(|k| k.to_string()).collect())) + .collect() + } + + #[test] + fn wildcards_collapse() { + let config = RoomSyncConfig::from_room_config( + 5, + &rs(&[ + ("m.a", "x"), + ("m.a", "y"), + ("*", "x"), + ("m.b", "*"), + ("m.b", "z"), + ]), + ); + assert_eq!( + config.required_state_map, + map(&[("m.a", &["y"]), ("*", &["x"]), ("m.b", &["*"])]) + ); + + let config = + RoomSyncConfig::from_room_config(5, &rs(&[("m.a", "x"), ("*", "*"), ("m.b", "y")])); + assert_eq!(config.required_state_map, map(&[("*", &["*"])])); + } + + #[test] + fn combine_takes_the_union() { + let a = RoomSyncConfig { + timeline_limit: 1, + required_state_map: map(&[("m.a", &["x"]), ("m.member", &["$LAZY"])]), + }; + let b = RoomSyncConfig { + timeline_limit: 3, + required_state_map: map(&[("*", &["x"]), ("m.c", &["*", "q"])]), + }; + let combined = a.combine(&b); + assert_eq!(combined.timeline_limit, 3); + assert_eq!( + combined.required_state_map, + map(&[("m.member", &["$LAZY"]), ("*", &["x"]), ("m.c", &["*"])]) + ); + } + + #[test] + fn full_state_needed_for_remote_members() { + let config = |entries: &[(&str, &[&str])]| RoomSyncConfig { + timeline_limit: 1, + required_state_map: map(entries), + }; + assert!(!config(&[("m.room.member", &["$LAZY", "$ME", "@a:here"])]) + .must_await_full_state("here") + .unwrap()); + assert!(config(&[("m.room.member", &["@a:there"])]) + .must_await_full_state("here") + .unwrap()); + assert!(config(&[("*", &["@a:there"])]) + .must_await_full_state("here") + .unwrap()); + assert!(config(&[("*", &[""])]) + .must_await_full_state("here") + .is_err()); + } +} diff --git a/rust/src/sliding_sync/connection.rs b/rust/src/sliding_sync/connection.rs new file mode 100644 index 0000000000..8198fcd158 --- /dev/null +++ b/rust/src/sliding_sync/connection.rs @@ -0,0 +1,562 @@ +/* + * 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 per-connection state: what a sliding sync connection has sent down +//! before (`synapse.types.handlers.sliding_sync`). +//! +//! Python's mutable twins wrap the previous state in a `ChainMap`, whose +//! top map is the request's updates. Here the mutable types hold the +//! previous maps behind an `Arc` and the updates beside them, and persist +//! writes only the updates, as Python's does. + +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; + +use pyo3::{exceptions::PyValueError, prelude::*, types::PyModule}; + +use crate::sliding_sync::config::{config_from_py, RoomSyncConfig}; +use crate::types::stream_token::RoomStreamToken; + +/// `LAZY_MEMBERS_UPDATE_INTERVAL`, in milliseconds: how stale a lazy +/// member's last-seen timestamp may get before it is rewritten. +pub const LAZY_MEMBERS_UPDATE_INTERVAL_MS: i64 = 60 * 60 * 1000; + +/// `HaveSentRoom`: whether the room was sent down the connection. The +/// valid transitions are never to live, live to previously, and previously +/// to live. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum HaveSentRoom { + /// Never sent, or forgotten. + Never, + /// Sent, with every update since. + Live, + /// Sent, with the updates after the token still to send. + Previously(T), +} + +impl HaveSentRoom { + pub fn is_never(&self) -> bool { + matches!(self, HaveSentRoom::Never) + } + + pub fn is_live(&self) -> bool { + matches!(self, HaveSentRoom::Live) + } +} + +/// `RoomStatusMap`: the status of each room for one stream. +#[derive(Debug, Clone)] +pub struct RoomStatusMap { + statuses: Arc>>, +} + +impl Default for RoomStatusMap { + fn default() -> Self { + RoomStatusMap { + statuses: Arc::new(HashMap::new()), + } + } +} + +impl RoomStatusMap { + pub fn new(statuses: HashMap>) -> Self { + RoomStatusMap { + statuses: Arc::new(statuses), + } + } + + pub fn have_sent_room(&self, room_id: &str) -> HaveSentRoom { + self.statuses + .get(room_id) + .cloned() + .unwrap_or(HaveSentRoom::Never) + } + + pub fn get_mutable(&self) -> MutableRoomStatusMap { + MutableRoomStatusMap { + base: Arc::clone(&self.statuses), + updates: HashMap::new(), + } + } + + pub fn len(&self) -> usize { + self.statuses.len() + } + + pub fn is_empty(&self) -> bool { + self.statuses.is_empty() + } + + pub fn iter(&self) -> impl Iterator)> { + self.statuses.iter().map(|(k, v)| (k.as_str(), v)) + } +} + +/// `MutableRoomStatusMap`: a [`RoomStatusMap`] and the request's updates +/// to it. +#[derive(Debug, Clone)] +pub struct MutableRoomStatusMap { + base: Arc>>, + updates: HashMap>, +} + +impl MutableRoomStatusMap { + pub fn have_sent_room(&self, room_id: &str) -> HaveSentRoom { + self.updates + .get(room_id) + .or_else(|| self.base.get(room_id)) + .cloned() + .unwrap_or(HaveSentRoom::Never) + } + + /// `get_updates`: only the changes made. + pub fn get_updates(&self) -> &HashMap> { + &self.updates + } + + /// `record_sent_rooms`: these rooms were sent in the response. A live + /// room records nothing. + pub fn record_sent_rooms<'a>(&mut self, room_ids: impl IntoIterator) { + for room_id in room_ids { + if self.have_sent_room(room_id).is_live() { + continue; + } + self.updates.insert(room_id.to_owned(), HaveSentRoom::Live); + } + } + + /// `record_unsent_rooms`: these rooms have updates that were not sent. + /// Only a live room changes, to previously at `from_token`. + pub fn record_unsent_rooms<'a>( + &mut self, + room_ids: impl IntoIterator, + from_token: &T, + ) { + for room_id in room_ids { + if !self.have_sent_room(room_id).is_live() { + continue; + } + self.updates.insert( + room_id.to_owned(), + HaveSentRoom::Previously(from_token.clone()), + ); + } + } + + /// `copy`: the flattened map. + pub fn copy(&self) -> RoomStatusMap { + let mut statuses = (*self.base).clone(); + statuses.extend(self.updates.iter().map(|(k, v)| (k.clone(), v.clone()))); + RoomStatusMap::new(statuses) + } + + pub fn len(&self) -> usize { + self.base.len() + + self + .updates + .keys() + .filter(|room_id| !self.base.contains_key(*room_id)) + .count() + } + + pub fn is_empty(&self) -> bool { + self.len() == 0 + } +} + +/// `PerConnectionState`: a snapshot of what the connection has sent down. +/// `R` is the receipts stream's token, a `MultiWriterStreamToken` in +/// Python. +#[derive(Debug, Clone)] +pub struct PerConnectionState { + /// When the connection was last used, in milliseconds. Only accurate + /// to `UPDATE_INTERVAL_LAST_USED_TS`. + pub last_used_ts: Option, + pub rooms: RoomStatusMap, + pub receipts: RoomStatusMap, + pub account_data: RoomStatusMap, + /// The config of each room sent down before. + pub room_configs: Arc>, +} + +impl Default for PerConnectionState { + fn default() -> Self { + PerConnectionState { + last_used_ts: None, + rooms: RoomStatusMap::default(), + receipts: RoomStatusMap::default(), + account_data: RoomStatusMap::default(), + room_configs: Arc::new(HashMap::new()), + } + } +} + +impl PerConnectionState { + pub fn get_mutable(&self) -> MutablePerConnectionState { + MutablePerConnectionState { + last_used_ts: self.last_used_ts, + rooms: self.rooms.get_mutable(), + receipts: self.receipts.get_mutable(), + account_data: self.account_data.get_mutable(), + room_configs_base: Arc::clone(&self.room_configs), + room_config_updates: HashMap::new(), + room_lazy_membership: HashMap::new(), + } + } + + pub fn len(&self) -> usize { + self.account_data.len() + self.rooms.len() + self.receipts.len() + self.room_configs.len() + } + + pub fn is_empty(&self) -> bool { + self.len() == 0 + } +} + +/// `RoomLazyMembershipChanges`: the lazy-loaded memberships one room +/// needed in the request. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct RoomLazyMembershipChanges { + /// The users whose membership the room lazily loaded, to when it was + /// last needed before, or `None` if this request sent it first. It + /// includes users not sent because they were sent before. + pub returned_user_id_to_last_seen_ts_map: HashMap>, + /// The users whose latest membership was not sent down. + pub invalidated_user_ids: HashSet, +} + +impl RoomLazyMembershipChanges { + /// `get_returned_user_ids_to_update`: the users whose last-seen + /// timestamp the database needs: sent first now, or not rewritten for + /// [`LAZY_MEMBERS_UPDATE_INTERVAL_MS`]. + pub fn get_returned_user_ids_to_update(&self, now_ms: i64) -> Vec<&str> { + self.returned_user_id_to_last_seen_ts_map + .iter() + .filter(|(_, last_seen_ts)| { + last_seen_ts.is_none_or(|ts| now_ms - ts >= LAZY_MEMBERS_UPDATE_INTERVAL_MS) + }) + .map(|(user_id, _)| user_id.as_str()) + .collect() + } + + /// `has_updates`: whether the room's lazy members need persisting. + pub fn has_updates(&self, now_ms: i64) -> bool { + !self.invalidated_user_ids.is_empty() + || !self.get_returned_user_ids_to_update(now_ms).is_empty() + } +} + +/// `MutablePerConnectionState`: the state and the request's updates. +#[derive(Debug, Clone)] +pub struct MutablePerConnectionState { + pub last_used_ts: Option, + pub rooms: MutableRoomStatusMap, + pub receipts: MutableRoomStatusMap, + pub account_data: MutableRoomStatusMap, + room_configs_base: Arc>, + room_config_updates: HashMap, + /// The lazy-loaded memberships each room needed in the request. + pub room_lazy_membership: HashMap, +} + +impl MutablePerConnectionState { + /// The room's config, the request's update first. + pub fn room_config(&self, room_id: &str) -> Option<&RoomSyncConfig> { + self.room_config_updates + .get(room_id) + .or_else(|| self.room_configs_base.get(room_id)) + } + + pub fn set_room_config(&mut self, room_id: &str, config: RoomSyncConfig) { + self.room_config_updates.insert(room_id.to_owned(), config); + } + + /// `get_room_config_updates`. + pub fn get_room_config_updates(&self) -> &HashMap { + &self.room_config_updates + } + + /// `has_updates`: whether persisting would change anything. A spurious + /// persist makes a new connection position, which invalidates caches. + pub fn has_updates(&self, now_ms: i64) -> bool { + !self.rooms.get_updates().is_empty() + || !self.receipts.get_updates().is_empty() + || !self.account_data.get_updates().is_empty() + || !self.room_config_updates.is_empty() + || self + .room_lazy_membership + .values() + .any(|change| change.has_updates(now_ms)) + } + + /// `copy`: the flattened state. + pub fn copy(&self) -> PerConnectionState { + let mut room_configs = (*self.room_configs_base).clone(); + room_configs.extend( + self.room_config_updates + .iter() + .map(|(k, v)| (k.clone(), v.clone())), + ); + PerConnectionState { + last_used_ts: self.last_used_ts, + rooms: self.rooms.copy(), + receipts: self.receipts.copy(), + account_data: self.account_data.copy(), + room_configs: Arc::new(room_configs), + } + } + + /// `len()` of Python's mutable state: the `ChainMap`s count each room + /// once. + pub fn len(&self) -> usize { + let room_configs = self.room_configs_base.len() + + self + .room_config_updates + .keys() + .filter(|room_id| !self.room_configs_base.contains_key(*room_id)) + .count(); + self.account_data.len() + self.rooms.len() + self.receipts.len() + room_configs + } + + pub fn is_empty(&self) -> bool { + self.len() == 0 + } +} + +/// A status from Python's `(flag, token)` test form. +fn status_from_py(flag: &str, token: Option) -> PyResult> { + match (flag, token) { + ("never", None) => Ok(HaveSentRoom::Never), + ("live", None) => Ok(HaveSentRoom::Live), + ("previously", Some(token)) => Ok(HaveSentRoom::Previously(token)), + _ => Err(PyValueError::new_err(format!( + "bad status {flag} {token:?}" + ))), + } +} + +fn status_to_py(status: &HaveSentRoom) -> (&'static str, Option) { + match status { + HaveSentRoom::Never => ("never", None), + HaveSentRoom::Live => ("live", None), + HaveSentRoom::Previously(token) => ("previously", Some(*token)), + } +} + +type PyStatuses = HashMap)>; + +fn statuses_from_py(statuses: PyStatuses) -> PyResult>> { + statuses + .into_iter() + .map(|(room_id, (flag, token))| Ok((room_id, status_from_py(&flag, token)?))) + .collect() +} + +fn statuses_to_py<'a>( + statuses: impl Iterator)>, +) -> HashMap)> { + statuses + .map(|(room_id, status)| (room_id.clone(), status_to_py(status))) + .collect() +} + +/// A `MutableRoomStatusMap[int]` through the Rust implementation: from +/// `statuses`, apply `ops` of `("sent", room_ids, None)` or `("unsent", +/// room_ids, token)`. Returns `(updates, flattened, len)`, each status as +/// `(flag, token)`. +#[pyfunction] +#[allow(clippy::type_complexity)] +fn room_status_map_ops_for_test( + statuses: PyStatuses, + ops: Vec<(String, Vec, Option)>, +) -> PyResult<( + HashMap)>, + HashMap)>, + usize, +)> { + let map = RoomStatusMap::new(statuses_from_py(statuses)?); + let mut mutable = map.get_mutable(); + for (op, room_ids, token) in &ops { + let room_ids = room_ids.iter().map(String::as_str); + match (op.as_str(), token) { + ("sent", None) => mutable.record_sent_rooms(room_ids), + ("unsent", Some(token)) => mutable.record_unsent_rooms(room_ids, token), + _ => return Err(PyValueError::new_err(format!("bad op {op}"))), + } + } + let flattened = mutable.copy(); + Ok(( + statuses_to_py(mutable.get_updates().iter()), + statuses_to_py(flattened.statuses.iter()), + mutable.len(), + )) +} + +/// `RoomLazyMembershipChanges` through the Rust implementation: returns +/// `(get_returned_user_ids_to_update, has_updates)` at `now_ms`. +#[pyfunction] +fn lazy_membership_changes_for_test( + returned_user_id_to_last_seen_ts_map: HashMap>, + invalidated_user_ids: HashSet, + now_ms: i64, +) -> (Vec, bool) { + let changes = RoomLazyMembershipChanges { + returned_user_id_to_last_seen_ts_map, + invalidated_user_ids, + }; + let mut to_update: Vec = changes + .get_returned_user_ids_to_update(now_ms) + .into_iter() + .map(str::to_owned) + .collect(); + to_update.sort(); + (to_update, changes.has_updates(now_ms)) +} + +/// `MutablePerConnectionState.has_updates` and `len` through the Rust +/// implementation, every stream's token an int (rooms' as a live +/// `RoomStreamToken` at that position). `room_configs` and +/// `room_config_updates` map a room to `(timeline_limit, {type: state +/// keys})`, and `lazy` a room to `(returned map, invalidated)`. +#[pyfunction] +#[allow(clippy::too_many_arguments, clippy::type_complexity)] +fn per_connection_state_for_test( + rooms: PyStatuses, + rooms_sent: Vec, + receipts: PyStatuses, + receipts_sent: Vec, + account_data: PyStatuses, + account_data_sent: Vec, + room_configs: HashMap>)>, + room_config_updates: HashMap>)>, + lazy: HashMap>, HashSet)>, + now_ms: i64, +) -> PyResult<(bool, usize, usize)> { + let to_room_tokens = |statuses: HashMap>| { + statuses + .into_iter() + .map(|(room_id, status)| { + let status = match status { + HaveSentRoom::Never => HaveSentRoom::Never, + HaveSentRoom::Live => HaveSentRoom::Live, + HaveSentRoom::Previously(token) => { + HaveSentRoom::Previously(RoomStreamToken::live(token)) + } + }; + (room_id, status) + }) + .collect() + }; + let state = PerConnectionState { + last_used_ts: None, + rooms: RoomStatusMap::new(to_room_tokens(statuses_from_py(rooms)?)), + receipts: RoomStatusMap::new(statuses_from_py(receipts)?), + account_data: RoomStatusMap::new(statuses_from_py(account_data)?), + room_configs: Arc::new( + room_configs + .into_iter() + .map(|(room_id, (limit, map))| (room_id, config_from_py(limit, map))) + .collect(), + ), + }; + let base_len = state.len(); + let mut mutable = state.get_mutable(); + mutable + .rooms + .record_sent_rooms(rooms_sent.iter().map(String::as_str)); + mutable + .receipts + .record_sent_rooms(receipts_sent.iter().map(String::as_str)); + mutable + .account_data + .record_sent_rooms(account_data_sent.iter().map(String::as_str)); + for (room_id, (limit, map)) in room_config_updates { + mutable.set_room_config(&room_id, config_from_py(limit, map)); + } + for (room_id, (returned, invalidated)) in lazy { + mutable.room_lazy_membership.insert( + room_id, + RoomLazyMembershipChanges { + returned_user_id_to_last_seen_ts_map: returned, + invalidated_user_ids: invalidated, + }, + ); + } + Ok((mutable.has_updates(now_ms), base_len, mutable.len())) +} + +pub fn register_test_functions(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_function(wrap_pyfunction!(room_status_map_ops_for_test, m)?)?; + m.add_function(wrap_pyfunction!(lazy_membership_changes_for_test, m)?)?; + m.add_function(wrap_pyfunction!(per_connection_state_for_test, m)?)?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn transitions() { + let map = RoomStatusMap::new(HashMap::from([ + ("!live".to_owned(), HaveSentRoom::Live), + ("!prev".to_owned(), HaveSentRoom::Previously(3)), + ])); + let mut mutable = map.get_mutable(); + mutable.record_sent_rooms(["!live", "!prev", "!new"]); + // A live room records nothing. + assert_eq!( + mutable.get_updates(), + &HashMap::from([ + ("!prev".to_owned(), HaveSentRoom::Live), + ("!new".to_owned(), HaveSentRoom::Live), + ]) + ); + + let mut mutable = map.get_mutable(); + mutable.record_unsent_rooms(["!live", "!prev", "!new"], &7); + assert_eq!( + mutable.get_updates(), + &HashMap::from([("!live".to_owned(), HaveSentRoom::Previously(7))]) + ); + assert_eq!(mutable.have_sent_room("!prev"), HaveSentRoom::Previously(3)); + assert_eq!(mutable.have_sent_room("!new"), HaveSentRoom::Never); + assert_eq!(mutable.len(), 2); + } + + #[test] + fn lazy_members_rewritten_after_the_interval() { + let changes = RoomLazyMembershipChanges { + returned_user_id_to_last_seen_ts_map: HashMap::from([ + ("@new:x".to_owned(), None), + ("@stale:x".to_owned(), Some(0)), + ("@fresh:x".to_owned(), Some(LAZY_MEMBERS_UPDATE_INTERVAL_MS)), + ]), + invalidated_user_ids: HashSet::new(), + }; + let mut to_update = + changes.get_returned_user_ids_to_update(LAZY_MEMBERS_UPDATE_INTERVAL_MS + 1); + to_update.sort(); + assert_eq!(to_update, vec!["@new:x", "@stale:x"]); + + let changes = RoomLazyMembershipChanges { + returned_user_id_to_last_seen_ts_map: HashMap::from([( + "@fresh:x".to_owned(), + Some(10), + )]), + invalidated_user_ids: HashSet::new(), + }; + assert!(!changes.has_updates(20)); + } +} diff --git a/rust/src/sliding_sync/encode.rs b/rust/src/sliding_sync/encode.rs new file mode 100644 index 0000000000..be88bbba38 --- /dev/null +++ b/rust/src/sliding_sync/encode.rs @@ -0,0 +1,212 @@ +/* + * 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 servlet's `encode_rooms` for the rooms Rust computed, so that their +//! events never become Python objects. + +use std::collections::HashMap; +use std::sync::Arc; + +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use pyo3::types::PyBytes; +use serde_json::{json, Map, Value}; + +use crate::events::json_object::JsonObject; +use crate::events::serialize::{ + serialize_event, BundleRef, SerializeEventConfig, SerializeLookups, +}; +use crate::events::Event; +use crate::sliding_sync::result::RoomResult; +use crate::sliding_sync::room_entry::RoomEntry; +use crate::sync::relations::Bundle; + +/// The lookups of one room: the bundles and redactions Rust loaded for it, +/// and the module callbacks' unsigned additions. +struct EntryLookups<'a> { + entry: &'a RoomEntry, + unsigned_additions: &'a HashMap, +} + +impl SerializeLookups for EntryLookups<'_> { + fn redaction(&self, event_id: &str) -> Option<&Event> { + self.entry.redactions.get(event_id) + } + + fn bundle(&self, event_id: &str) -> Option> { + self.entry + .bundles + .as_ref()? + .get(event_id) + .map(Bundle::as_ref) + } + + fn unsigned_additions(&self, event_id: &str) -> Option<&JsonObject> { + self.unsigned_additions.get(event_id) + } +} + +/// `encode_rooms` for one room. +pub fn encode_room( + entry: &RoomEntry, + time_now_ms: i64, + config: &SerializeEventConfig, + unsigned_additions: &HashMap, +) -> PyResult { + let lookups = EntryLookups { + entry, + unsigned_additions, + }; + + // Notifications are only done properly by clients, so these are + // placeholders. + let mut room = Map::new(); + room.insert("notification_count".to_owned(), json!(0)); + room.insert("highlight_count".to_owned(), json!(0)); + if let Some(bump_stamp) = entry.bump_stamp { + room.insert("bump_stamp".to_owned(), json!(bump_stamp)); + } + if let Some(joined_count) = entry.joined_count { + room.insert("joined_count".to_owned(), json!(joined_count)); + } + if let Some(invited_count) = entry.invited_count { + room.insert("invited_count".to_owned(), json!(invited_count)); + } + if let Some(name) = entry.name.as_deref().filter(|name| !name.is_empty()) { + room.insert("name".to_owned(), json!(name)); + } + if let Some(avatar) = entry.avatar.as_deref().filter(|avatar| !avatar.is_empty()) { + room.insert("avatar".to_owned(), json!(avatar)); + } + if !entry.heroes.is_empty() { + let heroes: Vec = entry + .heroes + .iter() + .map(|hero| { + let mut serialized = Map::new(); + serialized.insert("user_id".to_owned(), json!(hero.user_id)); + // The spec spells it "displayname". + if let Some(display_name) = &hero.display_name { + serialized.insert("displayname".to_owned(), json!(display_name)); + } + if let Some(avatar_url) = &hero.avatar_url { + serialized.insert("avatar_url".to_owned(), json!(avatar_url)); + } + Value::Object(serialized) + }) + .collect(); + room.insert("heroes".to_owned(), Value::Array(heroes)); + } + // Only sent when true, to save bandwidth. + if entry.initial { + room.insert("initial".to_owned(), json!(true)); + } + if entry.unstable_expanded_timeline { + room.insert("unstable_expanded_timeline".to_owned(), json!(true)); + } + if !entry.required_state.is_empty() { + let state = entry + .required_state + .iter() + .map(|event| { + serialize_event(event, time_now_ms, config, None, false, &lookups) + .map(Value::Object) + }) + .collect::>>()?; + room.insert("required_state".to_owned(), Value::Array(state)); + } + if !entry.timeline.is_empty() { + let with_bundles = entry.bundles.is_some(); + let timeline = entry + .timeline + .iter() + .map(|filtered| { + serialize_event( + &filtered.event, + time_now_ms, + config, + Some(&filtered.membership), + with_bundles, + &lookups, + ) + .map(Value::Object) + }) + .collect::>>()?; + room.insert("timeline".to_owned(), Value::Array(timeline)); + } + if let Some(limited) = entry.limited { + room.insert("limited".to_owned(), json!(limited)); + } + if let Some(prev_batch) = &entry.prev_batch { + room.insert("prev_batch".to_owned(), json!(prev_batch)); + } + if let Some(num_live) = entry.num_live { + room.insert("num_live".to_owned(), json!(num_live)); + } + // Absent on rooms that are not DMs. + if entry.is_dm { + room.insert("is_dm".to_owned(), json!(true)); + } + if !entry.stripped_state.is_empty() { + room.insert( + "invite_state".to_owned(), + Value::Array(entry.stripped_state.clone()), + ); + } + Ok(Value::Object(room)) +} + +/// `encode_rooms` for `rooms`, as JSON object members: `"room_id":{...}` +/// joined by commas, with no braces around them, so the servlet can splice +/// them into the rooms object next to rooms Python encoded. +/// `unsigned_additions` are the module callbacks' additions to the rooms' +/// events, by event ID. The encoding runs without the GIL. +#[pyfunction] +#[pyo3(signature = (rooms, time_now_ms, config, *, unsigned_additions = None))] +pub fn encode_rooms<'py>( + py: Python<'py>, + rooms: Vec>, + time_now_ms: i64, + config: &SerializeEventConfig, + unsigned_additions: Option>, +) -> PyResult> { + let unsigned_additions = unsigned_additions.unwrap_or_default(); + let entries: Vec> = rooms.iter().map(|room| room.get().entry.clone()).collect(); + let encoded = + py.detach(|| encode_rooms_json(&entries, time_now_ms, config, &unsigned_additions))?; + Ok(PyBytes::new(py, &encoded)) +} + +/// The members of the rooms object for `entries`. +fn encode_rooms_json( + entries: &[Arc], + time_now_ms: i64, + config: &SerializeEventConfig, + unsigned_additions: &HashMap, +) -> PyResult> { + let mut out = Vec::new(); + for (index, entry) in entries.iter().enumerate() { + if index > 0 { + out.push(b','); + } + let room = encode_room(entry, time_now_ms, config, unsigned_additions)?; + serde_json::to_writer(&mut out, &entry.room_id) + .and_then(|()| { + out.push(b':'); + serde_json::to_writer(&mut out, &room) + }) + .map_err(|err| PyValueError::new_err(format!("encoding room: {err}")))?; + } + Ok(out) +} diff --git a/rust/src/sliding_sync/extensions/mod.rs b/rust/src/sliding_sync/extensions/mod.rs new file mode 100644 index 0000000000..3c483c3e42 --- /dev/null +++ b/rust/src/sliding_sync/extensions/mod.rs @@ -0,0 +1,23 @@ +/* + * 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 sliding sync extensions (`synapse.handlers.sliding_sync.extensions`). + +use pyo3::{prelude::*, types::PyModule}; + +pub mod relevant; + +pub fn register_test_functions(m: &Bound<'_, PyModule>) -> PyResult<()> { + relevant::register_test_functions(m) +} diff --git a/rust/src/sliding_sync/extensions/relevant.rs b/rust/src/sliding_sync/extensions/relevant.rs new file mode 100644 index 0000000000..6504759a65 --- /dev/null +++ b/rust/src/sliding_sync/extensions/relevant.rs @@ -0,0 +1,120 @@ +/* + * 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: + * . + */ + +//! `find_relevant_room_ids_for_extension`: the rooms an extension covers, +//! from its `lists` and `rooms` and what the response holds. + +use std::collections::HashSet; +use std::sync::Arc; + +use anyhow::{bail, Context}; +use indexmap::IndexMap; +use pyo3::{exceptions::PyAssertionError, prelude::*, types::PyModule}; + +use crate::sliding_sync::room_lists::lists::SyncOp; + +/// The response's lists, by name, each with its operations. +pub type ActualLists = IndexMap, Vec>; + +/// The rooms of a list's single `SYNC` operation. Python asserts there is +/// exactly one, so a list with no ranges, or several, fails here. +fn single_op_rooms<'a>(name: &str, ops: &'a [SyncOp]) -> anyhow::Result<&'a [Arc]> { + match ops { + [op] => Ok(&op.room_ids), + _ => bail!("list {name} has {} operations, not one", ops.len()), + } +} + +/// `find_relevant_room_ids_for_extension`: of the response, the rooms +/// `requested_room_ids` names among `actual_room_ids` (the room +/// subscriptions), and the rooms of the lists `requested_lists` names. `*` +/// means all of them. +pub fn find_relevant_room_ids_for_extension( + requested_lists: Option<&[String]>, + requested_room_ids: Option<&[String]>, + actual_lists: &ActualLists, + actual_room_ids: &HashSet>, +) -> anyhow::Result>> { + let mut relevant: HashSet> = HashSet::new(); + + for room_id in requested_room_ids.unwrap_or_default() { + if room_id == "*" { + relevant.extend(actual_room_ids.iter().cloned()); + break; + } + if let Some(room_id) = actual_room_ids.get(room_id.as_str()) { + relevant.insert(Arc::clone(room_id)); + } + } + + for list_key in requested_lists.unwrap_or_default() { + if list_key == "*" { + for (name, ops) in actual_lists { + relevant.extend(single_op_rooms(name, ops)?.iter().cloned()); + } + break; + } + if let Some(ops) = actual_lists.get(list_key.as_str()) { + relevant.extend( + single_op_rooms(list_key, ops) + .context("an extension's list")? + .iter() + .cloned(), + ); + } + } + + Ok(relevant) +} + +/// `find_relevant_room_ids_for_extension` through the Rust implementation, +/// each list as its operations' room IDs. Raises `AssertionError` as +/// Python does. +#[pyfunction] +fn find_relevant_room_ids_for_test( + requested_lists: Option>, + requested_room_ids: Option>, + actual_lists: Vec<(String, Vec>)>, + actual_room_ids: HashSet, +) -> PyResult> { + let actual_lists: ActualLists = actual_lists + .into_iter() + .map(|(name, ops)| { + let ops = ops + .into_iter() + .map(|room_ids| SyncOp { + range: (0, 0), + room_ids: room_ids.into_iter().map(Arc::from).collect(), + }) + .collect(); + (Arc::from(name), ops) + }) + .collect(); + let actual_room_ids: HashSet> = actual_room_ids.into_iter().map(Arc::from).collect(); + let relevant = find_relevant_room_ids_for_extension( + requested_lists.as_deref(), + requested_room_ids.as_deref(), + &actual_lists, + &actual_room_ids, + ) + .map_err(|e| PyAssertionError::new_err(e.to_string()))?; + let mut relevant: Vec = relevant.iter().map(|r| r.to_string()).collect(); + relevant.sort(); + Ok(relevant) +} + +pub fn register_test_functions(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_function(wrap_pyfunction!(find_relevant_room_ids_for_test, m)?)?; + Ok(()) +} diff --git a/rust/src/sliding_sync/mod.rs b/rust/src/sliding_sync/mod.rs new file mode 100644 index 0000000000..f0cec1b6a3 --- /dev/null +++ b/rust/src/sliding_sync/mod.rs @@ -0,0 +1,682 @@ +/* + * 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: + * . + */ + +//! Simplified sliding sync (MSC4186) in Rust, reading the database through +//! the native Postgres driver. +//! +//! [`SlidingSyncHandler`] is the Python entry point, used by +//! `synapse.handlers.sliding_sync.SlidingSyncHandler` when +//! `rust_sliding_sync` is on. It shares the v2 sync handler's +//! [`SyncStore`], so the two read through the same caches. +//! +//! The handler's entry points follow the v2 handler's shape: a `start_*` +//! method reads what it needs from Python and returns a future, which the +//! public method runs off the GIL on the runtime and returns as a deferred +//! ([`SlidingSyncHandler::defer`]), and a `_blocking` twin runs to completion +//! for tests ([`SlidingSyncHandler::block_on`]). + +use std::collections::{HashMap, HashSet}; +use std::sync::{Arc, Mutex}; + +use anyhow::Context; +use futures::future::BoxFuture; +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; +pub mod encode; +pub mod extensions; +pub mod request; +pub mod required_state; +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. +#[pyclass(frozen)] +pub struct SlidingSyncHandler { + /// The Python `RustConnectionPool`. Its `_pool` is read per call, as + /// `close()`/`start()` replace it. + db_pool: Py, + runtime: RustRuntime, + /// The queries, over the Rust store's caches. The v2 sync handler's. + store: Arc, +} + +impl SlidingSyncHandler { + pub fn new(db_pool: Py, runtime: RustRuntime, store: Arc) -> Self { + SlidingSyncHandler { + db_pool, + runtime, + store, + } + } + + pub fn store(&self) -> &Arc { + &self.store + } + + /// The connection pool to read through, or an error if it is closed. + pub fn pool(&self, py: Python<'_>) -> PyResult { + let pool = self.db_pool.bind(py).getattr(intern!(py, "_pool"))?; + if pool.is_none() { + return Err(PyRuntimeError::new_err("the database pool is closed")); + } + Ok(pool.cast::()?.get().pool.clone()) + } + + /// Run `future` on the runtime and return a deferred for its result. + /// Cancelling the deferred stops the computation. The servlet is not + /// `@cancellable`, so today that only happens on shutdown. + pub fn defer<'py, O>( + &self, + py: Python<'py>, + future: BoxFuture<'static, PyResult>, + ) -> PyResult> + where + for<'a> O: IntoPyObject<'a> + Send + 'static, + { + create_deferred_with(py, &self.runtime, future, OnCancel::Abort) + } + + /// Run `future` to completion on the runtime, releasing the GIL. For the + /// `_blocking` entry points. + pub fn block_on( + &self, + py: Python<'_>, + future: BoxFuture<'static, PyResult>, + ) -> PyResult { + let handle = self.runtime.tokio_handle()?; + py.detach(|| handle.block_on(future)) + } + + /// 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 instances_of = |table: Arc<[(i64, Arc)]>| { + InstanceMap::new( + table.iter().cloned().chain( + instance_ids + .iter() + .map(|(name, id)| (*id, Arc::from(name.as_str()))), + ), + ) + }; + let mut instances = instances_of(reader.instance_map().await?); + 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()), + }; + // A stored token may name a writer added since the table + // was cached. + if stored.names_unknown_writer(&instances) { + reader.forget_instance_map(); + instances = instances_of(reader.instance_map().await?); + } + 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>, + 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 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 { + let reader = store.sliding_sync_reader(pool); + let outcomes = + room_entry::generate_room_entries(&reader, &request, &instance_ids, now_ms, &rooms) + .await + .map_err(anyhow_to_py)?; + let stats = SlidingSyncStats::from_db(&reader.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_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>, + state: &Bound<'py, ConnectionState>, + interested: &Bound<'py, PyInterestedRooms>, + ) -> PyResult> { + let future = self.start_room_entries(py, request, state, interested)?; + self.defer(py, future) + } + + /// As [`Self::generate_room_entries`], but blocks. + fn generate_room_entries_blocking( + &self, + py: Python<'_>, + request: &Bound<'_, SlidingSyncRequest>, + state: &Bound<'_, ConnectionState>, + interested: &Bound<'_, PyInterestedRooms>, + ) -> PyResult { + 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) + } +} + +/// What the Rust implementation did for one request, for the +/// `synapse_sliding_sync_rust_*` histograms. Python observes it once per +/// request. +#[pyclass(frozen, get_all, skip_from_py_object)] +#[derive(Debug, Default, Clone, Copy, PartialEq)] +pub struct SlidingSyncStats { + /// Time spent waiting on statements, in seconds. + pub db_time: f64, + /// Time spent waiting for a connection from the pool, in seconds. + pub checkout_time: f64, + /// Transactions begun. + pub transactions: u64, + /// Statements executed. + pub queries: u64, + /// Rooms whose entry Rust computed. + pub rooms_computed: u64, + /// Rooms whose lazy-loaded members were looked up in the connection + /// state. + pub lazy_member_lookups: u64, +} + +impl SlidingSyncStats { + /// The database half of the stats, from a reader's. + pub fn from_db(db: &SyncDbStats) -> Self { + SlidingSyncStats { + db_time: db.db_time.as_secs_f64(), + checkout_time: db.checkout_time.as_secs_f64(), + transactions: db.transactions, + queries: db.queries, + ..Default::default() + } + } +} + +#[pymethods] +impl SlidingSyncStats { + #[new] + #[pyo3(signature = (*, db_time=0.0, checkout_time=0.0, transactions=0, queries=0, rooms_computed=0, lazy_member_lookups=0))] + fn py_new( + db_time: f64, + checkout_time: f64, + transactions: u64, + queries: u64, + rooms_computed: u64, + lazy_member_lookups: u64, + ) -> Self { + SlidingSyncStats { + db_time, + checkout_time, + transactions, + queries, + rooms_computed, + lazy_member_lookups, + } + } +} + +/// Called when registering modules with python. +pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { + let child_module = PyModule::new(py, "sliding_sync")?; + 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)?; + connection::register_test_functions(&child_module)?; + extensions::register_test_functions(&child_module)?; + required_state::register_test_functions(&child_module)?; + room::register_test_functions(&child_module)?; + room_lists::register_test_functions(&child_module)?; + + m.add_submodule(&child_module)?; + + // We need to manually add the module to sys.modules to make `from + // synapse.synapse_rust import sliding_sync` work. + py.import("sys")? + .getattr("modules")? + .set_item("synapse.synapse_rust.sliding_sync", child_module)?; + + Ok(()) +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use super::*; + + #[test] + fn stats_from_db() { + let db = SyncDbStats { + checkouts: 3, + checkout_time: Duration::from_millis(5), + transactions: 2, + queries: 7, + db_time: Duration::from_millis(250), + ..Default::default() + }; + let stats = SlidingSyncStats::from_db(&db); + assert_eq!(stats.transactions, 2); + assert_eq!(stats.queries, 7); + assert!((stats.db_time - 0.25).abs() < 1e-9); + assert!((stats.checkout_time - 0.005).abs() < 1e-9); + assert_eq!(stats.rooms_computed, 0); + } +} diff --git a/rust/src/sliding_sync/request.rs b/rust/src/sliding_sync/request.rs new file mode 100644 index 0000000000..c5afe58e89 --- /dev/null +++ b/rust/src/sliding_sync/request.rs @@ -0,0 +1,548 @@ +/* + * 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: + * . + */ + +//! One sliding sync request as Rust sees it: the request body, plus what +//! Python adds to it. +//! +//! Python validates the body with pydantic (`SlidingSyncBody`), which also +//! produces the error messages. Rust then parses the same bytes with serde +//! into [`SlidingSyncBody`]. The two parses are held equal by +//! `tests/rest/client/sliding_sync/test_rust_request.py`. The types follow +//! the pydantic model field for field, including where it tells an explicit +//! `null` from a missing field (`slow_get_all_rooms`, and an extension's +//! `enabled`, `lists` and `rooms`). +//! +//! Rust only parses bodies pydantic has accepted, so it does not repeat the +//! validation: the `timeline_limit` maximum, the range bounds, the list key +//! length, the number of lists and the format of the `since` tokens. One +//! accepted body does not parse here: an integer outside `i64` (`u64` for a +//! range), which Python takes as an arbitrary-precision integer. Python then +//! computes the request itself. + +use std::collections::{HashMap, HashSet}; +use std::fmt; +use std::marker::PhantomData; +use std::sync::Arc; + +use pyo3::{exceptions::PyValueError, prelude::*}; +use pythonize::pythonize; +use serde::{ + de::{MapAccess, Visitor}, + ser::SerializeSeq, + Deserialize, Deserializer, Serialize, Serializer, +}; + +use crate::types::stream_token::RoomStreamToken; +use crate::visibility::SoftFailedPolicy; + +/// A JSON object read into its entries in document order. A repeated key +/// keeps its first position and its last value, as the Python dict that +/// `json.loads` builds does. +/// +/// Serialises as a list of `[key, value]` pairs, so the order survives the +/// trip to Python in tests. +#[derive(Debug, Clone, PartialEq)] +pub struct OrderedMap(pub Vec<(String, V)>); + +impl Default for OrderedMap { + fn default() -> Self { + OrderedMap(Vec::new()) + } +} + +impl OrderedMap { + pub fn iter(&self) -> impl Iterator { + self.0.iter().map(|(k, v)| (k.as_str(), v)) + } + + pub fn len(&self) -> usize { + self.0.len() + } + + pub fn is_empty(&self) -> bool { + self.0.is_empty() + } + + pub fn get(&self, key: &str) -> Option<&V> { + self.0.iter().find(|(k, _)| k == key).map(|(_, v)| v) + } +} + +impl<'de, V: Deserialize<'de>> Deserialize<'de> for OrderedMap { + fn deserialize>(deserializer: D) -> Result { + struct OrderedMapVisitor(PhantomData); + + impl<'de, V: Deserialize<'de>> Visitor<'de> for OrderedMapVisitor { + type Value = OrderedMap; + + fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result { + formatter.write_str("a JSON object") + } + + fn visit_map>(self, mut access: A) -> Result { + let mut entries: Vec<(String, V)> = Vec::new(); + let mut positions: HashMap = HashMap::new(); + while let Some((key, value)) = access.next_entry::()? { + match positions.get(&key) { + Some(&position) => entries[position].1 = value, + None => { + positions.insert(key.clone(), entries.len()); + entries.push((key, value)); + } + } + } + Ok(OrderedMap(entries)) + } + } + + deserializer.deserialize_map(OrderedMapVisitor(PhantomData)) + } +} + +impl Serialize for OrderedMap { + fn serialize(&self, serializer: S) -> Result { + let mut seq = serializer.serialize_seq(Some(self.0.len()))?; + for entry in &self.0 { + seq.serialize_element(entry)?; + } + seq.end() + } +} + +fn some_false() -> Option { + Some(false) +} + +fn all_lists_or_rooms() -> Option> { + Some(vec!["*".to_owned()]) +} + +fn hundred() -> i64 { + 100 +} + +/// `SlidingSyncBody`: the request body. +#[derive(Debug, Clone, Default, PartialEq, Deserialize, Serialize)] +pub struct SlidingSyncBody { + #[serde(default)] + pub conn_id: Option, + /// The sliding window lists, by the client's name for them. + #[serde(default)] + pub lists: Option>, + /// The room subscriptions, by room ID. + #[serde(default)] + pub room_subscriptions: Option>, + #[serde(default)] + pub extensions: Option, +} + +/// `SlidingSyncBody.SlidingSyncList`. +#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)] +pub struct SlidingSyncList { + /// `(event type, state key)` pairs, with the `*`, `$LAZY` and `$ME` + /// wildcards still in them. + pub required_state: Vec<(String, String)>, + /// At most 1000. Pydantic sets no minimum. + pub timeline_limit: i64, + /// Inclusive `[start, end]` ranges. `None` means the whole list. + #[serde(default)] + pub ranges: Option>, + /// `false` when missing, `None` for an explicit `null`. + #[serde(default = "some_false")] + pub slow_get_all_rooms: Option, + #[serde(default)] + pub filters: Option, +} + +/// `SlidingSyncBody.RoomSubscription`. +#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)] +pub struct RoomSubscription { + pub required_state: Vec<(String, String)>, + pub timeline_limit: i64, +} + +/// `SlidingSyncList.Filters`. Every field is optional, and a missing field +/// means no filter on it. +#[derive(Debug, Clone, Default, PartialEq, Deserialize, Serialize)] +pub struct Filters { + #[serde(default)] + pub is_dm: Option, + #[serde(default)] + pub spaces: Option>, + #[serde(default)] + pub is_encrypted: Option, + #[serde(default)] + pub is_invite: Option, + /// `None` inside the list is a room without a type. + #[serde(default)] + pub room_types: Option>>, + #[serde(default)] + pub not_room_types: Option>>, + #[serde(default)] + pub room_name_like: Option, + #[serde(default)] + pub tags: Option>, + #[serde(default)] + pub not_tags: Option>, +} + +/// `SlidingSyncBody.Extensions`. An unknown extension is ignored. +#[derive(Debug, Clone, Default, PartialEq, Deserialize, Serialize)] +pub struct Extensions { + #[serde(default)] + pub to_device: Option, + #[serde(default)] + pub e2ee: Option, + #[serde(default)] + pub account_data: Option, + #[serde(default)] + pub receipts: Option, + #[serde(default)] + pub typing: Option, + /// MSC4308. Only the unstable name is accepted. + #[serde(default, rename = "io.element.msc4308.thread_subscriptions")] + pub thread_subscriptions: Option, + /// MSC4354. `None` when absent: pydantic rejects an explicit `null`. + #[serde(default, rename = "org.matrix.msc4354.sticky_events")] + pub sticky_events: Option, + /// MSC4262. As `sticky_events`. + #[serde(default, rename = "org.matrix.msc4262.profiles")] + pub profiles: Option, +} + +/// `Extensions.ToDeviceExtension` (MSC3885). +#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)] +pub struct ToDeviceExtension { + #[serde(default = "some_false")] + pub enabled: Option, + #[serde(default = "hundred")] + pub limit: i64, + /// The device inbox stream position, as a string of digits. + #[serde(default)] + pub since: Option, +} + +/// `Extensions.E2eeExtension` (MSC3884). +#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)] +pub struct E2eeExtension { + #[serde(default = "some_false")] + pub enabled: Option, +} + +/// The account data (MSC3959), receipts (MSC3960) and typing (MSC3961) +/// extensions, which share their fields. +#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)] +pub struct ListsAndRoomsExtension { + #[serde(default = "some_false")] + pub enabled: Option, + /// The list names to cover. `["*"]`, every list, when missing. + #[serde(default = "all_lists_or_rooms")] + pub lists: Option>, + /// The room subscriptions to cover. `["*"]`, every one, when missing. + #[serde(default = "all_lists_or_rooms")] + pub rooms: Option>, +} + +/// `Extensions.ThreadSubscriptionsExtension` (MSC4308). +#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)] +pub struct ThreadSubscriptionsExtension { + #[serde(default = "some_false")] + pub enabled: Option, + #[serde(default = "hundred")] + pub limit: i64, +} + +/// `Extensions.StickyEventsExtension` (MSC4354). +#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)] +pub struct StickyEventsExtension { + #[serde(default)] + pub enabled: bool, + #[serde(default = "hundred")] + pub limit: i64, + /// A `sticky_` token, unparsed. + #[serde(default)] + pub since: Option, +} + +/// `Extensions.ProfilesExtension` (MSC4262). +#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)] +pub struct ProfilesExtension { + #[serde(default)] + pub enabled: bool, + #[serde(default)] + pub fields: Option>, +} + +impl SlidingSyncBody { + /// Parse a body pydantic has accepted. + pub fn parse(body: &[u8]) -> Result { + serde_json::from_slice(body) + } +} + +/// Where an incremental request starts: the `pos` the client sent, less +/// the stream keys Rust does not read yet. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct FromToken { + pub room_key: RoomStreamToken, + pub connection_position: i64, +} + +/// What every part of a request's computation reads, fixed for the whole +/// request. +#[derive(Debug)] +pub struct SlidingSyncRequestData { + pub body: SlidingSyncBody, + pub user_id: Arc, + pub device_id: Option>, + /// `None` for an initial request. + pub from: Option, + /// The room key of the token the response is computed up to. + pub to_room_key: RoomStreamToken, + /// The rest of that token's string form, after the room key and its + /// `_`. A room's `prev_batch` is a new room key followed by it. + pub to_suffix: Arc, + /// `Some` for a server admin: what they asked to see of soft-failed + /// events. + pub admin_policy: Option, + /// `have_finished_sliding_sync_background_jobs`. Rust covers only the + /// path that reads the sliding sync tables, so a request made before + /// they are filled is Python's. + pub have_finished_bg_jobs: bool, + /// `rooms_to_exclude_from_sync`. + pub rooms_to_exclude: Arc>>, + pub msc4306_enabled: bool, + pub msc4354_enabled: bool, + /// `include_profile_updates_in_sync`, which gates MSC4262. + pub profiles_enabled: bool, +} + +/// One sliding sync request. Built by Python in +/// `SlidingSyncHandler._build_rust_request` for each computation of the +/// response, as the token it runs up to changes between the first try and +/// the one after waiting. +#[pyclass(frozen, skip_from_py_object)] +pub struct SlidingSyncRequest { + pub(crate) data: Arc, +} + +#[pymethods] +impl SlidingSyncRequest { + /// Raises `ValueError` if the body does not parse. + #[new] + #[pyo3(signature = (body, *, user_id, device_id, from_room_key, connection_position, to_room_key, to_suffix, is_server_admin, return_soft_failed_events, return_policy_server_spammy_events, have_finished_bg_jobs, rooms_to_exclude, msc4306_enabled, msc4354_enabled, profiles_enabled))] + #[allow(clippy::too_many_arguments)] + fn py_new( + body: &[u8], + user_id: &str, + device_id: Option<&str>, + from_room_key: Option, + connection_position: Option, + to_room_key: RoomStreamToken, + to_suffix: &str, + is_server_admin: bool, + return_soft_failed_events: bool, + return_policy_server_spammy_events: bool, + have_finished_bg_jobs: bool, + rooms_to_exclude: HashSet, + msc4306_enabled: bool, + msc4354_enabled: bool, + profiles_enabled: bool, + ) -> PyResult { + let body = SlidingSyncBody::parse(body) + .map_err(|e| PyValueError::new_err(format!("parsing the sliding sync body: {e}")))?; + let from = match (from_room_key, connection_position) { + (Some(room_key), Some(connection_position)) => Some(FromToken { + room_key, + connection_position, + }), + (None, None) => None, + _ => { + return Err(PyValueError::new_err( + "from_room_key and connection_position go together", + )) + } + }; + let admin_policy = is_server_admin.then_some(SoftFailedPolicy { + return_soft_failed_events, + return_policy_server_spammy_events, + }); + Ok(SlidingSyncRequest { + data: Arc::new(SlidingSyncRequestData { + body, + user_id: Arc::from(user_id), + device_id: device_id.map(Arc::from), + from, + to_room_key, + to_suffix: Arc::from(to_suffix), + admin_policy, + have_finished_bg_jobs, + rooms_to_exclude: Arc::new(rooms_to_exclude.into_iter().map(Arc::from).collect()), + msc4306_enabled, + msc4354_enabled, + profiles_enabled, + }), + }) + } + + /// The parsed body, with every default filled in and each JSON object + /// keyed by name (lists, room subscriptions) as a list of `[name, + /// value]` pairs. For the differential test. + fn body_for_test<'py>(&self, py: Python<'py>) -> PyResult> { + Ok(pythonize(py, &self.data.body)?) + } + + #[getter] + fn user_id(&self) -> &str { + &self.data.user_id + } + + #[getter] + fn device_id(&self) -> Option<&str> { + self.data.device_id.as_deref() + } + + #[getter] + #[allow(clippy::wrong_self_convention)] + fn from_room_key(&self) -> Option { + self.data.from.as_ref().map(|from| from.room_key.clone()) + } + + #[getter] + fn connection_position(&self) -> Option { + self.data.from.as_ref().map(|from| from.connection_position) + } + + #[getter] + fn to_room_key(&self) -> RoomStreamToken { + self.data.to_room_key.clone() + } + + #[getter] + fn to_suffix(&self) -> &str { + &self.data.to_suffix + } + + #[getter] + fn is_server_admin(&self) -> bool { + self.data.admin_policy.is_some() + } + + #[getter] + fn have_finished_bg_jobs(&self) -> bool { + self.data.have_finished_bg_jobs + } + + #[getter] + fn rooms_to_exclude(&self) -> HashSet { + self.data + .rooms_to_exclude + .iter() + .map(|room_id| room_id.to_string()) + .collect() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn empty_body() { + let body = SlidingSyncBody::parse(b"{}").unwrap(); + assert_eq!(body, SlidingSyncBody::default()); + } + + #[test] + fn list_defaults_and_nulls() { + let body = SlidingSyncBody::parse( + br#"{ + "lists": { + "a": {"required_state": [["m.room.name", ""]], "timeline_limit": 5}, + "b": {"required_state": [], "timeline_limit": 0, "slow_get_all_rooms": null, + "ranges": [[0, 9], [20, 29]], "filters": {"room_types": [null, "m.space"]}} + } + }"#, + ) + .unwrap(); + let lists = body.lists.unwrap(); + let a = lists.get("a").unwrap(); + assert_eq!(a.slow_get_all_rooms, Some(false)); + assert_eq!(a.ranges, None); + assert_eq!( + a.required_state, + vec![("m.room.name".to_owned(), String::new())] + ); + let b = lists.get("b").unwrap(); + assert_eq!(b.slow_get_all_rooms, None); + assert_eq!(b.ranges, Some(vec![(0, 9), (20, 29)])); + assert_eq!( + b.filters.as_ref().unwrap().room_types, + Some(vec![None, Some("m.space".to_owned())]) + ); + } + + #[test] + fn ordered_map_keeps_document_order_and_the_last_value() { + let map: OrderedMap = serde_json::from_str(r#"{"z": 1, "a": 2, "z": 3}"#).unwrap(); + assert_eq!(map.0, vec![("z".to_owned(), 3), ("a".to_owned(), 2)]); + } + + #[test] + fn extension_defaults() { + let body = SlidingSyncBody::parse( + br#"{"extensions": { + "to_device": {}, + "account_data": {"enabled": true, "lists": null}, + "io.element.msc4308.thread_subscriptions": {"enabled": true}, + "thread_subscriptions": {"enabled": false}, + "org.matrix.msc4354.sticky_events": {"since": "sticky_4"}, + "unknown": {"enabled": true} + }}"#, + ) + .unwrap(); + let extensions = body.extensions.unwrap(); + let to_device = extensions.to_device.unwrap(); + assert_eq!(to_device.enabled, Some(false)); + assert_eq!(to_device.limit, 100); + let account_data = extensions.account_data.unwrap(); + assert_eq!(account_data.lists, None); + assert_eq!(account_data.rooms, Some(vec!["*".to_owned()])); + assert_eq!(extensions.thread_subscriptions.unwrap().enabled, Some(true)); + let sticky = extensions.sticky_events.unwrap(); + assert!(!sticky.enabled); + assert_eq!(sticky.since.as_deref(), Some("sticky_4")); + assert_eq!(extensions.profiles, None); + assert_eq!(extensions.e2ee, None); + } + + #[test] + fn unknown_fields_are_ignored() { + let body = SlidingSyncBody::parse( + br#"{"conn_id": "c", "foo": [1, 2], + "room_subscriptions": {"!r:x": {"required_state": [], "timeline_limit": 1, "bar": null}}}"#, + ) + .unwrap(); + assert_eq!(body.conn_id.as_deref(), Some("c")); + assert_eq!(body.room_subscriptions.unwrap().len(), 1); + } + + #[test] + fn out_of_range_integers_do_not_parse() { + assert!(SlidingSyncBody::parse( + br#"{"lists": {"a": {"required_state": [], "timeline_limit": 1, "ranges": [[0, 18446744073709551616]]}}}"# + ) + .is_err()); + } +} diff --git a/rust/src/sliding_sync/required_state.rs b/rust/src/sliding_sync/required_state.rs new file mode 100644 index 0000000000..8a5eeeaa03 --- /dev/null +++ b/rust/src/sliding_sync/required_state.rs @@ -0,0 +1,756 @@ +/* + * 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: + * . + */ + +//! What a room's `required_state` asks for: the state filter and member +//! sets it derives from the timeline (`_required_state_request`), and how +//! it changed since the connection last sent the room +//! (`_required_state_changes`). + +use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; +use std::sync::Arc; + +use pyo3::{prelude::*, types::PyModule}; + +use crate::sliding_sync::config::{RequiredStateMap, LAZY, ME, MEMBER, WILDCARD}; +use crate::types::state_filter::StateFilter; + +/// `MAX_NUMBER_PREVIOUS_STATE_KEYS_TO_REMEMBER`: the most state keys of one +/// type the connection remembers having sent. +pub const MAX_NUMBER_PREVIOUS_STATE_KEYS_TO_REMEMBER: usize = 100; + +const ROOM_AVATAR: &str = "m.room.avatar"; + +/// `_RequiredStateChangesReturn`. +#[derive(Debug, Clone, PartialEq)] +pub struct RequiredStateChanges { + /// The required state map to store in the room config, or `None` if it + /// is unchanged. + pub changed_required_state_map: Option, + /// The current state to fetch on top of the deltas, for what the + /// request newly asks for. + pub added_state_filter: StateFilter, + /// Users sent down explicitly before, now lazy-loaded, to add to the + /// lazy members table. + pub extra_users_to_add_to_lazy_cache: HashSet, + /// Users whose membership changed but is not sent down, to remove from + /// the lazy members table. + pub lazy_members_invalidated: HashSet, +} + +fn from_types(added: &[(String, Option)]) -> StateFilter { + StateFilter::from_types(added.iter().map(|(t, k)| (t.as_str(), k.as_deref()))) +} + +/// `_required_state_changes`: compare the required state the connection +/// last stored for the room with the request's. +/// +/// `state_deltas` are the keys of the room's state deltas in the token +/// range, in the order of Python's dict. +/// +/// Two places depend on iteration order in Python, and the Rust picks one +/// of Python's possible outcomes: +/// - The state keys kept when a type has more than +/// [`MAX_NUMBER_PREVIOUS_STATE_KEYS_TO_REMEMBER`]: Python slices a set, +/// Rust takes the first in sorted order. +/// - The check for explicit members converted to lazy ones reads a loop +/// variable that leaked from the loops before it. With state deltas +/// that is the request's state keys for the last changed type (the +/// order of the deltas). Without, Python's is the last type of a set +/// union, in hash order, and Rust takes `m.room.member`'s. +pub fn required_state_changes( + user_id: &str, + prev_required_state_map: &RequiredStateMap, + request_required_state_map: &RequiredStateMap, + previously_returned_lazy_user_ids: &HashSet, + request_lazy_load_user_ids: &HashSet, + state_deltas: &[(String, String)], +) -> RequiredStateChanges { + let empty = BTreeSet::new(); + + // Lazy members whose membership changed but is not sent down. + let lazy_members_invalidated: HashSet = state_deltas + .iter() + .filter(|(event_type, state_key)| { + event_type == MEMBER + && !request_lazy_load_user_ids.contains(state_key) + && previously_returned_lazy_user_ids.contains(state_key) + }) + .map(|(_, state_key)| state_key.clone()) + .collect(); + + if prev_required_state_map == request_required_state_map { + // Only lazy members can need fetching. + let added: Vec<(String, Option)> = request_lazy_load_user_ids + .difference(previously_returned_lazy_user_ids) + .map(|user_id| (MEMBER.to_owned(), Some(user_id.clone()))) + .collect(); + return RequiredStateChanges { + changed_required_state_map: None, + added_state_filter: if added.is_empty() { + StateFilter::none() + } else { + from_types(&added) + }, + extra_users_to_add_to_lazy_cache: HashSet::new(), + lazy_members_invalidated, + }; + } + + let prev_wildcard = prev_required_state_map.get(WILDCARD).unwrap_or(&empty); + let request_wildcard = request_required_state_map.get(WILDCARD).unwrap_or(&empty); + + // Everything was fetched before, so nothing is new now. + if prev_wildcard.contains(WILDCARD) { + return RequiredStateChanges { + changed_required_state_map: Some(request_required_state_map.clone()), + added_state_filter: StateFilter::none(), + extra_users_to_add_to_lazy_cache: HashSet::new(), + lazy_members_invalidated, + }; + } + // A wildcard type gained state keys: fetch everything. + if !request_wildcard.is_subset(prev_wildcard) { + return RequiredStateChanges { + changed_required_state_map: Some(request_required_state_map.clone()), + added_state_filter: StateFilter::all(), + extra_users_to_add_to_lazy_cache: HashSet::new(), + lazy_members_invalidated, + }; + } + // A wildcard type only lost state keys. + if !prev_wildcard.is_subset(request_wildcard) { + return RequiredStateChanges { + changed_required_state_map: Some(request_required_state_map.clone()), + added_state_filter: StateFilter::none(), + extra_users_to_add_to_lazy_cache: HashSet::new(), + lazy_members_invalidated, + }; + } + + let mut changes: BTreeMap> = BTreeMap::new(); + let mut added: Vec<(String, Option)> = Vec::new(); + + // The changed state keys by type, the types in the order they first + // appear in the deltas. + let mut changed_types: Vec<&str> = Vec::new(); + let mut changed_types_to_state_keys: HashMap<&str, BTreeSet> = HashMap::new(); + for (event_type, state_key) in state_deltas { + changed_types_to_state_keys + .entry(event_type.as_str()) + .or_insert_with(|| { + changed_types.push(event_type.as_str()); + BTreeSet::new() + }) + .insert(state_key.clone()); + } + let delta_keys: HashSet<(&str, &str)> = state_deltas + .iter() + .map(|(t, k)| (t.as_str(), k.as_str())) + .collect(); + + // First, what has been added. + let all_types: BTreeSet<&String> = prev_required_state_map + .keys() + .chain(request_required_state_map.keys()) + .collect(); + for event_type in all_types { + let old_state_keys = prev_required_state_map.get(event_type).unwrap_or(&empty); + let request_state_keys = request_required_state_map.get(event_type).unwrap_or(&empty); + let changed_state_keys = changed_types_to_state_keys + .get(event_type.as_str()) + .unwrap_or(&empty); + + if old_state_keys == request_state_keys { + continue; + } + if request_state_keys.is_subset(old_state_keys) { + // Nothing added. Removals are below. + continue; + } + + // Adding `$LAZY` is always recorded. + if event_type == MEMBER + && !old_state_keys.contains(LAZY) + && request_state_keys.contains(LAZY) + { + changes.insert(event_type.clone(), request_state_keys.clone()); + continue; + } + + // A removed state key is only forgotten once its state changes. + let invalidated_state_keys: BTreeSet<&String> = old_state_keys + .difference(request_state_keys) + .filter(|key| changed_state_keys.contains(*key)) + .collect(); + let inheritable_previous_state_keys: BTreeSet<&String> = old_state_keys + .iter() + .filter(|key| { + key.as_str() != WILDCARD + && key.as_str() != LAZY + && !invalidated_state_keys.contains(key) + }) + .collect(); + + let mut changed: BTreeSet = request_state_keys.clone(); + changed.extend(inheritable_previous_state_keys.iter().map(|k| (*k).clone())); + if changed.len() > MAX_NUMBER_PREVIOUS_STATE_KEYS_TO_REMEMBER { + changed = request_state_keys.clone(); + if request_state_keys.len() < MAX_NUMBER_PREVIOUS_STATE_KEYS_TO_REMEMBER { + // Python takes an arbitrary slice of the set, without + // leaving out the requested keys first. + let room = MAX_NUMBER_PREVIOUS_STATE_KEYS_TO_REMEMBER - request_state_keys.len(); + changed.extend( + inheritable_previous_state_keys + .iter() + .take(room) + .map(|k| (*k).clone()), + ); + } + } + changes.insert(event_type.clone(), changed); + + if old_state_keys.contains(WILDCARD) { + // Everything of the type was fetched before. + continue; + } + + if request_state_keys.contains(WILDCARD) { + added.push((event_type.clone(), None)); + } else { + for state_key in request_state_keys.difference(old_state_keys) { + if state_key == ME { + added.push((event_type.clone(), Some(user_id.to_owned()))); + } else if state_key == LAZY { + // Lazy members are handled below. + } else if event_type == MEMBER { + // Only explicit members not sent down before. + if !previously_returned_lazy_user_ids.contains(state_key) { + added.push((event_type.clone(), Some(state_key.clone()))); + } + } else { + added.push((event_type.clone(), Some(state_key.clone()))); + } + } + } + } + + let mut previously_required_state_members: HashSet<&str> = prev_required_state_map + .get(MEMBER) + .map(|keys| keys.iter().map(String::as_str).collect()) + .unwrap_or_default(); + if previously_required_state_members.contains(ME) { + previously_required_state_members.insert(user_id); + } + + // The lazy members not returned before, explicitly or lazily. + for required_user_id in request_lazy_load_user_ids { + if previously_returned_lazy_user_ids.contains(required_user_id) + || previously_required_state_members.contains(required_user_id.as_str()) + { + continue; + } + added.push((MEMBER.to_owned(), Some(required_user_id.clone()))); + } + + let added_state_filter = from_types(&added); + + // Then, what has been removed. + for event_type in &changed_types { + let mut changed_state_keys = changed_types_to_state_keys[event_type].clone(); + let old_state_keys = prev_required_state_map.get(*event_type).unwrap_or(&empty); + let request_state_keys = request_required_state_map + .get(*event_type) + .unwrap_or(&empty); + + if old_state_keys == request_state_keys { + continue; + } + + // `$ME` changes with the user's own state. + if changed_state_keys.contains(user_id) { + changed_state_keys.insert(ME.to_owned()); + } + + let invalidated_state_keys: BTreeSet = old_state_keys + .difference(request_state_keys) + .filter(|key| changed_state_keys.contains(*key)) + .cloned() + .collect(); + + // Additions were handled above. + if !request_state_keys.is_subset(old_state_keys) { + continue; + } + + if old_state_keys.contains(WILDCARD) != request_state_keys.contains(WILDCARD) { + changes.insert((*event_type).to_owned(), request_state_keys.clone()); + continue; + } + + if *event_type == MEMBER { + let old_lazy = old_state_keys.contains(LAZY); + let request_lazy = request_state_keys.contains(LAZY); + if old_lazy != request_lazy { + changes.insert((*event_type).to_owned(), request_state_keys.clone()); + continue; + } + if (old_lazy || request_lazy) && !invalidated_state_keys.is_empty() { + changes.insert((*event_type).to_owned(), request_state_keys.clone()); + continue; + } + } + + if !invalidated_state_keys.is_empty() { + changes.insert( + (*event_type).to_owned(), + old_state_keys + .difference(&invalidated_state_keys) + .cloned() + .collect(), + ); + } + } + + // The loop variable Python reads here: see the doc comment. + let leaked_type = changed_types.last().copied().unwrap_or(MEMBER); + let leaked_request_state_keys = request_required_state_map + .get(leaked_type) + .unwrap_or(&empty); + + // Explicit members, now lazy-loaded, that need no fetching. + let mut extra_users_to_add_to_lazy_cache = HashSet::new(); + let membership_changes = changes.get(MEMBER).is_some_and(|keys| !keys.is_empty()); + if membership_changes && leaked_request_state_keys.contains(LAZY) { + for state_key in prev_required_state_map.get(MEMBER).unwrap_or(&empty) { + if state_key == WILDCARD || state_key == LAZY { + continue; + } + let state_key = if state_key == ME { user_id } else { state_key }; + if !delta_keys.contains(&(MEMBER, state_key)) { + extra_users_to_add_to_lazy_cache.insert(state_key.to_owned()); + } + } + } + + let changed_required_state_map = if changes.is_empty() { + None + } else { + let mut map = prev_required_state_map.clone(); + for (event_type, state_keys) in changes { + if state_keys.is_empty() { + map.remove(&event_type); + } else { + map.insert(event_type, state_keys); + } + } + Some(map) + }; + + RequiredStateChanges { + changed_required_state_map, + added_state_filter, + extra_users_to_add_to_lazy_cache, + lazy_members_invalidated, + } +} + +/// One timeline event, as [`required_state_request`] reads it. +pub struct TimelineMember<'a> { + pub sender: &'a str, + pub event_type: &'a str, + pub state_key: Option<&'a str>, +} + +/// `_RequiredStateRequest`. +#[derive(Debug, Clone, PartialEq)] +pub struct RequiredStateRequest { + /// The state the request asks for. + pub required_state_filter: StateFilter, + /// Whether `m.room.member` asks for `$LAZY`. + pub lazy_load_room_members: bool, + /// The lazy-loaded members, less the explicit ones. + pub lazy_load_user_ids: HashSet, + /// The members asked for by user ID (or `$ME`). + pub explicit_user_state: HashSet, +} + +/// `_required_state_request`: the state filter and member sets a room's +/// `required_state` derives, for a user with `membership` in the room and +/// the timeline returned. `state_delta_keys` are the keys of the room's +/// state deltas. +pub fn required_state_request<'a>( + user_id: &str, + membership: &str, + required_state_map: &RequiredStateMap, + timeline: &[TimelineMember<'a>], + limited: Option, + initial: bool, + state_delta_keys: impl IntoIterator, +) -> RequiredStateRequest { + let mut explicit_user_state: HashSet = HashSet::new(); + let mut lazy_load_user_ids: HashSet = HashSet::new(); + let mut lazy_load_room_members = false; + let mut required_state_filter = StateFilter::none(); + + // No required state for invites and knocks, only stripped state. + if membership != "invite" && membership != "knock" { + let wildcard_keys = required_state_map.get(WILDCARD); + if wildcard_keys.is_some() { + // `*` state key of `*`, or a wildcard type, which the state + // filter cannot express: everything. + required_state_filter = StateFilter::all(); + } else { + let mut required_state_types: Vec<(String, Option)> = Vec::new(); + let mut state_delta_keys = Some(state_delta_keys); + for (state_type, state_key_set) in required_state_map { + for state_key in state_key_set { + if state_key == WILDCARD { + required_state_types.push((state_type.clone(), None)); + } else if state_type == MEMBER && state_key == LAZY { + lazy_load_room_members = true; + + // Everyone in the timeline: senders, and the + // targets of membership events. + let mut timeline_membership: HashSet<&str> = HashSet::new(); + for event in timeline { + timeline_membership.insert(event.sender); + if event.event_type == MEMBER { + if let Some(state_key) = event.state_key { + timeline_membership.insert(state_key); + } + } + } + lazy_load_user_ids + .extend(timeline_membership.iter().map(|u| u.to_string())); + + if limited == Some(true) || initial { + for user_id in &timeline_membership { + required_state_types + .push((MEMBER.to_owned(), Some(user_id.to_string()))); + } + } else { + // Every membership change, for clients that + // keep the full member list. + required_state_types.push((MEMBER.to_owned(), None)); + if let Some(keys) = state_delta_keys.take() { + lazy_load_user_ids.extend( + keys.into_iter() + .filter(|(event_type, _)| *event_type == MEMBER) + .map(|(_, state_key)| state_key.to_owned()), + ); + } + } + } else { + let normalized = if state_key == ME { user_id } else { state_key }; + if state_type == MEMBER { + explicit_user_state.insert(normalized.to_owned()); + } + required_state_types + .push((state_type.clone(), Some(normalized.to_owned()))); + } + } + } + required_state_filter = from_types(&required_state_types); + } + } + + for user_id in &explicit_user_state { + lazy_load_user_ids.remove(user_id); + } + + RequiredStateRequest { + required_state_filter, + lazy_load_room_members, + lazy_load_user_ids, + explicit_user_state, + } +} + +/// `_state_filter_with_meta`: `required_state_filter` plus the heroes' +/// memberships and, with `include_avatar`, the room avatar. +pub fn state_filter_with_meta( + required_state_filter: &StateFilter, + hero_user_ids: &[Arc], + include_avatar: bool, +) -> StateFilter { + if *required_state_filter == StateFilter::all() { + return StateFilter::all(); + } + let mut types: Vec<(&str, Option<&str>)> = hero_user_ids + .iter() + .map(|user_id| (MEMBER, Some(&**user_id))) + .collect(); + if include_avatar { + types.push((ROOM_AVATAR, Some(""))); + } + for (event_type, state_keys) in required_state_filter.types() { + match state_keys { + None => types.push((event_type, None)), + Some(keys) => types.extend(keys.iter().map(|key| (&**event_type, Some(&**key)))), + } + } + StateFilter::new( + StateFilter::from_types(types).types().clone(), + required_state_filter.include_others(), + ) +} + +type PyRequiredStateMap = BTreeMap>; +type PyStateFilter = (BTreeMap>>, bool); + +fn map_from_py(map: PyRequiredStateMap) -> RequiredStateMap { + map.into_iter() + .map(|(k, v)| (k, v.into_iter().collect())) + .collect() +} + +fn map_to_py(map: &RequiredStateMap) -> PyRequiredStateMap { + map.iter() + .map(|(k, v)| (k.clone(), v.iter().cloned().collect())) + .collect() +} + +/// A [`StateFilter`] as `({type: state keys or None}, include_others)`. +pub fn state_filter_to_py(filter: &StateFilter) -> PyStateFilter { + ( + filter + .types() + .iter() + .map(|(t, keys)| { + ( + t.to_string(), + keys.as_ref().map(|keys| { + let mut keys: Vec = keys.iter().map(|k| k.to_string()).collect(); + keys.sort(); + keys + }), + ) + }) + .collect(), + filter.include_others(), + ) +} + +fn sorted(set: HashSet) -> Vec { + let mut v: Vec = set.into_iter().collect(); + v.sort(); + v +} + +/// `_required_state_changes` through the Rust implementation. Returns +/// `(changed_required_state_map, added_state_filter, +/// extra_users_to_add_to_lazy_cache, lazy_members_invalidated)`. +#[pyfunction] +#[allow(clippy::type_complexity)] +fn required_state_changes_for_test( + user_id: &str, + prev_required_state_map: PyRequiredStateMap, + request_required_state_map: PyRequiredStateMap, + previously_returned_lazy_user_ids: HashSet, + request_lazy_load_user_ids: HashSet, + state_deltas: Vec<(String, String)>, +) -> ( + Option, + PyStateFilter, + Vec, + Vec, +) { + let changes = required_state_changes( + user_id, + &map_from_py(prev_required_state_map), + &map_from_py(request_required_state_map), + &previously_returned_lazy_user_ids, + &request_lazy_load_user_ids, + &state_deltas, + ); + ( + changes.changed_required_state_map.as_ref().map(map_to_py), + state_filter_to_py(&changes.added_state_filter), + sorted(changes.extra_users_to_add_to_lazy_cache), + sorted(changes.lazy_members_invalidated), + ) +} + +/// `_required_state_request` through the Rust implementation, with the +/// timeline as `(sender, type, state_key)`. Returns `(required_state_filter, +/// lazy_load_room_members, lazy_load_user_ids, explicit_user_state)`. +#[pyfunction] +#[allow(clippy::too_many_arguments)] +fn required_state_request_for_test( + user_id: &str, + membership: &str, + required_state_map: PyRequiredStateMap, + timeline: Vec<(String, String, Option)>, + limited: Option, + initial: bool, + state_delta_keys: Vec<(String, String)>, +) -> (PyStateFilter, bool, Vec, Vec) { + let timeline: Vec> = timeline + .iter() + .map(|(sender, event_type, state_key)| TimelineMember { + sender, + event_type, + state_key: state_key.as_deref(), + }) + .collect(); + let request = required_state_request( + user_id, + membership, + &map_from_py(required_state_map), + &timeline, + limited, + initial, + state_delta_keys + .iter() + .map(|(t, k)| (t.as_str(), k.as_str())), + ); + ( + state_filter_to_py(&request.required_state_filter), + request.lazy_load_room_members, + sorted(request.lazy_load_user_ids), + sorted(request.explicit_user_state), + ) +} + +/// `_state_filter_with_meta` through the Rust implementation. +#[pyfunction] +fn state_filter_with_meta_for_test( + required_state_filter: StateFilter, + hero_user_ids: Vec, + include_avatar: bool, +) -> PyStateFilter { + let heroes: Vec> = hero_user_ids.into_iter().map(Arc::from).collect(); + state_filter_to_py(&state_filter_with_meta( + &required_state_filter, + &heroes, + include_avatar, + )) +} + +pub fn register_test_functions(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_function(wrap_pyfunction!(required_state_changes_for_test, m)?)?; + m.add_function(wrap_pyfunction!(required_state_request_for_test, m)?)?; + m.add_function(wrap_pyfunction!(state_filter_with_meta_for_test, m)?)?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn map(entries: &[(&str, &[&str])]) -> RequiredStateMap { + entries + .iter() + .map(|(t, keys)| (t.to_string(), keys.iter().map(|k| k.to_string()).collect())) + .collect() + } + + fn set(items: &[&str]) -> HashSet { + items.iter().map(|s| s.to_string()).collect() + } + + #[test] + fn unchanged_config_fetches_new_lazy_members() { + let config = map(&[(MEMBER, &[LAZY])]); + let changes = required_state_changes( + "@me:x", + &config, + &config, + &set(&["@a:x"]), + &set(&["@a:x", "@b:x"]), + &[(MEMBER.to_owned(), "@c:x".to_owned())], + ); + assert_eq!(changes.changed_required_state_map, None); + assert_eq!( + changes.added_state_filter, + StateFilter::from_types([(MEMBER, Some("@b:x"))]) + ); + assert!(changes.lazy_members_invalidated.is_empty()); + } + + #[test] + fn removed_key_is_kept_until_its_state_changes() { + let prev = map(&[("m.a", &["x", "y"])]); + let request = map(&[("m.a", &["x"])]); + let none = HashSet::new(); + let changes = required_state_changes("@me:x", &prev, &request, &none, &none, &[]); + assert_eq!(changes.changed_required_state_map, None); + + let deltas = [("m.a".to_owned(), "y".to_owned())]; + let changes = required_state_changes("@me:x", &prev, &request, &none, &none, &deltas); + assert_eq!(changes.changed_required_state_map, Some(request)); + } + + #[test] + fn explicit_members_converted_to_lazy_join_the_lazy_cache() { + let prev = map(&[(MEMBER, &["@a:x", ME])]); + let request = map(&[(MEMBER, &[LAZY])]); + let none = HashSet::new(); + let changes = required_state_changes("@me:x", &prev, &request, &none, &none, &[]); + assert_eq!( + changes.extra_users_to_add_to_lazy_cache, + set(&["@a:x", "@me:x"]) + ); + } + + #[test] + fn remembered_keys_are_capped() { + let old: Vec = (0..150).map(|i| format!("@u{i}:x")).collect(); + let old_refs: Vec<&str> = old.iter().map(String::as_str).collect(); + let prev = map(&[("m.a", &old_refs)]); + let request = map(&[("m.a", &["new"])]); + let none = HashSet::new(); + let changes = required_state_changes("@me:x", &prev, &request, &none, &none, &[]); + let kept = &changes.changed_required_state_map.unwrap()["m.a"]; + assert_eq!(kept.len(), MAX_NUMBER_PREVIOUS_STATE_KEYS_TO_REMEMBER); + assert!(kept.contains("new")); + } + + #[test] + fn lazy_members_from_the_timeline() { + let config = map(&[(MEMBER, &[LAZY, "@a:x"]), ("m.room.name", &[""])]); + let timeline = [ + TimelineMember { + sender: "@a:x", + event_type: "m.room.message", + state_key: None, + }, + TimelineMember { + sender: "@b:x", + event_type: MEMBER, + state_key: Some("@c:x"), + }, + ]; + let request = + required_state_request("@me:x", "join", &config, &timeline, Some(true), false, []); + assert!(request.lazy_load_room_members); + assert_eq!(request.lazy_load_user_ids, set(&["@b:x", "@c:x"])); + assert_eq!(request.explicit_user_state, set(&["@a:x"])); + assert!(request.required_state_filter.matches(MEMBER, "@c:x")); + assert!(request.required_state_filter.matches("m.room.name", "")); + + // An unlimited incremental sync returns every membership change. + let request = required_state_request( + "@me:x", + "join", + &config, + &timeline, + Some(false), + false, + [(MEMBER, "@d:x")], + ); + assert!(request.required_state_filter.matches(MEMBER, "@anyone:x")); + assert!(request.lazy_load_user_ids.contains("@d:x")); + } +} diff --git a/rust/src/sliding_sync/result.rs b/rust/src/sliding_sync/result.rs new file mode 100644 index 0000000000..b542f29d83 --- /dev/null +++ b/rust/src/sliding_sync/result.rs @@ -0,0 +1,255 @@ +/* + * 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: + * . + */ + +//! 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::BTreeSet; +use std::sync::{Arc, Mutex}; + +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::{MutablePerConnectionState, RoomLazyMembershipChanges}; +use crate::sliding_sync::room_entry::{RoomEntry, RoomOutcome}; +use crate::sliding_sync::SlidingSyncStats; +use crate::sync::relations::Bundle; + +/// A `required_state_map` from Python: a mapping of types to sets of +/// state keys. +pub fn required_state_map_from_py(map: &Bound<'_, PyAny>) -> PyResult { + let py = map.py(); + let mut result = RequiredStateMap::new(); + for item in map.call_method0(intern!(py, "items"))?.try_iter()? { + let (state_type, state_keys): (String, Bound<'_, PyAny>) = item?.extract()?; + let mut keys = BTreeSet::new(); + for key in state_keys.try_iter()? { + keys.insert(key?.extract::()?); + } + result.insert(state_type, keys); + } + Ok(result) +} + +/// 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)] +pub struct RoomResult { + pub(crate) entry: Arc, +} + +#[pymethods] +impl RoomResult { + fn __bool__(&self) -> bool { + self.entry.has_updates() + } + + #[getter] + fn room_id(&self) -> &str { + &self.entry.room_id + } + + #[getter] + fn initial(&self) -> bool { + self.entry.initial + } + + #[getter] + fn unstable_expanded_timeline(&self) -> bool { + self.entry.unstable_expanded_timeline + } + + #[getter] + fn timeline_event_ids(&self) -> Vec<&str> { + self.entry + .timeline + .iter() + .map(|filtered| filtered.event.event_id()) + .collect() + } + + #[getter] + fn timeline_senders(&self) -> Vec<&str> { + self.entry + .timeline + .iter() + .map(|filtered| filtered.event.sender()) + .collect() + } + + /// The state keys of the membership events in `required_state`. + #[getter] + fn required_state_member_keys(&self) -> Vec<&str> { + self.entry + .required_state + .iter() + .filter(|event| event.type_() == MEMBER) + .filter_map(Event::state_key) + .collect() + } + + #[getter] + fn hero_user_ids(&self) -> Vec<&str> { + self.entry + .heroes + .iter() + .map(|hero| hero.user_id.as_str()) + .collect() + } + + /// Every event the room's encoding serializes, for the module + /// callbacks that add to their unsigned data. + fn events_to_serialize<'py>(&self, py: Python<'py>) -> PyResult> { + let bundled = self + .entry + .bundles + .iter() + .flat_map(|bundles| bundles.values().flat_map(Bundle::events)); + let events = self + .entry + .timeline + .iter() + .map(|filtered| &filtered.event) + .chain(&self.entry.required_state) + .chain(bundled) + .map(|event| Bound::new(py, event.shallow_clone())) + .collect::>>()?; + PyList::new(py, events) + } +} + +/// What [`crate::sliding_sync::room_entry::generate_room_entries`] made of +/// the rooms. +#[pyclass(frozen, skip_from_py_object)] +pub struct RoomsResult { + /// The entries to send. Taken by [`RoomsResult::take_entries`]. + entries: Mutex>>, + /// The rooms computed with nothing to send. + empty: Vec>, + /// `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<(Arc, RoomLazyMembershipChanges)>, + fallbacks: Vec<(Arc, String)>, + errors: Vec<(Arc, String)>, + stats: SlidingSyncStats, +} + +impl RoomsResult { + /// The result of `outcomes`. An entry with nothing to send is dropped + /// from an incremental response. + pub fn new( + outcomes: Vec<(Arc, RoomOutcome)>, + incremental: bool, + mut stats: SlidingSyncStats, + ) -> Self { + let mut entries = Vec::new(); + let mut empty = Vec::new(); + let mut room_config_updates = Vec::new(); + let mut lazy_memberships = Vec::new(); + let mut fallbacks = Vec::new(); + let mut errors = Vec::new(); + for (room_id, outcome) in outcomes { + match outcome { + RoomOutcome::Computed(computed) => { + stats.rooms_computed += 1; + if computed.lazy_member_lookup { + stats.lazy_member_lookups += 1; + } + let mutations = computed.mutations; + if let Some(config) = mutations.room_config { + room_config_updates.push((Arc::clone(&room_id), config)); + } + lazy_memberships.push((Arc::clone(&room_id), mutations.lazy_membership)); + if computed.entry.has_updates() || !incremental { + entries.push(Arc::new(computed.entry)); + } else { + empty.push(room_id); + } + } + RoomOutcome::Fallback(reason) => fallbacks.push((room_id, reason)), + RoomOutcome::Error(error) => errors.push((room_id, error)), + } + } + RoomsResult { + entries: Mutex::new(entries), + empty, + room_config_updates, + lazy_memberships, + fallbacks, + errors, + 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] +impl RoomsResult { + /// The entries to send as `RoomResult`s. A second call returns an + /// empty list. + fn take_entries<'py>(&self, py: Python<'py>) -> PyResult> { + let entries = std::mem::take(&mut *lock_ignoring_poison(&self.entries)); + let items = entries + .into_iter() + .map(|entry| Bound::new(py, RoomResult { entry })) + .collect::>>()?; + PyList::new(py, items) + } + + #[getter] + fn empty(&self) -> Vec<&str> { + self.empty.iter().map(|room_id| &**room_id).collect() + } + + /// `(room_id, reason)` for each room Python must compute. + #[getter] + fn fallbacks(&self) -> Vec<(&str, &str)> { + self.fallbacks + .iter() + .map(|(room_id, reason)| (&**room_id, reason.as_str())) + .collect() + } + + /// `(room_id, error)` for each room Rust failed on. + #[getter] + fn errors(&self) -> Vec<(&str, &str)> { + self.errors + .iter() + .map(|(room_id, error)| (&**room_id, error.as_str())) + .collect() + } + + #[getter] + fn stats(&self) -> SlidingSyncStats { + self.stats + } +} diff --git a/rust/src/sliding_sync/room.rs b/rust/src/sliding_sync/room.rs new file mode 100644 index 0000000000..93bf10d183 --- /dev/null +++ b/rust/src/sliding_sync/room.rs @@ -0,0 +1,473 @@ +/* + * 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 decisions `get_room_sync_data` makes for one room, apart from its +//! reads: where its updates start, its live events, its stripped state, +//! name, avatar and heroes, its bump stamp, and the config the connection +//! records for it. + +use std::collections::BTreeMap; + +use pyo3::{exceptions::PyValueError, prelude::*, types::PyModule}; +use serde_json::{json, Map, Value}; + +use crate::sliding_sync::config::{ + config_from_py, to_test_tuple, RequiredStateMap, RoomSyncConfig, +}; +use crate::sliding_sync::connection::HaveSentRoom; +use crate::types::stream_token::{persisted_after, RoomStreamToken}; + +/// `SLIDING_SYNC_DEFAULT_BUMP_EVENT_TYPES`: the event types that move a +/// room up the list. +pub const BUMP_EVENT_TYPES: [&str; 7] = [ + "m.room.create", + "m.room.message", + "m.room.encrypted", + "m.sticker", + "m.call.invite", + "m.poll.start", + "m.beacon_info", +]; + +/// `_timeline_bounds`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TimelineBounds { + /// The lower bound of the room's updates, `None` for the whole room. + pub from_bound: Option, + /// Whether the room is sent in full. + pub initial: bool, + /// Whether the timeline ignores `from_bound`, as the timeline limit + /// grew (`unstable_expanded_timeline`). + pub ignore_timeline_bound: bool, +} + +/// `_timeline_bounds`: where the room's updates start, from what the +/// connection sent before. `from_room_key` is the `pos` token's room key, +/// and `skip_previous_bounds` sends the room as if never sent, when newly +/// joined or state reset out of it. +pub fn timeline_bounds( + room_status: &HaveSentRoom, + from_room_key: Option<&RoomStreamToken>, + skip_previous_bounds: bool, + prev_room_sync_config: Option<&RoomSyncConfig>, + room_sync_config: &RoomSyncConfig, +) -> TimelineBounds { + let mut bounds = TimelineBounds { + from_bound: None, + initial: true, + ignore_timeline_bound: false, + }; + let Some(from_room_key) = from_room_key else { + return bounds; + }; + if skip_previous_bounds { + return bounds; + } + match room_status { + HaveSentRoom::Live => { + bounds.from_bound = Some(from_room_key.clone()); + bounds.initial = false; + } + HaveSentRoom::Previously(last_token) => { + bounds.from_bound = Some(last_token.clone()); + bounds.initial = false; + } + HaveSentRoom::Never => {} + } + if let Some(prev) = prev_room_sync_config { + if prev.timeline_limit < room_sync_config.timeline_limit { + bounds.ignore_timeline_bound = true; + } + } + bounds +} + +/// `num_live`: how many of the timeline's last events are after the `pos` +/// token, given each event's `(instance_name, stream_ordering)` in +/// timeline order. Zero on an initial request. +pub fn num_live<'a>( + positions: impl DoubleEndedIterator, + from_room_key: Option<&RoomStreamToken>, +) -> i64 { + let Some(from_room_key) = from_room_key else { + return 0; + }; + positions + .rev() + .take_while(|(instance_name, stream)| { + persisted_after(instance_name, *stream, from_room_key) + }) + .count() as i64 +} + +/// `strip_event`: the parts of a state event stripped state carries. +pub fn strip_event(event_type: &str, state_key: &str, content: &Value, sender: &str) -> Value { + json!({ + "type": event_type, + "state_key": state_key, + "content": content, + "sender": sender, + }) +} + +/// `_stripped_state`: the stripped state of an invite or knock, from the +/// event's `membership` and `unsigned`, then `stripped_event`, the event +/// itself stripped. A non-list `invite_room_state`/`knock_room_state` is +/// ignored. +pub fn stripped_state( + membership: Option<&str>, + unsigned: &Map, + stripped_event: Value, +) -> Vec { + let key = match membership { + Some("invite") => Some("invite_room_state"), + Some("knock") => Some("knock_room_state"), + _ => None, + }; + let mut stripped = match key.and_then(|key| unsigned.get(key)) { + Some(Value::Array(state)) => state.clone(), + _ => Vec::new(), + }; + stripped.push(stripped_event); + stripped +} + +/// A string field of an event's content: `None` when missing or not a +/// string, as the Python's `isinstance(value, str)` checks. +pub fn string_field<'a>(content: &'a Map, key: &str) -> Option<&'a str> { + content.get(key).and_then(Value::as_str) +} + +/// Whether the room needs heroes: it has no name (an empty one counts as +/// none), and is sent in full or its membership changed. +pub fn needs_heroes(room_name: Option<&str>, initial: bool, membership_changed: bool) -> bool { + room_name.is_none_or(str::is_empty) && (initial || membership_changed) +} + +/// `SlidingSyncResult.RoomResult.StrippedHero`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct StrippedHero { + pub user_id: String, + pub display_name: Option, + pub avatar_url: Option, +} + +/// The heroes of `hero_user_ids` that `member_content` finds a membership +/// event for, with its display name and avatar. +pub fn assemble_heroes<'a>( + hero_user_ids: &[&str], + member_content: impl Fn(&str) -> Option<&'a Map>, +) -> Vec { + hero_user_ids + .iter() + .filter_map(|user_id| { + let content = member_content(user_id)?; + Some(StrippedHero { + user_id: (*user_id).to_owned(), + display_name: string_field(content, "displayname").map(str::to_owned), + avatar_url: string_field(content, "avatar_url").map(str::to_owned), + }) + }) + .collect() +} + +/// The first half of `_get_bump_stamp`: the stream ordering of the +/// timeline's last bump event, given each event's `(type, +/// stream_ordering)` in timeline order. A backfilled event's negative +/// stream ordering is skipped. +pub fn timeline_bump_stamp<'a>( + timeline: impl DoubleEndedIterator, +) -> Option { + timeline + .rev() + .find(|(event_type, stream)| BUMP_EVENT_TYPES.contains(event_type) && *stream > 0) + .map(|(_, stream)| stream) +} + +/// What `_get_bump_stamp` does with the room's `bump_stamp` from the +/// sliding sync tables. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum StoreBumpStamp { + /// The answer. + Done(Option), + /// Look up the room's last bump event before the `to_token`, and pass + /// its stream ordering to [`bump_stamp_from_last_event`]. + LookupLastEvent, +} + +/// The second half of `_get_bump_stamp`, outside the timeline: +/// `latest_room_bump_stamp` is the tables' `bump_stamp`, and +/// `min_to_token_position` the `to_token`'s minimum stream position. +pub fn store_bump_stamp( + latest_room_bump_stamp: Option, + have_finished_bg_jobs: bool, + min_to_token_position: i64, +) -> StoreBumpStamp { + match latest_room_bump_stamp { + None if have_finished_bg_jobs => StoreBumpStamp::Done(None), + // Only a bump stamp wholly before the token can be trusted. + Some(stamp) if stamp < min_to_token_position => { + StoreBumpStamp::Done(Some(stamp).filter(|stamp| *stamp > 0)) + } + _ => StoreBumpStamp::LookupLastEvent, + } +} + +/// The bump stamp from the room's last bump event before the `to_token`. +pub fn bump_stamp_from_last_event(last_event_stream: Option) -> Option { + last_event_stream.filter(|stream| *stream > 0) +} + +/// `always_return_bump_stamp`: whether the response must carry a bump +/// stamp even when it is unchanged. `limited` is `None` with no timeline. +pub fn always_return_bump_stamp(membership: &str, limited: Option, initial: bool) -> bool { + membership != "join" || limited != Some(false) || initial +} + +/// The bump stamp the room's entry carries: `new_bump_stamp` from +/// `_get_bump_stamp` (for joined rooms), else the membership event's +/// position if one must be returned, and never negative. +pub fn final_bump_stamp( + new_bump_stamp: Option, + always_return: bool, + membership_stream: i64, +) -> Option { + let bump_stamp = match new_bump_stamp { + Some(stamp) => Some(stamp), + None if always_return => Some(membership_stream), + None => None, + }; + // A negative stream ordering cannot be compared with positive ones. It + // only happens for a membership in a room left, deleted and rejoined. + bump_stamp.map(|stamp| stamp.max(0)) +} + +/// `_room_config_to_persist`: the config the connection records for the +/// room after sending it, or `None` to keep the stored one. An empty +/// `changed_required_state_map` records the request's map, but counts as a +/// change, as in Python. +pub fn room_config_to_persist( + room_sync_config: &RoomSyncConfig, + prev_room_sync_config: Option<&RoomSyncConfig>, + changed_required_state_map: Option<&RequiredStateMap>, + ignore_timeline_bound: bool, + limited: Option, +) -> Option { + let required_state_map = match changed_required_state_map { + Some(changed) if !changed.is_empty() => changed.clone(), + _ => room_sync_config.required_state_map.clone(), + }; + let config = RoomSyncConfig { + timeline_limit: room_sync_config.timeline_limit, + required_state_map, + }; + if ignore_timeline_bound { + return Some(config); + } + let Some(prev) = prev_room_sync_config else { + return Some(config); + }; + // A smaller limit is only recorded when the timeline was limited: the + // client then lacks the events a later larger limit would send. + if limited == Some(true) && prev.timeline_limit > room_sync_config.timeline_limit { + return Some(config); + } + if changed_required_state_map.is_some() { + return Some(config); + } + None +} + +type PyConfig = (i64, BTreeMap>); + +/// `_timeline_bounds` through the Rust implementation, the status as +/// `(flag, token)` and tokens as live stream positions. Returns +/// `(from_bound, initial, ignore_timeline_bound)`. +#[pyfunction] +fn timeline_bounds_for_test( + room_status: (String, Option), + from_stream: Option, + skip_previous_bounds: bool, + prev_room_sync_config: Option, + room_sync_config: PyConfig, +) -> PyResult<(Option, bool, bool)> { + let status = match (room_status.0.as_str(), room_status.1) { + ("never", _) => HaveSentRoom::Never, + ("live", _) => HaveSentRoom::Live, + ("previously", Some(token)) => HaveSentRoom::Previously(RoomStreamToken::live(token)), + _ => return Err(PyValueError::new_err("bad status")), + }; + let from = from_stream.map(RoomStreamToken::live); + let prev = prev_room_sync_config.map(|(limit, map)| config_from_py(limit, map)); + let bounds = timeline_bounds( + &status, + from.as_ref(), + skip_previous_bounds, + prev.as_ref(), + &config_from_py(room_sync_config.0, room_sync_config.1), + ); + Ok(( + bounds.from_bound.map(|token| token.stream), + bounds.initial, + bounds.ignore_timeline_bound, + )) +} + +/// `_stripped_state` through the Rust implementation, with the event as +/// JSON. Returns the stripped state as JSON. +#[pyfunction] +fn stripped_state_for_test(event: &str) -> PyResult { + let event: Value = + serde_json::from_str(event).map_err(|e| PyValueError::new_err(e.to_string()))?; + let empty = Map::new(); + let content = event + .get("content") + .cloned() + .unwrap_or(Value::Object(Map::new())); + let membership = content.get("membership").and_then(Value::as_str); + let unsigned = event + .get("unsigned") + .and_then(Value::as_object) + .unwrap_or(&empty); + let field = |key: &str| event.get(key).and_then(Value::as_str).unwrap_or_default(); + let stripped = strip_event(field("type"), field("state_key"), &content, field("sender")); + Ok(Value::Array(stripped_state(membership, unsigned, stripped)).to_string()) +} + +/// `_get_bump_stamp` through the Rust implementation, the timeline as +/// `(type, stream_ordering)`, the store as `latest_room_bump_stamp`, +/// `have_finished_bg_jobs` and `last_event_stream`. +#[pyfunction] +fn get_bump_stamp_for_test( + timeline: Vec<(String, i64)>, + check_outside_timeline: bool, + latest_room_bump_stamp: Option, + have_finished_bg_jobs: bool, + to_stream: i64, + last_event_stream: Option, +) -> Option { + if let Some(stamp) = timeline_bump_stamp(timeline.iter().map(|(t, s)| (t.as_str(), *s))) { + return Some(stamp); + } + if !check_outside_timeline { + return None; + } + match store_bump_stamp(latest_room_bump_stamp, have_finished_bg_jobs, to_stream) { + StoreBumpStamp::Done(stamp) => stamp, + StoreBumpStamp::LookupLastEvent => bump_stamp_from_last_event(last_event_stream), + } +} + +/// `_room_config_to_persist` through the Rust implementation. +#[pyfunction] +fn room_config_to_persist_for_test( + room_sync_config: PyConfig, + prev_room_sync_config: Option, + changed_required_state_map: Option>>, + ignore_timeline_bound: bool, + limited: Option, +) -> Option { + let changed = changed_required_state_map.map(|map| config_from_py(0, map).required_state_map); + let prev = prev_room_sync_config.map(|(limit, map)| config_from_py(limit, map)); + room_config_to_persist( + &config_from_py(room_sync_config.0, room_sync_config.1), + prev.as_ref(), + changed.as_ref(), + ignore_timeline_bound, + limited, + ) + .map(|config| to_test_tuple(&config)) +} + +pub fn register_test_functions(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_function(wrap_pyfunction!(timeline_bounds_for_test, m)?)?; + m.add_function(wrap_pyfunction!(stripped_state_for_test, m)?)?; + m.add_function(wrap_pyfunction!(get_bump_stamp_for_test, m)?)?; + m.add_function(wrap_pyfunction!(room_config_to_persist_for_test, m)?)?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn live_events_counted_from_the_end() { + let from = RoomStreamToken::live(5); + let positions = [ + ("master", 3), + ("master", 6), + ("master", 4), + ("master", 7), + ("master", 8), + ]; + assert_eq!(num_live(positions.iter().copied(), Some(&from)), 2); + assert_eq!(num_live(positions.iter().copied(), None), 0); + } + + #[test] + fn names_and_heroes() { + let content = json!({"name": "x", "url": 3, "displayname": "Alice"}); + let content = content.as_object().unwrap(); + assert_eq!(string_field(content, "name"), Some("x")); + assert_eq!(string_field(content, "url"), None); + assert!(needs_heroes(Some(""), true, false)); + assert!(!needs_heroes(Some("x"), true, true)); + assert!(!needs_heroes(None, false, false)); + + let heroes = assemble_heroes(&["@a:x", "@b:x"], |user_id| { + (user_id == "@a:x").then_some(content) + }); + assert_eq!( + heroes, + vec![StrippedHero { + user_id: "@a:x".to_owned(), + display_name: Some("Alice".to_owned()), + avatar_url: None, + }] + ); + } + + #[test] + fn bump_stamps() { + let timeline = [ + ("m.room.message", 4), + ("m.room.member", 9), + ("m.room.message", -2), + ]; + assert_eq!(timeline_bump_stamp(timeline.iter().copied()), Some(4)); + assert!(always_return_bump_stamp("leave", Some(false), false)); + assert!(always_return_bump_stamp("join", None, false)); + assert!(!always_return_bump_stamp("join", Some(false), false)); + assert_eq!(final_bump_stamp(None, true, -5), Some(0)); + assert_eq!(final_bump_stamp(None, false, 5), None); + assert_eq!(final_bump_stamp(Some(7), false, 5), Some(7)); + } + + #[test] + fn stripped_state_from_unsigned() { + let unsigned = + json!({"invite_room_state": [{"type": "m.room.name"}], "knock_room_state": 3}); + let unsigned = unsigned.as_object().unwrap(); + let me = json!({"type": "m.room.member"}); + assert_eq!( + stripped_state(Some("invite"), unsigned, me.clone()).len(), + 2 + ); + assert_eq!( + stripped_state(Some("knock"), unsigned, me.clone()), + vec![me] + ); + } +} diff --git a/rust/src/sliding_sync/room_entry.rs b/rust/src/sliding_sync/room_entry.rs new file mode 100644 index 0000000000..f6edd5d462 --- /dev/null +++ b/rust/src/sliding_sync/room_entry.rs @@ -0,0 +1,974 @@ +/* + * 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: + * . + */ + +//! `get_room_sync_data`: one room's entry in a sliding sync response, in +//! the order of the Python, and the batch of rooms Python hands over in +//! one call. +//! +//! The pure decisions are in [`crate::sliding_sync::room`] and +//! [`crate::sliding_sync::required_state`]. This module does the reads +//! between them. + +use std::collections::{HashMap, HashSet}; +use std::panic::AssertUnwindSafe; +use std::sync::Arc; + +use anyhow::{bail, Context}; +use futures::{stream, FutureExt, StreamExt}; +use indexmap::IndexMap; +use serde_json::{Map, Value}; + +use crate::events::Event; +use crate::sliding_sync::config::{RoomSyncConfig, MEMBER}; +use crate::sliding_sync::connection::{HaveSentRoom, RoomLazyMembershipChanges}; +use crate::sliding_sync::request::SlidingSyncRequestData; +use crate::sliding_sync::required_state::{ + required_state_changes, required_state_request, state_filter_with_meta, TimelineMember, +}; +use crate::sliding_sync::room::{ + always_return_bump_stamp, bump_stamp_from_last_event, final_bump_stamp, needs_heroes, num_live, + room_config_to_persist, store_bump_stamp, strip_event, stripped_state, timeline_bounds, + timeline_bump_stamp, StoreBumpStamp, StrippedHero, TimelineBounds, BUMP_EVENT_TYPES, +}; +use crate::sliding_sync::room_lists::RoomForUser; +use crate::state::{lookup, StateKey, StateMap}; +use crate::storage::events::RedactBehaviour; +use crate::storage::sliding_sync_store::SlidingSyncReader; +use crate::storage::sync_store::{PaginationOrder, RoomSummary}; +use crate::sync::client_events::{filter_events_for_client, ClientFilterContext}; +use crate::sync::relations::{bundled_aggregations, Bundle}; +use crate::sync::room::Fallback; +use crate::sync::rooms::panic_message; +use crate::sync::summary::extract_heroes_from_room_summary; +use crate::types::state_filter::StateFilter; +use crate::types::stream_token::{Direction, RoomStreamToken}; +use crate::visibility::FilteredEvent; + +const M_ROOM_NAME: &str = "m.room.name"; +const M_ROOM_AVATAR: &str = "m.room.avatar"; + +/// How many rooms are computed at once, as `concurrently_execute` does. +const ROOM_ENTRY_CONCURRENCY: usize = 20; + +/// What `current_sync_for_user` passes `get_room_sync_data` for one room, +/// besides the request. +#[derive(Debug, Clone)] +pub struct RoomInputData { + /// The user's membership in the room at the `to_token`. + pub membership: RoomForUser, + pub room_sync_config: RoomSyncConfig, + /// What the connection last recorded for the room. + pub prev_room_sync_config: Option, + /// Whether the connection sent the room before. + pub room_status: HaveSentRoom, + pub newly_joined: bool, + pub newly_left: bool, + pub is_dm: bool, +} + +/// `SlidingSyncResult.RoomResult`, holding the events Rust loaded. +pub struct RoomEntry { + pub room_id: Arc, + /// May be empty, which counts as an update but is not sent. + pub name: Option, + pub avatar: Option, + pub heroes: Vec, + pub is_dm: bool, + pub initial: bool, + pub unstable_expanded_timeline: bool, + pub required_state: Vec, + pub timeline: Vec, + /// `None` unless the timeline is limited. + pub bundles: Option, Bundle>>, + /// The redactions of the redacted events among the room's, which the + /// serializer adds to their unsigned data. + pub redactions: HashMap, Event>, + pub stripped_state: Vec, + /// The `prev_batch` token as a string. + pub prev_batch: Option, + pub limited: Option, + pub num_live: Option, + pub bump_stamp: Option, + pub joined_count: Option, + pub invited_count: Option, +} + +impl RoomEntry { + /// `RoomResult.__bool__`: whether the entry holds an update to return + /// at once. + pub fn has_updates(&self) -> bool { + self.initial + || self.name.is_some() + || self.avatar.is_some() + || !self.heroes.is_empty() + || self.joined_count.is_some() + || self.invited_count.is_some() + || !self.required_state.is_empty() + || !self.timeline.is_empty() + || !self.stripped_state.is_empty() + } +} + +/// What `get_room_sync_data` writes to the new connection state for the +/// room. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct RoomMutations { + /// `room_configs[room_id]`, if it changes. + pub room_config: Option, + /// `room_lazy_membership[room_id]`, always set. + pub lazy_membership: RoomLazyMembershipChanges, +} + +/// One room computed. +pub struct RoomComputed { + pub entry: RoomEntry, + pub mutations: RoomMutations, + /// Whether the room's lazy-loaded members were looked up in the + /// connection state. + pub lazy_member_lookup: bool, +} + +/// How one room came out. +pub enum RoomOutcome { + Computed(Box), + /// Python must compute the room: [`Fallback`]. + Fallback(String), + /// Computing the room failed. + Error(String), +} + +/// What computing one room reads, shared by the rooms of a request. +pub struct RoomContext<'r, 'a> { + pub reader: &'r SlidingSyncReader<'a>, + pub request: &'r SlidingSyncRequestData, + pub client_filter: ClientFilterContext, + /// `get_id_for_instance` for the writers the rooms' tokens may name. + pub instance_ids: &'r HashMap, + /// The time the visibility checks compare with. + pub now_ms: i64, +} + +/// One room's view of the request: the room, and where the user's view of +/// it ends. +struct Room<'r> { + input: &'r RoomInputData, + room_id: &'r str, + membership: &'r str, + /// Where a leave or ban stops what the user sees: the membership's + /// position, as `event_pos.to_room_stream_token()`. + membership_token: RoomStreamToken, +} + +impl Room<'_> { + fn has_left(&self) -> bool { + matches!(self.membership, "leave" | "ban") + } +} + +/// The events of `event_ids`, by state key. Events that are not found or +/// rejected are left out. +/// +/// An event of the room's timeline is that event, with the previous content +/// the timeline's load gave it. Python's loads find the same objects through +/// its map of the events still held, whatever the event cache holds. +async fn state_events( + reader: &SlidingSyncReader<'_>, + timeline: &HashMap<&str, &Arc>, + event_ids: &[Arc], +) -> anyhow::Result> { + let (held, to_load): (Vec<&Arc>, Vec<&Arc>) = event_ids + .iter() + .partition(|event_id| timeline.contains_key(&***event_id)); + let to_load: Vec> = to_load.into_iter().cloned().collect(); + let loaded = reader + .events(&to_load, RedactBehaviour::Redact, false, false) + .await?; + let held = held + .into_iter() + .map(|event_id| timeline[&**event_id].shallow_clone()); + Ok(loaded + .into_iter() + .chain(held) + .filter_map(|event| { + let key = StateKey::new(event.type_(), event.state_key()?); + Some((key, event)) + }) + .collect()) +} + +/// `get_event`: the event, which must exist. +async fn one_event(reader: &SlidingSyncReader<'_>, event_id: &Arc) -> anyhow::Result { + reader + .events( + std::slice::from_ref(event_id), + RedactBehaviour::Redact, + false, + false, + ) + .await? + .into_iter() + .next() + .with_context(|| format!("Could not find event {event_id}")) +} + +/// `get_current_state_ids_at`: the room's state as the user sees it, at +/// their leave or ban, or the current state. +async fn state_ids_at( + context: &RoomContext<'_, '_>, + room: &Room<'_>, + state_filter: &StateFilter, +) -> anyhow::Result { + if room.has_left() { + context + .reader + .state_ids_at(room.room_id, &room.membership_token, state_filter) + .await + } else { + context + .reader + .current_state_ids(room.room_id, state_filter) + .await + } +} + +/// `get_current_state_at`: as [`state_ids_at`], with the events. +async fn state_at( + context: &RoomContext<'_, '_>, + room: &Room<'_>, + timeline: &HashMap<&str, &Arc>, + state_filter: &StateFilter, +) -> anyhow::Result> { + let state_ids = state_ids_at(context, room, state_filter).await?; + let event_ids: Vec> = state_ids.into_values().collect(); + state_events(context.reader, timeline, &event_ids).await +} + +/// `filter` with the room name added, as the first send reads it. +fn with_room_name(filter: StateFilter) -> StateFilter { + if filter.matches(M_ROOM_NAME, "") { + return filter; + } + let mut types = filter.types().clone(); + types + .entry(Arc::from(M_ROOM_NAME)) + .or_insert_with(|| Some(HashSet::new())) + .get_or_insert_with(HashSet::new) + .insert(Arc::from("")); + StateFilter::new(types, filter.include_others()) +} + +/// `_stripped_state` of an invite or knock event. +fn stripped_state_of(event: &Event) -> anyhow::Result> { + let mut unsigned = Map::new(); + let unsigned_ref = event.unsigned_ref(); + if let Some(state) = unsigned_ref.invite_room_state()? { + unsigned.insert("invite_room_state".to_owned(), Value::Array(state)); + } + if let Some(state) = unsigned_ref.knock_room_state()? { + unsigned.insert("knock_room_state".to_owned(), Value::Array(state)); + } + let content = Value::Object( + event + .content() + .as_map() + .iter() + .map(|(key, value)| (key.to_string(), value.clone())) + .collect(), + ); + let membership = event + .content() + .get_field("membership") + .and_then(Value::as_str); + let stripped = strip_event( + event.type_(), + event.state_key().unwrap_or_default(), + &content, + event.sender(), + ); + Ok(stripped_state(membership, &unsigned, stripped)) +} + +/// A string field of an event's content, as the Python's `isinstance(value, +/// str)` checks read it. +fn content_str(event: &Event, key: &str) -> Option { + event + .content() + .get_field(key) + .and_then(Value::as_str) + .map(str::to_owned) +} + +/// `_get_bump_stamp`. +async fn bump_stamp( + context: &RoomContext<'_, '_>, + room_id: &str, + timeline: &[FilteredEvent], + check_outside_timeline: bool, +) -> anyhow::Result> { + let in_timeline = timeline_bump_stamp(timeline.iter().filter_map(|filtered| { + let stream = filtered.event.internal_metadata().stream_ordering()?; + Some((filtered.event.type_(), stream.get())) + })); + if in_timeline.is_some() { + return Ok(in_timeline); + } + if !check_outside_timeline { + return Ok(None); + } + let request = context.request; + let latest = context.reader.latest_bump_stamp(room_id).await?; + match store_bump_stamp( + latest, + request.have_finished_bg_jobs, + request.to_room_key.stream, + ) { + StoreBumpStamp::Done(stamp) => Ok(stamp), + StoreBumpStamp::LookupLastEvent => { + let types: Vec = BUMP_EVENT_TYPES.iter().map(|t| t.to_string()).collect(); + let last = context + .reader + .last_event_of_types_before(room_id, &request.to_room_key, &types) + .await?; + Ok(bump_stamp_from_last_event(last.map(|(_, pos)| pos.stream))) + } + } +} + +/// The redactions of the redacted events of `entry`. +async fn load_redactions( + reader: &SlidingSyncReader<'_>, + entry: &RoomEntry, +) -> anyhow::Result, Event>> { + let mut redaction_ids: Vec> = Vec::new(); + let bundled = entry + .bundles + .iter() + .flat_map(|bundles| bundles.values().flat_map(Bundle::events)); + let events = entry + .timeline + .iter() + .map(|filtered| &filtered.event) + .chain(&entry.required_state) + .chain(bundled); + for event in events { + if let Some(redacted_by) = event.internal_metadata().redacted_by()? { + let redacted_by: Arc = Arc::from(redacted_by); + if !redaction_ids.contains(&redacted_by) { + redaction_ids.push(redacted_by); + } + } + } + Ok(reader + .events(&redaction_ids, RedactBehaviour::Redact, false, false) + .await? + .into_iter() + .map(|event| (Arc::from(event.event_id()), event)) + .collect()) +} + +/// `get_room_sync_data`. +pub async fn room_sync_data( + context: &RoomContext<'_, '_>, + input: &RoomInputData, +) -> anyhow::Result { + let reader = context.reader; + let request = context.request; + let user_id = &*request.user_id; + let member = &input.membership; + let room = Room { + input, + room_id: &member.room_id, + membership: &member.membership, + membership_token: RoomStreamToken::live(member.event_pos.stream), + }; + let room_id = room.room_id; + let membership = room.membership; + let config = &input.room_sync_config; + let prev_config = input.prev_room_sync_config.as_ref(); + let from = request.from.as_ref(); + + // Only a state reset out of the room leaves a membership with no event. + let state_reset_out_of_room = member.event_id.is_none(); + if state_reset_out_of_room && membership != "leave" { + bail!("{room_id}: membership {membership} with no event"); + } + + let TimelineBounds { + from_bound, + initial, + ignore_timeline_bound, + } = timeline_bounds( + &room.input.room_status, + from.map(|from| &from.room_key), + input.newly_joined || state_reset_out_of_room, + prev_config, + config, + ); + + // The timeline. Invites and knocks only have stripped state. + let mut timeline: Vec = Vec::new(); + // The timeline's events as loaded, before visibility. + let mut loaded_timeline: Vec> = Vec::new(); + let mut bundles = None; + let mut limited = None; + let mut prev_batch = None; + let mut live = None; + if config.timeline_limit > 0 && !matches!(membership, "invite" | "knock") { + let to_bound = if room.has_left() { + &room.membership_token + } else { + &request.to_room_key + }; + let timeline_from_bound = if ignore_timeline_bound { + None + } else { + from_bound.as_ref() + }; + // A historical section of the timeline in topological order, as + // `/messages` returns it, or the updates since the bound in the + // order the server received them. + let order = if timeline_from_bound.is_none() { + PaginationOrder::Topological + } else { + PaginationOrder::Stream + }; + let (mut events, new_room_key, page_limited) = reader + .paginate_room_events( + room_id, + to_bound, + timeline_from_bound, + Direction::Backwards, + config.timeline_limit, + order, + ) + .await?; + events.reverse(); + loaded_timeline = events.into_iter().map(Arc::new).collect(); + timeline = filter_events_for_client( + reader, + &context.client_filter, + &loaded_timeline, + &HashSet::new(), + membership != "join", + context.now_ms, + ) + .await?; + + let mut positions = Vec::new(); + if from.is_some() { + for filtered in &timeline { + let metadata = filtered.event.internal_metadata(); + let stream = metadata.stream_ordering().with_context(|| { + format!("event {} has no stream ordering", filtered.event.event_id()) + })?; + let instance_name = metadata.instance_name().with_context(|| { + format!("event {} has no instance name", filtered.event.event_id()) + })?; + positions.push((instance_name, stream.get())); + } + } + live = Some(num_live( + positions + .iter() + .map(|(instance_name, stream)| (instance_name.as_str(), *stream)), + from.map(|from| &from.room_key), + )); + + // With a limited timeline, the client lacks the events it would + // need to aggregate relations itself. + if page_limited { + bundles = Some( + bundled_aggregations( + reader, + &timeline, + user_id, + &context.client_filter.ignored_users, + ) + .await?, + ); + } + limited = Some(page_limited); + + let room_key = new_room_key + .to_string_with_instances(|name| { + let id = context.instance_ids.get(name).copied(); + if id.is_none() { + // A writer added since the table was cached. The next + // request reads it again. + reader.forget_instance_map(); + } + id + }) + .map_err(|e| anyhow::anyhow!("{room_id}: prev_batch: {e}"))?; + prev_batch = Some(format!("{room_key}_{}", request.to_suffix)); + } + + // Python holds no reference to a redacted event, which is loaded again. + let held: HashMap<&str, &Arc> = loaded_timeline + .iter() + .filter(|event| matches!(event.internal_metadata().redacted_by(), Ok(None))) + .map(|event| (event.event_id(), event)) + .collect(); + + // Invites and knocks carry the stripped state the potential joiner + // identifies the room by. + let mut stripped = Vec::new(); + if matches!(membership, "invite" | "knock") { + let event_id = member + .event_id + .as_ref() + .with_context(|| format!("{room_id}: {membership} with no event"))?; + stripped = stripped_state_of(&one_event(reader, event_id).await?)?; + } + + // The changes to the current state in the token range, in the order of + // Python's dict. + let mut deltas: IndexMap> = IndexMap::new(); + let mut membership_changed = false; + let mut name_changed = false; + let mut avatar_changed = false; + if !initial { + let from_bound = from_bound + .as_ref() + .with_context(|| format!("{room_id}: incremental with no from bound"))?; + + // An out-of-band leave, such as a rejected remote invite, is not in + // the state deltas: it is passed in as newly left. + if input.newly_left { + if let Some(event_id) = &member.event_id { + membership_changed = true; + let leave = one_event(reader, event_id).await?; + if let Some(state_key) = leave.state_key() { + deltas.insert( + StateKey::new(leave.type_(), state_key), + Arc::clone(event_id), + ); + } + } + } + + let to_bound = match membership { + "leave" | "ban" => Some(&room.membership_token), + "join" => Some(&request.to_room_key), + "invite" | "knock" => None, + other => bail!("Unexpected membership {other} that we don't know how to handle yet"), + }; + let state_deltas = match to_bound { + Some(to_bound) => { + reader + .state_deltas_by_stream_id(room_id, Some(from_bound), Some(to_bound)) + .await? + } + None => Vec::new(), + }; + for delta in state_deltas { + // A state reset leaves a delta with no event, which is not + // handled yet. + if let Some(event_id) = &delta.event_id { + deltas.insert( + StateKey(Arc::clone(&delta.event_type), Arc::clone(&delta.state_key)), + Arc::clone(event_id), + ); + } + if &*delta.event_type == MEMBER { + membership_changed = true; + } else if &*delta.event_type == M_ROOM_NAME && delta.state_key.is_empty() { + name_changed = true; + } else if &*delta.event_type == M_ROOM_AVATAR && delta.state_key.is_empty() { + avatar_changed = true; + } + } + } + + // The state the request asks for, from the timeline and the deltas + // before any are added below. + let delta_keys: Vec<(String, String)> = deltas + .keys() + .map(|key| (key.0.to_string(), key.1.to_string())) + .collect(); + let timeline_members: Vec> = timeline + .iter() + .map(|filtered| TimelineMember { + sender: filtered.event.sender(), + event_type: filtered.event.type_(), + state_key: filtered.event.state_key(), + }) + .collect(); + let state_request = required_state_request( + user_id, + membership, + &config.required_state_map, + &timeline_members, + limited, + initial, + delta_keys.iter().map(|(t, k)| (t.as_str(), k.as_str())), + ); + drop(timeline_members); + let required_state_filter = &state_request.required_state_filter; + let lazy_load_user_ids = &state_request.lazy_load_user_ids; + + // On the first send, the name and the required state come from one read + // of the room's state, split in memory below. Python reads the name and + // then the required state, each from its cache of the room's state. + let initial_filter = with_room_name(state_filter_with_meta(required_state_filter, &[], true)); + let initial_state_ids = if initial { + Some(state_ids_at(context, &room, &initial_filter).await?) + } else { + None + }; + + let mut room_name = None; + if initial || name_changed { + let name_filter = StateFilter::from_types([(M_ROOM_NAME, Some(""))]); + let name_state = match &initial_state_ids { + Some(state_ids) => { + let ids: Vec> = + name_filter.filter_state(state_ids).into_values().collect(); + state_events(reader, &held, &ids).await? + } + None => state_at(context, &room, &held, &name_filter).await?, + }; + room_name = + lookup(&name_state, (M_ROOM_NAME, "")).and_then(|event| content_str(event, "name")); + } + + // Heroes, when the room has no name, on the first send or when the + // membership changed. The summary also gives the member counts. + let mut summary: Option> = None; + let mut hero_user_ids: Vec> = Vec::new(); + if needs_heroes(room_name.as_deref(), initial, membership_changed) { + // The summary of a room left, invited to or knocked on is not known. + let details = if membership == "join" { + reader.room_summary(room_id).await? + } else { + Arc::new(RoomSummary::new()) + }; + hero_user_ids = extract_heroes_from_room_summary(&details, user_id); + summary = Some(details); + } + + let mut joined_count = None; + let mut invited_count = None; + if (initial || membership_changed) && membership == "join" { + let counts = match &summary { + Some(details) => { + let count = |m: &str| details.get(m).map_or(0, |summary| summary.count); + (count("join"), count("invite")) + } + None => { + let counts = reader.cached_member_counts(room_id).await?; + let count = |m: &str| counts.get(m).copied().unwrap_or(0); + (count("join"), count("invite")) + } + }; + joined_count = Some(counts.0); + invited_count = Some(counts.1); + } + + let state_filter = state_filter_with_meta( + required_state_filter, + &hero_user_ids, + initial || avatar_changed, + ); + + // Every lazy-loaded member counts as sent first in this request, until + // the connection says otherwise. + let mut returned: HashMap> = if state_request.lazy_load_room_members { + lazy_load_user_ids + .iter() + .map(|user_id| (user_id.clone(), None)) + .collect() + } else { + HashMap::new() + }; + let mut invalidated = HashSet::new(); + let mut changed_required_state_map = None; + let mut lazy_member_lookup = false; + + let room_state: HashMap; + // The heroes' memberships, when not in `room_state`. + let mut hero_state: Option> = None; + if initial { + let Some(state_ids) = &initial_state_ids else { + anyhow::bail!("{room_id}: first send without its state read"); + }; + let mut ids = state_filter.filter_state(state_ids); + // The heroes' memberships the first read did not cover. + let missing_heroes: Vec<(&str, Option<&str>)> = hero_user_ids + .iter() + .filter(|user_id| !initial_filter.matches(MEMBER, user_id)) + .map(|user_id| (MEMBER, Some(&**user_id))) + .collect(); + if !missing_heroes.is_empty() { + let heroes = StateFilter::from_types(missing_heroes); + ids.extend(state_ids_at(context, &room, &heroes).await?); + } + let ids: Vec> = ids.into_values().collect(); + room_state = state_events(reader, &held, &ids).await?; + } else { + let from = from.with_context(|| format!("{room_id}: incremental with no from token"))?; + if let Some(prev_config) = prev_config { + // The members whose state may be sent: explicit, lazy-loaded, + // and those whose membership changed. + let mut all_required_user_state: HashSet<&str> = state_request + .explicit_user_state + .iter() + .chain(lazy_load_user_ids) + .map(String::as_str) + .collect(); + for (state_type, state_key) in &delta_keys { + if state_type == MEMBER { + all_required_user_state.insert(state_key); + } + } + + let mut previously_returned: HashSet = HashSet::new(); + if !all_required_user_state.is_empty() { + lazy_member_lookup = true; + let user_ids: Vec> = all_required_user_state + .iter() + .map(|u| Arc::from(*u)) + .collect(); + let last_seen = reader + .connection_lazy_members(from.connection_position, room_id, &user_ids) + .await?; + for (user_id, last_seen_ts) in last_seen { + if lazy_load_user_ids.contains(&*user_id) { + returned.insert(user_id.to_string(), Some(last_seen_ts)); + } + previously_returned.insert(user_id.to_string()); + } + } + + let changes = required_state_changes( + user_id, + &prev_config.required_state_map, + &config.required_state_map, + &previously_returned, + lazy_load_user_ids, + &delta_keys, + ); + changed_required_state_map = changes.changed_required_state_map; + invalidated = changes.lazy_members_invalidated; + // Explicit members converted to lazy ones: when they were last + // sent is not known, so a timestamp already there is kept. + for user_id in changes.extra_users_to_add_to_lazy_cache { + returned.entry(user_id).or_insert(None); + } + if changes.added_state_filter != StateFilter::none() { + // State the request newly asks for, on top of the deltas. + let added = state_ids_at(context, &room, &changes.added_state_filter).await?; + deltas.extend(added); + } + } + + let delta_state: StateMap = deltas + .iter() + .map(|(key, event_id)| (key.clone(), Arc::clone(event_id))) + .collect(); + let event_ids: Vec> = state_filter + .filter_state(&delta_state) + .into_values() + .collect(); + room_state = state_events(reader, &held, &event_ids).await?; + + if !hero_user_ids.is_empty() { + let hero_filter = + StateFilter::from_types(hero_user_ids.iter().map(|u| (MEMBER, Some(&**u)))); + hero_state = Some(state_at(context, &room, &held, &hero_filter).await?); + } + } + + let room_avatar = + lookup(&room_state, (M_ROOM_AVATAR, "")).and_then(|event| content_str(event, "url")); + + let hero_membership_state = hero_state.as_ref().unwrap_or(&room_state); + let heroes: Vec = hero_user_ids + .iter() + .filter_map(|hero| { + let event = lookup(hero_membership_state, (MEMBER, hero))?; + Some(StrippedHero { + user_id: hero.to_string(), + display_name: content_str(event, "displayname"), + avatar_url: content_str(event, "avatar_url"), + }) + }) + .collect(); + + let mut required_state: Vec = Vec::new(); + if *required_state_filter != StateFilter::none() { + let mut keys: Vec<&StateKey> = room_state + .keys() + .filter(|key| required_state_filter.matches(&key.0, &key.1)) + .collect(); + keys.sort(); + required_state = keys + .into_iter() + .map(|key| room_state[key].shallow_clone()) + .collect(); + } + + // The bump stamp: the last bump event for a joined room, else the + // membership's position if the client needs one. + let always_return = always_return_bump_stamp(membership, limited, initial); + let new_bump_stamp = if membership == "join" { + bump_stamp(context, room_id, &timeline, always_return).await? + } else { + None + }; + let bump_stamp = final_bump_stamp(new_bump_stamp, always_return, member.event_pos.stream); + + let room_config = room_config_to_persist( + config, + prev_config, + changed_required_state_map.as_ref(), + ignore_timeline_bound, + limited, + ); + + let mut entry = RoomEntry { + room_id: Arc::clone(&member.room_id), + name: room_name, + avatar: room_avatar, + heroes, + is_dm: input.is_dm, + initial, + unstable_expanded_timeline: ignore_timeline_bound, + required_state, + timeline, + bundles, + redactions: HashMap::new(), + stripped_state: stripped, + prev_batch, + limited, + num_live: live, + bump_stamp, + joined_count, + invited_count, + }; + entry.redactions = load_redactions(reader, &entry).await?; + + Ok(RoomComputed { + entry, + mutations: RoomMutations { + room_config, + lazy_membership: RoomLazyMembershipChanges { + returned_user_id_to_last_seen_ts_map: returned, + invalidated_user_ids: invalidated, + }, + }, + lazy_member_lookup, + }) +} + +/// One room's outcome. A panic is this room's error, not the request's. +async fn room_outcome( + context: &RoomContext<'_, '_>, + input: &RoomInputData, +) -> (Arc, RoomOutcome) { + let outcome = match AssertUnwindSafe(room_sync_data(context, input)) + .catch_unwind() + .await + { + Ok(Ok(computed)) => RoomOutcome::Computed(Box::new(computed)), + Ok(Err(err)) => match err.downcast::() { + Ok(fallback) => RoomOutcome::Fallback(fallback.0), + Err(err) => RoomOutcome::Error(format!("{err:#}")), + }, + Err(panic) => RoomOutcome::Error(format!("panicked: {}", panic_message(&*panic))), + }; + (Arc::clone(&input.membership.room_id), outcome) +} + +/// The entries of `rooms`, twenty at a time, in the order they finish. +pub async fn generate_room_entries( + reader: &SlidingSyncReader<'_>, + request: &SlidingSyncRequestData, + instance_ids: &HashMap, + now_ms: i64, + rooms: &[RoomInputData], +) -> anyhow::Result, RoomOutcome)>> { + let ignored_users = reader.ignored_users(&request.user_id).await?; + let context = RoomContext { + reader, + request, + client_filter: ClientFilterContext { + user_id: Arc::clone(&request.user_id), + ignored_users, + admin_policy: request.admin_policy, + msc4354_enabled: request.msc4354_enabled, + }, + instance_ids, + now_ms, + }; + let outcomes: Vec<_> = rooms + .iter() + .map(|room| room_outcome(&context, room)) + .collect(); + Ok(stream::iter(outcomes) + .buffer_unordered(ROOM_ENTRY_CONCURRENCY) + .collect() + .await) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn empty_entry() -> RoomEntry { + RoomEntry { + room_id: Arc::from("!r:test"), + name: None, + avatar: None, + heroes: Vec::new(), + is_dm: true, + initial: false, + unstable_expanded_timeline: true, + required_state: Vec::new(), + timeline: Vec::new(), + bundles: None, + redactions: HashMap::new(), + stripped_state: Vec::new(), + prev_batch: Some("s1_0".to_owned()), + limited: Some(true), + num_live: Some(0), + bump_stamp: Some(5), + joined_count: None, + invited_count: None, + } + } + + #[test] + fn what_counts_as_an_update() { + // As in Python: a DM flag, an expanded timeline, a token or a bump + // stamp alone are not updates. + assert!(!empty_entry().has_updates()); + // An empty name is, though it is not sent. + let entry = RoomEntry { + name: Some(String::new()), + ..empty_entry() + }; + assert!(entry.has_updates()); + let entry = RoomEntry { + joined_count: Some(0), + ..empty_entry() + }; + assert!(entry.has_updates()); + let entry = RoomEntry { + initial: true, + ..empty_entry() + }; + assert!(entry.has_updates()); + } +} diff --git a/rust/src/sliding_sync/room_lists/filter.rs b/rust/src/sliding_sync/room_lists/filter.rs new file mode 100644 index 0000000000..3f9dd444eb --- /dev/null +++ b/rust/src/sliding_sync/room_lists/filter.rs @@ -0,0 +1,169 @@ +/* + * 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: + * . + */ + +//! Which of a user's rooms a sync covers, and a list's filters. + +use std::collections::{HashMap, HashSet}; + +use pyo3::{exceptions::PyValueError, prelude::*, types::PyModule}; + +use crate::sliding_sync::request::Filters; +use crate::sliding_sync::room_lists::{room_from_py, PyRoomForUser, RoomForUser}; + +/// A filter Python does not implement either: `spaces` and +/// `room_name_like` raise `NotImplementedError` there. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct NotImplementedFilter(pub &'static str); + +impl std::fmt::Display for NotImplementedFilter { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "the {} filter is not implemented", self.0) + } +} + +impl std::error::Error for NotImplementedFilter {} + +/// `filter_membership_for_sync`: whether the sync covers the room. Every +/// membership but a leave does, and a leave only when it is newly left or +/// a kick. A leave with no sender is a state reset. +pub fn filter_membership_for_sync(user_id: &str, room: &RoomForUser, newly_left: bool) -> bool { + &*room.membership != "leave" + || newly_left + || room + .sender + .as_deref() + .is_some_and(|sender| sender != user_id) +} + +/// `filter_rooms_using_tables`: the rooms of `rooms` that pass `filters`, +/// in their order, plus any the user was state reset out of that the +/// connection has sent before (`sent_before` is whether it has, from the +/// 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], + filters: &Filters, + dm_room_ids: &HashSet, + room_tags: &HashMap>, + sent_before: impl Fn(&str) -> bool, +) -> Result, NotImplementedFilter> { + if filters.spaces.is_some() { + return Err(NotImplementedFilter("spaces")); + } + if filters.room_name_like.is_some() { + return Err(NotImplementedFilter("room_name_like")); + } + + let no_tags = HashSet::new(); + let passes = |room: &RoomForUser| -> bool { + if let Some(is_dm) = filters.is_dm { + if dm_room_ids.contains(&*room.room_id) != is_dm { + return false; + } + } + if let Some(is_encrypted) = filters.is_encrypted { + // An unknown encryption status fails either way. + if !room.has_known_state || room.is_encrypted != is_encrypted { + return false; + } + } + if let Some(is_invite) = filters.is_invite { + if (&*room.membership == "invite") != is_invite { + return false; + } + } + if filters.room_types.is_some() || filters.not_room_types.is_some() { + if !room.has_known_state { + return false; + } + let room_type = room.room_type.as_deref(); + if let Some(room_types) = &filters.room_types { + if !room_types.iter().any(|t| t.as_deref() == room_type) { + return false; + } + } + if let Some(not_room_types) = &filters.not_room_types { + if not_room_types.iter().any(|t| t.as_deref() == room_type) { + return false; + } + } + } + if filters.tags.is_some() || filters.not_tags.is_some() { + let tags = room_tags.get(&*room.room_id).unwrap_or(&no_tags); + if let Some(wanted) = &filters.tags { + if !wanted.iter().any(|tag| tags.contains(tag)) { + return false; + } + } + if let Some(unwanted) = &filters.not_tags { + if unwanted.iter().any(|tag| tags.contains(tag)) { + return false; + } + } + } + true + }; + + Ok(rooms + .iter() + .copied() + .filter(|room| { + passes(room) + // State reset out of a room the client knows: it must hear + // of the reset whatever the filters. + || (room.event_id.is_none() && sent_before(&room.room_id)) + }) + .collect()) +} + +/// `filter_membership_for_sync` through the Rust implementation. +#[pyfunction] +fn filter_membership_for_sync_for_test( + user_id: &str, + room: PyRoomForUser, + newly_left: bool, +) -> bool { + filter_membership_for_sync(user_id, &room_from_py(room), newly_left) +} + +/// `filter_rooms_using_tables` through the Rust implementation, with +/// `filters` as JSON. Returns the room IDs that pass, or raises +/// `NotImplementedError`. +#[pyfunction] +fn filter_rooms_for_test( + rooms: Vec, + filters: &str, + dm_room_ids: HashSet, + room_tags: HashMap>, + sent_before: HashSet, +) -> PyResult> { + 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) + }) + .map_err(|e| pyo3::exceptions::PyNotImplementedError::new_err(e.to_string()))?; + Ok(filtered + .iter() + .map(|room| room.room_id.to_string()) + .collect()) +} + +pub fn register_test_functions(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_function(wrap_pyfunction!(filter_membership_for_sync_for_test, m)?)?; + m.add_function(wrap_pyfunction!(filter_rooms_for_test, m)?)?; + Ok(()) +} 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..62295230eb --- /dev/null +++ b/rust/src/sliding_sync/room_lists/interested.rs @@ -0,0 +1,913 @@ +/* + * 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::*, types::PyDict}; +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. +/// +/// `get_max_pos_of_last_change` is a lookup in the cache's `_entity_to_key` +/// dict, so this looks each room up in that dict directly. A method call per +/// joined room was most of the GIL time of an incremental request on a +/// large account. +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 entity_to_key = cache.getattr(pyo3::intern!(py, "_entity_to_key"))?; + let entity_to_key = entity_to_key.cast::()?; + let mut positions = HashMap::with_capacity(room_ids.len()); + for room_id in room_ids { + if let Some(position) = entity_to_key.get_item(&**room_id)? { + positions.insert(Arc::clone(room_id), position.extract()?); + } + } + 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 new file mode 100644 index 0000000000..691c7c60cb --- /dev/null +++ b/rust/src/sliding_sync/room_lists/lists.rs @@ -0,0 +1,225 @@ +/* + * 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 list's ranges to its `SYNC` operations, the room configs the lists +//! and room subscriptions combine into, and which of those rooms may have +//! updates to send (`_filter_relevant_rooms_to_send`). + +use std::collections::HashSet; +use std::sync::Arc; + +use indexmap::IndexMap; + +use crate::sliding_sync::config::RoomSyncConfig; +use crate::sliding_sync::connection::{HaveSentRoom, PerConnectionState}; + +/// `MINIMUM_NOT_USED_AGE_EXPIRY`, in milliseconds: a connection with too +/// many rooms to update is only reset after this long unused. +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 +/// the lists and room subscriptions add them. +pub type RelevantRoomMap = IndexMap, RoomSyncConfig>; + +/// `SlidingWindowList.Operation`: always a `SYNC`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SyncOp { + pub range: (u64, u64), + pub room_ids: Vec>, +} + +/// Whether a list with these `ranges` and `count` rooms needs sorting: +/// not when one range covers every room from the start, and not without +/// ranges. +pub fn needs_sort(ranges: &[(u64, u64)], count: usize) -> bool { + match ranges { + [] => false, + [(0, end)] => (*end as u128) < (count as u128).saturating_sub(1), + _ => true, + } +} + +/// The number of top rooms the ranges need sorted: the end of the last +/// one, inclusive. +pub fn sort_limit(ranges: &[(u64, u64)]) -> usize { + ranges + .iter() + .map(|(_, end)| { + usize::try_from(*end) + .unwrap_or(usize::MAX) + .saturating_add(1) + }) + .max() + .unwrap_or(0) +} + +/// `_sync_list_ops`: the `SYNC` operations of a list's `ranges` over its +/// sorted rooms, adding each room in a range to `relevant_room_map` with +/// the list's config combined with the room's existing one. +/// +/// As in Python, once a room combines with an existing config, the rooms +/// after it in the list take the combined config too. +pub fn sync_list_ops( + ranges: &[(u64, u64)], + sorted_room_ids: &[Arc], + room_sync_config: &RoomSyncConfig, + relevant_room_map: &mut RelevantRoomMap, +) -> Vec { + let mut room_sync_config = room_sync_config.clone(); + let mut ops = Vec::with_capacity(ranges.len()); + for &(start, end) in ranges { + // An end before the start takes no rooms, as in Python. + let max_num_rooms = if end < start { + 0 + } else { + (end - start).saturating_add(1) + }; + let skip = usize::try_from(start).unwrap_or(usize::MAX); + let mut room_ids = Vec::new(); + for room_id in sorted_room_ids.iter().skip(skip) { + if room_ids.len() as u64 >= max_num_rooms { + break; + } + if let Some(existing) = relevant_room_map.get(room_id) { + room_sync_config = existing.combine(&room_sync_config); + } + relevant_room_map.insert(Arc::clone(room_id), room_sync_config.clone()); + room_ids.push(Arc::clone(room_id)); + } + ops.push(SyncOp { + range: (start, end), + room_ids, + }); + } + ops +} + +/// Add a room subscription's room to `relevant_room_map`, combined with +/// the room's existing config. +pub fn add_room_subscription( + relevant_room_map: &mut RelevantRoomMap, + room_id: &Arc, + room_sync_config: RoomSyncConfig, +) { + let config = match relevant_room_map.get(room_id) { + Some(existing) => existing.combine(&room_sync_config), + None => room_sync_config, + }; + relevant_room_map.insert(Arc::clone(room_id), config); +} + +/// `SlidingSyncUnknownPosition`: the connection is reset, as it has too +/// many rooms to update or its position is not one the store knows. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct UnknownPosition; + +impl std::fmt::Display for UnknownPosition { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("unknown sliding sync connection position") + } +} + +impl std::error::Error for UnknownPosition {} + +/// The rooms of an incremental sync that may have updates, the decision +/// half of `_filter_relevant_rooms_to_send`: rooms whose config expanded, +/// rooms never sent or with updates left over, and +/// `rooms_that_have_updates`, the rooms the events stream has new events +/// in since the `from_token`. Their order is `relevant_room_map`'s. +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(); + for (room_id, room_config) in relevant_room_map { + if let Some(prev) = previous_state.room_configs.get(&**room_id) { + // Timeline limit increases and required state additions. + if prev.combine(room_config) != *prev { + rooms_should_send.insert(room_id); + continue; + } + } + match previous_state.rooms.have_sent_room(room_id) { + HaveSentRoom::Never | HaveSentRoom::Previously(_) => { + rooms_should_send.insert(room_id); + } + // Everything up to the `from_token` was sent. + HaveSentRoom::Live => {} + } + } + + // 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 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); + } + } + } + + Ok(relevant_room_map + .keys() + .filter(|room_id| { + rooms_should_send.contains(room_id) || rooms_that_have_updates.contains(*room_id) + }) + .cloned() + .collect()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn config(limit: i64) -> RoomSyncConfig { + RoomSyncConfig { + timeline_limit: limit, + ..Default::default() + } + } + + #[test] + fn full_range_needs_no_sort() { + assert!(!needs_sort(&[], 5)); + assert!(!needs_sort(&[(0, 4)], 5)); + assert!(!needs_sort(&[(0, 0)], 0)); + assert!(!needs_sort(&[(0, 10)], 5)); + assert!(needs_sort(&[(0, 3)], 5)); + assert!(needs_sort(&[(1, 10)], 5)); + assert!(needs_sort(&[(0, 10), (0, 10)], 5)); + assert_eq!(sort_limit(&[(0, 3), (5, 9)]), 10); + } + + #[test] + fn combined_config_carries_on_down_the_list() { + let rooms: Vec> = ["!a", "!b", "!c"].into_iter().map(Arc::from).collect(); + let mut relevant = RelevantRoomMap::new(); + relevant.insert(Arc::from("!b"), config(10)); + let ops = sync_list_ops(&[(0, 1), (2, 5)], &rooms, &config(1), &mut relevant); + assert_eq!(ops[0].room_ids, rooms[..2].to_vec()); + assert_eq!(ops[1].room_ids, rooms[2..].to_vec()); + assert_eq!(relevant["!a"].timeline_limit, 1); + assert_eq!(relevant["!b"].timeline_limit, 10); + // Python's leak: `!c` takes `!b`'s combined config. + assert_eq!(relevant["!c"].timeline_limit, 10); + } +} diff --git a/rust/src/sliding_sync/room_lists/mod.rs b/rust/src/sliding_sync/room_lists/mod.rs new file mode 100644 index 0000000000..3948438325 --- /dev/null +++ b/rust/src/sliding_sync/room_lists/mod.rs @@ -0,0 +1,231 @@ +/* + * 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: which rooms a request is interested in, filtered, +//! sorted and cut to the requested ranges +//! (`synapse.handlers.sliding_sync.room_lists`). + +use std::collections::{BTreeMap, HashMap, HashSet}; +use std::sync::Arc; + +use pyo3::{ + exceptions::{PyKeyError, PyValueError}, + prelude::*, + types::PyModule, +}; + +use crate::sliding_sync::config::{config_from_py, to_test_tuple, RoomSyncConfig}; +use crate::sliding_sync::connection::{HaveSentRoom, PerConnectionState, RoomStatusMap}; +use crate::storage::sync_store::PersistedEventPosition; +use crate::types::stream_token::RoomStreamToken; + +pub mod filter; +pub mod interested; +pub mod lists; +pub mod sort; + +/// A user's membership in a room at the `to_token`, as the sliding sync +/// tables record it: `RoomsForUserSlidingSync`, which also stands for +/// `RoomsForUser` and `RoomsForUserStateReset` where only their fields are +/// read. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RoomForUser { + pub room_id: Arc, + /// `None` after a state reset. + pub sender: Option>, + pub membership: Arc, + /// `None` after a state reset. + pub event_id: Option>, + pub event_pos: PersistedEventPosition, + /// `None` when the `rooms` row has no version, which only the reads + /// that keep unknown versions return. + pub room_version_id: Option>, + pub has_known_state: bool, + pub room_type: Option>, + pub is_encrypted: bool, +} + +/// A [`RoomForUser`] from the test form `(room_id, sender, membership, +/// event_id, stream, has_known_state, room_type, is_encrypted)`. +pub(crate) type PyRoomForUser = ( + String, + Option, + String, + Option, + i64, + bool, + Option, + bool, +); + +pub(crate) fn room_from_py(room: PyRoomForUser) -> RoomForUser { + let (room_id, sender, membership, event_id, stream, has_known_state, room_type, is_encrypted) = + room; + 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("master"), + stream, + }, + room_version_id: Some(Arc::from("10")), + has_known_state, + room_type: room_type.map(Arc::from), + is_encrypted, + } +} + +type PyConfig = (i64, BTreeMap>); + +/// `sort_rooms` through the Rust implementation, with the stream change +/// cache as `earliest_cache_position` and `cache` (room to its max position +/// of last change), and the database as `recheck` (room to the stream of +/// its last event before the `to_token`) and `joined_positions`. Returns +/// the sorted room IDs. Raises `KeyError` as Python does. +#[pyfunction] +#[pyo3(signature = (rooms, earliest_cache_position, cache, to_stream, limit, recheck, joined_positions))] +fn sort_rooms_for_test( + rooms: Vec, + earliest_cache_position: i64, + cache: HashMap, + to_stream: i64, + limit: Option, + recheck: HashMap>, + joined_positions: HashMap, +) -> PyResult> { + let rooms: Vec = rooms.into_iter().map(room_from_py).collect(); + let rooms: Vec<&RoomForUser> = rooms.iter().collect(); + let mut plan = sort::SortPlan::new(&rooms, earliest_cache_position, |room_id| { + cache.get(room_id).copied() + }); + for room_id in plan.rooms_to_recheck(limit, to_stream) { + let result = recheck + .get(&*room_id) + .ok_or_else(|| PyValueError::new_err(format!("no recheck for {room_id}")))?; + plan.apply_recheck(&room_id, *result); + } + let sorted = match plan.early_result(&rooms, limit) { + Some(sorted) => sorted, + None => { + let joined: HashMap, i64> = sort::SortPlan::joined_rooms(&rooms) + .into_iter() + .filter_map(|room_id| { + joined_positions + .get(&*room_id) + .map(|position| (room_id, *position)) + }) + .collect(); + plan.sort_all(&rooms, &joined) + .map_err(|e| PyKeyError::new_err(e.to_string()))? + } + }; + Ok(sorted.iter().map(|room| room.room_id.to_string()).collect()) +} + +/// `_sync_list_ops` through the Rust implementation, the relevant room map +/// as a list of `(room_id, config)`. Returns `(ops, relevant_room_map)`, +/// each op as `(range, room_ids)`. +#[pyfunction] +#[allow(clippy::type_complexity)] +fn sync_list_ops_for_test( + ranges: Vec<(u64, u64)>, + sorted_room_ids: Vec, + room_sync_config: PyConfig, + relevant_room_map: Vec<(String, PyConfig)>, +) -> (Vec<((u64, u64), Vec)>, Vec<(String, PyConfig)>) { + let mut relevant: lists::RelevantRoomMap = relevant_room_map + .into_iter() + .map(|(room_id, (limit, map))| (Arc::from(room_id), config_from_py(limit, map))) + .collect(); + let sorted: Vec> = sorted_room_ids.into_iter().map(Arc::from).collect(); + let ops = lists::sync_list_ops( + &ranges, + &sorted, + &config_from_py(room_sync_config.0, room_sync_config.1), + &mut relevant, + ); + ( + ops.into_iter() + .map(|op| { + ( + op.range, + op.room_ids.iter().map(|r| r.to_string()).collect(), + ) + }) + .collect(), + relevant + .iter() + .map(|(room_id, config)| (room_id.to_string(), to_test_tuple(config))) + .collect(), + ) +} + +/// The decision of `_filter_relevant_rooms_to_send` through the Rust +/// implementation, each room's status as `(flag, token)` with a token an +/// int. Returns the room IDs to send, or raises `ValueError` for an +/// unknown position. +#[pyfunction] +fn rooms_to_send_for_test( + relevant_room_map: Vec<(String, PyConfig)>, + room_configs: HashMap, + statuses: HashMap)>, + last_used_ts: Option, + rooms_that_have_updates: HashSet, + now_ms: i64, +) -> PyResult> { + let relevant: lists::RelevantRoomMap = relevant_room_map + .into_iter() + .map(|(room_id, (limit, map))| (Arc::from(room_id), config_from_py(limit, map))) + .collect(); + let mut rooms = HashMap::new(); + for (room_id, (flag, token)) in statuses { + let status = match (flag.as_str(), token) { + ("never", _) => HaveSentRoom::Never, + ("live", _) => HaveSentRoom::Live, + (_, Some(token)) => HaveSentRoom::Previously(RoomStreamToken::live(token)), + _ => return Err(PyValueError::new_err("bad status")), + }; + rooms.insert(room_id, status); + } + let room_configs: HashMap = room_configs + .into_iter() + .map(|(room_id, (limit, map))| (room_id, config_from_py(limit, map))) + .collect(); + let state: PerConnectionState = PerConnectionState { + last_used_ts, + rooms: RoomStatusMap::new(rooms), + room_configs: Arc::new(room_configs), + ..Default::default() + }; + let updates: HashSet> = rooms_that_have_updates.into_iter().map(Arc::from).collect(); + 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<()> { + filter::register_test_functions(m)?; + m.add_function(wrap_pyfunction!(sort_rooms_for_test, m)?)?; + m.add_function(wrap_pyfunction!(sync_list_ops_for_test, m)?)?; + m.add_function(wrap_pyfunction!(rooms_to_send_for_test, m)?)?; + Ok(()) +} diff --git a/rust/src/sliding_sync/room_lists/sort.rs b/rust/src/sliding_sync/room_lists/sort.rs new file mode 100644 index 0000000000..7af21bcac3 --- /dev/null +++ b/rust/src/sliding_sync/room_lists/sort.rs @@ -0,0 +1,169 @@ +/* + * 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: + * . + */ + +//! `sort_rooms`: a list's rooms by the stream ordering of the last event +//! the user can see in each, latest first. +//! +//! Python's method reads the events stream change cache and the database +//! as it goes. Here [`SortPlan`] is the decision logic, and the caller +//! fetches what it asks for between the steps: +//! +//! 1. [`SortPlan::new`] reads the cached positions. +//! 2. With a `limit`, [`SortPlan::rooms_to_recheck`] are cached positions +//! past the `to_token`, whose last event before it the caller looks up +//! and passes to [`SortPlan::apply_recheck`]. +//! 3. [`SortPlan::early_result`] is the answer if the cache covered the +//! top `limit` rooms. +//! 4. Otherwise the caller fetches the last event positions of +//! [`SortPlan::joined_rooms`] and [`SortPlan::sort_all`] sorts. + +use std::collections::HashMap; +use std::sync::Arc; + +use anyhow::anyhow; + +use crate::sliding_sync::room_lists::RoomForUser; + +/// The positions [`sort_rooms`](self) has decided on so far. +#[derive(Debug, Clone)] +pub struct SortPlan { + /// The earliest position the events stream change cache knows. + earliest_cache_position: i64, + /// The last activity of each room the user is not joined to: their + /// membership event, as they see nothing after it. + last_activity: HashMap, i64>, + /// The positions known to be in the stream change cache's range. + cached_positions: HashMap, i64>, +} + +impl SortPlan { + /// Read the positions the cache gives. `max_pos_of_last_change` is + /// the cache's `get_max_pos_of_last_change`. + pub fn new( + rooms: &[&RoomForUser], + earliest_cache_position: i64, + max_pos_of_last_change: impl Fn(&str) -> Option, + ) -> Self { + let mut last_activity = HashMap::new(); + let mut cached_positions = HashMap::new(); + for room in rooms { + if &*room.membership == "join" { + if let Some(position) = max_pos_of_last_change(&room.room_id) { + cached_positions.insert(Arc::clone(&room.room_id), position); + } + } else { + // FIXME, as in Python: invited or knocking users might see + // further, depending on the history visibility. + let stream = room.event_pos.stream; + last_activity.insert(Arc::clone(&room.room_id), stream); + if stream > earliest_cache_position { + cached_positions.insert(Arc::clone(&room.room_id), stream); + } + } + } + SortPlan { + earliest_cache_position, + last_activity, + cached_positions, + } + } + + /// Whether the cached positions cover the top `limit` rooms. + fn enough_cached(&self, limit: Option) -> bool { + limit.is_some_and(|limit| self.cached_positions.len() >= limit) + } + + /// The rooms whose cached position is past `to_stream`, which need + /// their last event before the `to_token` looked up. Only with a + /// `limit` the cache could answer. + pub fn rooms_to_recheck(&self, limit: Option, to_stream: i64) -> Vec> { + if !self.enough_cached(limit) { + return Vec::new(); + } + self.cached_positions + .iter() + .filter(|(_, position)| **position > to_stream) + .map(|(room_id, _)| Arc::clone(room_id)) + .collect() + } + + /// Record a recheck: `last_event_stream` is the stream ordering of the + /// room's last event before the `to_token`, if any. A position outside + /// the cache's range drops the room from the cached positions. + pub fn apply_recheck(&mut self, room_id: &str, last_event_stream: Option) { + match last_event_stream { + Some(stream) if stream > self.earliest_cache_position => { + self.cached_positions.insert(Arc::from(room_id), stream); + } + _ => { + self.cached_positions.remove(room_id); + } + } + } + + /// The top rooms from the cache alone, if it covers `limit` of them: + /// the rooms with a cached position, latest first, ties in `rooms` + /// order. + pub fn early_result<'a>( + &self, + rooms: &[&'a RoomForUser], + limit: Option, + ) -> Option> { + if !self.enough_cached(limit) { + return None; + } + let mut sorted: Vec<&RoomForUser> = rooms + .iter() + .filter(|room| self.cached_positions.contains_key(&room.room_id)) + .copied() + .collect(); + sorted.sort_by_key(|room| std::cmp::Reverse(self.cached_positions[&room.room_id])); + Some(sorted) + } + + /// The joined rooms whose last event position the full sort needs. + pub fn joined_rooms(rooms: &[&RoomForUser]) -> Vec> { + rooms + .iter() + .filter(|room| &*room.membership == "join") + .map(|room| Arc::clone(&room.room_id)) + .collect() + } + + /// Sort every room, given `joined_positions`, the last event position + /// before the `to_token` of each joined room. Fails for a joined room + /// without one, where Python raises `KeyError`. + pub fn sort_all<'a>( + mut self, + rooms: &[&'a RoomForUser], + joined_positions: &HashMap, i64>, + ) -> anyhow::Result> { + self.last_activity.extend( + joined_positions + .iter() + .map(|(room_id, position)| (Arc::clone(room_id), *position)), + ); + let mut keyed = Vec::with_capacity(rooms.len()); + for room in rooms { + let position = self + .last_activity + .get(&room.room_id) + .copied() + .ok_or_else(|| anyhow!("no last event position for {}", room.room_id))?; + keyed.push((position, *room)); + } + keyed.sort_by_key(|(position, _)| std::cmp::Reverse(*position)); + Ok(keyed.into_iter().map(|(_, room)| room).collect()) + } +} diff --git a/rust/src/sliding_sync/state.rs b/rust/src/sliding_sync/state.rs new file mode 100644 index 0000000000..6a2fe42b6f --- /dev/null +++ b/rust/src/sliding_sync/state.rs @@ -0,0 +1,813 @@ +/* + * 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::*, + sync::PyOnceLock, + types::{PyDict, PyTuple}, +}; + +use crate::caches::lock_ignoring_poison; +use crate::sliding_sync::config::RoomSyncConfig; +use crate::sliding_sync::connection::{ + HaveSentRoom, MutablePerConnectionState, MutableRoomStatusMap, 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() + } +} + +impl StoredConnectionState { + /// Whether a stored room token names a writer `instances` lacks. + pub fn names_unknown_writer(&self, instances: &InstanceMap) -> bool { + self.rooms.values().any(|status| match status { + HaveSentRoom::Previously(token) => { + RoomStreamToken::parse_with_instances(token, |id| instances.name(id)).is_err() + && RoomStreamToken::parse_with_instances(token, |_| Some(Arc::from(""))).is_ok() + } + HaveSentRoom::Never | HaveSentRoom::Live => false, + }) + } +} + +/// 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, + } + } +} + +/// Python's `HaveSentRoom` class, with one `LIVE` and one `NEVER` value. The +/// class is frozen, so those two serve every room. The extensions ask for +/// each room in their lists, so building them per call showed in requests +/// with hundreds of rooms. +struct PyHaveSentRoom { + class: Py, + live: Py, + never: Py, +} + +static PY_HAVE_SENT_ROOM: PyOnceLock = PyOnceLock::new(); + +/// 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 cached = PY_HAVE_SENT_ROOM.get_or_try_init(py, || -> PyResult<_> { + let class = py + .import(intern!(py, "synapse.types.handlers.sliding_sync"))? + .getattr(intern!(py, "HaveSentRoom"))?; + Ok(PyHaveSentRoom { + live: class.call_method0(intern!(py, "live"))?.unbind(), + never: class.call_method0(intern!(py, "never"))?.unbind(), + class: class.unbind(), + }) + })?; + match status { + HaveSentRoom::Never => Ok(cached.never.bind(py).clone()), + HaveSentRoom::Live => Ok(cached.live.bind(py).clone()), + HaveSentRoom::Previously(token) => cached + .class + .bind(py) + .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>, + )> { + // Copy out under the lock, then build the Python objects without it. + let (last_used_ts, configs, sent) = { + let data = lock_ignoring_poison(&self.data); + let previous = &data.previous; + let mut configs = HashMap::new(); + let mut sent = Vec::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() { + sent.push((room_id, status)); + } + } + (previous.last_used_ts, configs, sent) + }; + let statuses = statuses_to_py(py, sent, |token| token.into_pyobject(py))? + .into_iter() + .collect(); + Ok((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. The statuses are copied out first, so the lock is + /// not held while calling into Python. + fn statuses<'py>( + &self, + py: Python<'py>, + room_ids: Option<&[String]>, + ) -> PyResult)>> { + let state = self.state.bind(py).get(); + let instances = &state.instances; + macro_rules! copy { + ($stream:ident) => {{ + let data = lock_ignoring_poison(&state.data); + copy_statuses( + &data.previous.$stream, + &data.new.$stream, + self.mutable, + room_ids, + ) + }}; + } + match self.stream { + Stream::Rooms => statuses_to_py(py, copy!(rooms), |token: &RoomStreamToken| { + token.into_pyobject(py) + }), + Stream::Receipts => statuses_to_py(py, copy!(receipts), |token: &String| { + multi_writer_to_py(py, token, instances) + }), + Stream::AccountData => statuses_to_py(py, copy!(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", + )) + } + } +} + +/// The statuses of `room_ids`, or every room's when `None`, from +/// `previous`, or from `new` when `mutable`. +fn copy_statuses( + previous: &RoomStatusMap, + new: &MutableRoomStatusMap, + mutable: bool, + room_ids: Option<&[String]>, +) -> Vec<(String, HaveSentRoom)> { + match room_ids { + Some(room_ids) => room_ids + .iter() + .map(|room_id| { + let status = if mutable { + new.have_sent_room(room_id) + } else { + previous.have_sent_room(room_id) + }; + (room_id.clone(), status) + }) + .collect(), + None if mutable => new + .copy() + .iter() + .map(|(room_id, status)| (room_id.to_owned(), status.clone())) + .collect(), + None => previous + .iter() + .map(|(room_id, status)| (room_id.to_owned(), status.clone())) + .collect(), + } +} + +fn statuses_to_py<'py, T>( + py: Python<'py>, + statuses: Vec<(String, HaveSentRoom)>, + token_to_py: impl Fn(&T) -> PyResult>, +) -> PyResult)>> { + statuses + .into_iter() + .map(|(room_id, status)| Ok((room_id, have_sent_room_to_py(py, &status, &token_to_py)?))) + .collect() +} + +/// 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 unknown_writers_in_stored_tokens() { + let stored = |token: &str| StoredConnectionState { + rooms: HashMap::from([( + "!a:x".to_owned(), + HaveSentRoom::Previously(token.to_owned()), + )]), + ..Default::default() + }; + let known = InstanceMap::new([(3, Arc::from("w3"))]); + let empty = InstanceMap::default(); + + assert!(stored("m5~3.6").names_unknown_writer(&empty)); + assert!(!stored("m5~3.6").names_unknown_writer(&known)); + assert!(!stored("s5").names_unknown_writer(&empty)); + // A token that is broken otherwise is not a missing writer. + assert!(!stored("x5").names_unknown_writer(&empty)); + assert!(!StoredConnectionState::default().names_unknown_writer(&empty)); + } + + #[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/state/v2.rs b/rust/src/state/v2.rs index f6f7db1c49..10a062d696 100644 --- a/rust/src/state/v2.rs +++ b/rust/src/state/v2.rs @@ -686,7 +686,7 @@ fn push_py_json_list(out: &mut String, items: &[&str]) { } /// Encode a string as Python's `json.dumps` does with `ensure_ascii=True`. -fn push_py_json_str(out: &mut String, s: &str) { +pub(crate) fn push_py_json_str(out: &mut String, s: &str) { use std::fmt::Write; out.push('"'); diff --git a/rust/src/storage/caches.rs b/rust/src/storage/caches.rs index ffe50f2fe7..6dd5f0573b 100644 --- a/rust/src/storage/caches.rs +++ b/rust/src/storage/caches.rs @@ -49,13 +49,19 @@ //! `synapse.util.caches.mirror_rust_cache`. They are filled by the sync //! store's readers. -use std::{collections::HashSet, sync::Arc}; +use std::{ + collections::{HashMap, HashSet}, + sync::Arc, +}; use crate::caches::{python::CacheHandle, registry::CacheRegistry, Cache}; use crate::federation::client::limiter::RetryTimings; use crate::state::StateMap; use crate::storage::events::CachedEvent; -use crate::storage::sync_store::{RetentionRow, RoomNotifCounts, RoomSummary}; +use crate::storage::sliding_sync_store::{ + ConnectionPositionKey, RoomType, SnapshotRooms, StoredConnectionState, UserTags, +}; +use crate::storage::sync_store::{RelatedEvent, RetentionRow, RoomNotifCounts, RoomSummary}; pub type EventCache = Cache, CachedEvent>; @@ -145,11 +151,70 @@ pub struct StoreCaches { pub unread_counts: Arc, /// `get_destination_retry_timings`. pub destination_retry_timings: Arc, Option>>, + /// `get_member_counts`. + pub member_counts: Arc, + /// `get_global_account_data_by_type_for_user`, keyed on user and type, + /// and invalidated by user alone as well, as a `tree=True` cache. + pub global_account_data: Arc, + /// `ignored_users`. + pub ignored_users: Arc, + /// `get_tags_for_user`. + pub tags_for_user: Arc, Arc>>, + /// `get_room_type`. + pub room_types: Arc, RoomType>>, + /// `get_room_version_id`. A room without one is not cached, as Python + /// raises for it. + pub room_version_ids: Arc, Arc>>, + /// `get_partial_rooms`, under its one key, weighed by rooms. + pub partial_rooms: Arc>>>>, + + // Sliding sync. + /// `get_sliding_sync_rooms_for_user_from_membership_snapshots`, weighed + /// by rooms. + pub sliding_sync_rooms_for_user: Arc, Arc>>, + /// `_get_max_event_pos`. `None` for a room the batch load did not find, + /// as Python's `cachedList` stores. + pub max_event_pos: Arc, Option>>, + + // Relations. + /// The IDs and count of `get_thread_summary`: the latest threaded reply + /// and the number of replies, `None` for an event with no thread. + pub thread_summaries: Arc, + /// `get_references_for_event`, as `(event_id, sender)` oldest first. + pub references: Arc, Arc<[RelatedEvent]>>>, + /// The ID of `get_applicable_edit`, `None` for an event with no edit. + pub applicable_edits: Arc, Option>>>, + /// `get_and_clear_connection_positions`: a memo of each connection + /// position's state, so the transaction that loads it runs once per + /// position. Python never invalidates it, and a position's stored state + /// never changes once it is loaded, so this mirrors nothing. Weighed by + /// rooms, as Python's `iterable=True` cache is. + pub connection_states: Arc>>, + /// The `instance_map` table, under its one key. Rows are only ever + /// added, so this mirrors nothing: a reader that misses a writer drops + /// it and reads again. + pub instance_map: Arc, } +/// Room ID to its number of members by membership. +pub type MemberCountsCache = Cache, Arc>>; + +/// `(user_id, account_data_type)` to the user's global account data of +/// that type. +pub type GlobalAccountDataCache = Cache<(Arc, Arc), Option>>; + /// User ID to the rooms they are joined to. pub type RoomsForUserCache = Cache, Arc>>>; +/// `ignored_users`, by the ignoring user. +pub type IgnoredUsersCache = Cache, Arc>>>; + +/// The latest threaded reply and the reply count, by thread root. +pub type ThreadSummaryCache = Cache, Option<(Arc, i64)>>; + +/// The `instance_map` table's `(instance_id, instance_name)` rows. +pub type InstanceMapCache = Cache<(), Arc<[(i64, Arc)]>>; + /// `(room_id, user_id)` to the user's unread counts in the room. pub type UnreadCountsCache = Cache<(Arc, Arc), Arc>; @@ -194,6 +259,30 @@ impl StoreCaches { 5000, |counts| counts.threads.len() + 1, ), + member_counts: registry.new_cache("rust_get_member_counts", 1000), + global_account_data: registry + .new_cache("rust_get_global_account_data_by_type_for_user", 5000), + ignored_users: registry.new_cache("rust_ignored_users", 5000), + tags_for_user: registry.new_cache("rust_get_tags_for_user", 1000), + room_types: registry.new_cache("rust_get_room_type", 10_000), + room_version_ids: registry.new_cache("rust_get_room_version_id", 10_000), + partial_rooms: registry + .new_weighted_cache("rust_get_partial_rooms", 10_000, |rooms| rooms.len()), + sliding_sync_rooms_for_user: registry.new_weighted_cache( + "rust_get_sliding_sync_rooms_for_user_from_membership_snapshots", + 10_000, + |rooms| rooms.len(), + ), + max_event_pos: registry.new_cache("rust_get_max_event_pos", 1000), + thread_summaries: registry.new_cache("rust_get_thread_summary", 10_000), + references: registry.new_cache("rust_get_references_for_event", 10_000), + applicable_edits: registry.new_cache("rust_get_applicable_edit", 10_000), + instance_map: registry.new_cache("rust_instance_map", 1), + connection_states: registry.new_weighted_cache( + "rust_get_and_clear_connection_positions", + 100_000, + |state| state.len(), + ), } } @@ -227,6 +316,24 @@ impl StoreCaches { "get_destination_retry_timings", &self.destination_retry_timings, ), + mirror("get_member_counts", &self.member_counts), + mirror( + "get_global_account_data_by_type_for_user", + &self.global_account_data, + ), + mirror("ignored_users", &self.ignored_users), + mirror("get_tags_for_user", &self.tags_for_user), + mirror("get_room_type", &self.room_types), + mirror("get_room_version_id", &self.room_version_ids), + mirror("get_partial_rooms", &self.partial_rooms), + mirror( + "get_sliding_sync_rooms_for_user_from_membership_snapshots", + &self.sliding_sync_rooms_for_user, + ), + mirror("_get_max_event_pos", &self.max_event_pos), + mirror("get_thread_summary", &self.thread_summaries), + mirror("get_references_for_event", &self.references), + mirror("get_applicable_edit", &self.applicable_edits), ] } diff --git a/rust/src/storage/mod.rs b/rust/src/storage/mod.rs index b14ab715ef..47216408cc 100644 --- a/rust/src/storage/mod.rs +++ b/rust/src/storage/mod.rs @@ -20,6 +20,7 @@ pub mod federation_store; pub mod native_pool; pub mod pool_querier; pub mod python; +pub mod sliding_sync_store; pub mod store; pub mod stream_ids; pub mod sync_store; diff --git a/rust/src/storage/sliding_sync_store.rs b/rust/src/storage/sliding_sync_store.rs new file mode 100644 index 0000000000..77a42705f6 --- /dev/null +++ b/rust/src/storage/sliding_sync_store.rs @@ -0,0 +1,2292 @@ +/* + * 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: + * . + */ + +//! What the sliding sync handler reads from the database, and the two +//! transactions that load and store a connection's state. +//! +//! [`SlidingSyncReader`] is a [`SyncReader`] with the sliding sync queries +//! on top, so one reader serves a request's v2 reads (timelines, state, +//! events) and its sliding sync ones, counting both in one set of stats. +//! Each query is a port of the store method it is named after, and +//! `tests/storage/test_sliding_sync_store.py` runs both against the same +//! database. +//! +//! Only the path Python takes once the sliding sync background updates +//! have finished is ported: the request falls back to Python as a whole +//! before then. +//! +//! The reads that consult Python's in-memory stream change caches take +//! what they need from them as arguments, read once per request: whether +//! the user's membership may have changed, and which rooms have new +//! events. The caches that mirror a Python `@cached` function are in +//! [`StoreCaches`](super::caches::StoreCaches). + +use std::{ + collections::{BTreeMap, HashMap, HashSet}, + ops::Deref, + sync::Arc, +}; + +use anyhow::Context; +use futures::{future::BoxFuture, FutureExt}; +use indexmap::IndexMap; +use log::warn; +use pyo3::{prelude::*, types::PyTuple}; +use tokio::runtime::Handle; +use tokio_postgres::{error::SqlState, Row}; + +use crate::caches::python::PyCacheKey; +use crate::database::postgres::pool::{ConnectionPool, PyConnectionPool}; +use crate::room_versions::is_known_room_version; +use crate::sliding_sync::config::{RequiredStateMap, RoomSyncConfig}; +use crate::sliding_sync::connection::{HaveSentRoom, RoomLazyMembershipChanges}; +use crate::sliding_sync::room_lists::{lists::UnknownPosition, RoomForUser}; +use crate::sliding_sync::SlidingSyncHandler; +use crate::state::StateKey; +use crate::storage::events::RedactBehaviour; +use crate::storage::sync_store::{ + anyhow_to_py, Borrowed, Interner, PersistedEventPosition, SyncReader, SyncStore, + SELECT_MANY_BATCH_SIZE, +}; +use crate::types::state_filter::StateFilter; +use crate::types::stream_token::{filter_results_by_stream, RoomStreamToken}; + +/// How often a connection's `last_used_ts` is rewritten when it is used, +/// in milliseconds: `UPDATE_INTERVAL_LAST_USED_TS`. +pub const UPDATE_INTERVAL_LAST_USED_TS_MS: i64 = 5 * 60 * 1000; + +/// How many times a write transaction is tried, as `runInteraction` tries +/// 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, 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>; + +/// `get_room_type`'s answer. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum RoomType { + /// `ROOM_UNKNOWN_SENTINEL`: the server has no current state for the + /// room, or no create event in it. + Unknown, + /// The create event's `type`, `None` without one. A `type` that is not + /// a string reads as none, as the sliding sync tables record it. + Known(Option>), +} + +/// `RoomsForUserStateReset`: a meaningful change to the user's membership +/// of a room, from `get_sliding_sync_membership_changes`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct MembershipChange { + pub room_id: Arc, + /// `None` after a state reset. + pub sender: Option>, + pub membership: Arc, + /// `None` after a state reset. + pub event_id: Option>, + pub event_pos: PersistedEventPosition, + pub room_version_id: Option>, +} + +/// `CurrentStateDeltaMembership`: a change to the user's membership in a +/// room's current state, and the membership it replaced. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DeltaMembershipChange { + pub room_id: Arc, + /// `None` when the server left the room or a state reset removed the + /// user. + pub event_id: Option>, + pub event_pos: PersistedEventPosition, + pub membership: Arc, + pub sender: Option>, + pub prev_event_id: Option>, + pub prev_event_pos: Option, + pub prev_membership: Option>, + pub prev_sender: Option>, +} + +/// `PerConnectionStateDB`: a connection's state as stored, with each +/// stream's tokens as their strings. Loading it gives the whole state at a +/// position. Storing one stores the request's updates. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct StoredConnectionState { + /// When the connection was last used, before this load. + pub last_used_ts: Option, + pub rooms: HashMap>, + pub receipts: HashMap>, + pub account_data: HashMap>, + pub room_configs: HashMap, +} + +impl StoredConnectionState { + /// `len()` of the `PerConnectionState` it becomes: what the memo of + /// loaded states weighs it as. + pub fn len(&self) -> usize { + self.rooms.len() + self.receipts.len() + self.account_data.len() + self.room_configs.len() + } + + pub fn is_empty(&self) -> bool { + self.len() == 0 + } +} + +/// The updates of one request to a connection's state, for +/// [`SlidingSyncReader::persist_connection_state`]: each stream's status +/// updates with their tokens as strings, the room configs set, and the +/// rooms' lazy-loaded memberships. +#[derive(Debug, Clone, Default)] +pub struct ConnectionStateUpdates { + pub rooms: HashMap>, + pub receipts: HashMap>, + pub account_data: HashMap>, + pub room_configs: HashMap, + pub room_lazy_membership: HashMap, +} + +/// The key of a loaded connection state: `(user_id, device_id, conn_id, +/// connection_position)`, the arguments of +/// `get_and_clear_connection_positions`. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct ConnectionPositionKey { + pub user_id: Arc, + pub device_id: Arc, + pub conn_id: Arc, + pub connection_position: i64, +} + +impl PyCacheKey for ConnectionPositionKey { + fn from_py_key(key: &Bound<'_, PyTuple>) -> Option { + let (user_id, device_id, conn_id, connection_position): (String, String, String, i64) = + key.extract().ok()?; + Some(ConnectionPositionKey { + user_id: Arc::from(user_id), + device_id: Arc::from(device_id), + conn_id: Arc::from(conn_id), + connection_position, + }) + } +} + +impl SyncStore { + /// A sliding sync reader for one request, over `pool`. + pub fn sliding_sync_reader(&self, pool: ConnectionPool) -> SlidingSyncReader<'_> { + SlidingSyncReader { + reader: self.reader(pool), + } + } +} + +/// The queries of one sliding sync request: the v2 reader's, through +/// `Deref`, and the sliding sync ones. See the module docs. +pub struct SlidingSyncReader<'a> { + reader: SyncReader<'a>, +} + +impl<'a> Deref for SlidingSyncReader<'a> { + type Target = SyncReader<'a>; + + fn deref(&self) -> &SyncReader<'a> { + &self.reader + } +} + +/// The columns of a membership snapshot, as `RoomsForUserSlidingSync` +/// reads them, less the room type and encryption, which each query takes +/// from its own tables. +const SNAPSHOT_COLUMNS: &str = "m.room_id, m.sender, m.membership, m.membership_event_id, + r.room_version, m.event_instance_name, m.event_stream_ordering, m.has_known_state"; + +/// A [`RoomForUser`] from a row of [`SNAPSHOT_COLUMNS`], the room type and +/// whether the room is encrypted. +fn snapshot_from_row(row: &Row, interner: &mut Interner) -> anyhow::Result { + let room_id: String = row.try_get(0)?; + let sender: String = row.try_get(1)?; + let membership: String = row.try_get(2)?; + let event_id: String = row.try_get(3)?; + let room_version: Option = row.try_get(4)?; + let instance_name: String = row.try_get(5)?; + let stream: i64 = row.try_get(6)?; + let has_known_state: bool = row.try_get(7)?; + let room_type: Option = row.try_get(8)?; + let is_encrypted: bool = row.try_get(9)?; + Ok(RoomForUser { + room_id: Arc::from(room_id), + sender: Some(Arc::from(sender)), + membership: interner.intern(&membership), + event_id: Some(Arc::from(event_id)), + event_pos: PersistedEventPosition { + instance_name: interner.intern(&instance_name), + stream, + }, + room_version_id: room_version.map(|version| interner.intern(&version)), + has_known_state, + room_type: room_type.map(|room_type| interner.intern(&room_type)), + is_encrypted, + }) +} + +/// Whether a snapshot's room version is one this server knows. A room +/// without a version is not. +fn has_known_version(room: &RoomForUser) -> bool { + room.room_version_id + .as_deref() + .is_some_and(is_known_room_version) +} + +/// A position from a nullable instance name: a null one is a row from +/// before sharded persisters, which Python reads as `"master"`. +fn position( + instance_name: Option, + stream: i64, + interner: &mut Interner, +) -> PersistedEventPosition { + PersistedEventPosition { + instance_name: interner.intern(instance_name.as_deref().unwrap_or("master")), + stream, + } +} + +/// Fail on a historic token, as the Python asserts. +fn assert_live(token: &RoomStreamToken) -> anyhow::Result<()> { + anyhow::ensure!(!token.is_historic(), "expected a live token, got {token:?}"); + Ok(()) +} + +/// Whether `e` is a serialization failure or a deadlock, which +/// `runInteraction` retries. +fn is_retryable(e: &anyhow::Error) -> bool { + e.chain().any(|cause| { + cause + .downcast_ref::() + .and_then(|e| e.code()) + .is_some_and(|code| { + *code == SqlState::T_R_SERIALIZATION_FAILURE + || *code == SqlState::T_R_DEADLOCK_DETECTED + }) + }) +} + +/// A room status as `sliding_sync_connection_streams` stores it: +/// `HaveSentRoomFlag`'s value and the token. +fn status_to_row(status: &HaveSentRoom) -> (&'static str, Option<&str>) { + match status { + HaveSentRoom::Never => ("never", None), + HaveSentRoom::Live => ("live", None), + HaveSentRoom::Previously(token) => ("previously", Some(token)), + } +} + +/// A room status from its row. A `previously` without a token is an error, +/// where Python would carry the `None` on until something read it. +fn status_from_row( + status: &str, + last_token: Option, +) -> anyhow::Result> { + match (status, last_token) { + ("never", _) => Ok(HaveSentRoom::Never), + ("live", _) => Ok(HaveSentRoom::Live), + ("previously", Some(token)) => Ok(HaveSentRoom::Previously(token)), + (status, token) => anyhow::bail!("invalid room status {status:?} with token {token:?}"), + } +} + +/// A room's required state as `sliding_sync_connection_required_state` +/// stores it: the sorted `[type, state_key]` pairs, encoded as Python's +/// `json_encoder` does, so the same state from either side dedups to one +/// row. +pub fn encode_required_state(map: &RequiredStateMap) -> String { + let mut out = String::from("["); + let mut first = true; + // The map and its sets are sorted, so this is the order of Python's + // `sorted` over the pairs. + for (event_type, state_keys) in map { + for state_key in state_keys { + if !first { + out.push(','); + } + first = false; + out.push('['); + crate::state::v2::push_py_json_str(&mut out, event_type); + out.push(','); + crate::state::v2::push_py_json_str(&mut out, state_key); + out.push(']'); + } + } + out.push(']'); + out +} + +/// The reverse of [`encode_required_state`]. +fn decode_required_state(json: &str) -> anyhow::Result { + let pairs: Vec<(String, String)> = + serde_json::from_str(json).with_context(|| format!("required state {json}"))?; + let mut map: RequiredStateMap = BTreeMap::new(); + for (event_type, state_key) in pairs { + map.entry(event_type).or_default().insert(state_key); + } + Ok(map) +} + +impl<'a> SlidingSyncReader<'a> { + // Membership snapshots. + + /// `get_sliding_sync_rooms_for_user_from_membership_snapshots`: the + /// user's rooms from their membership snapshots, less forgotten rooms, + /// rooms they left themselves, and rooms of unknown versions. + pub async fn rooms_from_snapshots(&self, user_id: &str) -> anyhow::Result> { + let key: Arc = Arc::from(user_id); + self.store() + .caches() + .sliding_sync_rooms_for_user + .get_or_load(Arc::clone(&key), || async { + let sql = format!( + "SELECT {SNAPSHOT_COLUMNS}, + COALESCE(j.room_type, m.room_type), + COALESCE(j.is_encrypted, m.is_encrypted) + FROM sliding_sync_membership_snapshots AS m + INNER JOIN rooms AS r USING (room_id) + LEFT JOIN sliding_sync_joined_rooms AS j + ON (j.room_id = m.room_id AND m.membership = 'join') + WHERE user_id = $1 + AND m.forgotten = 0 + AND (m.membership != 'leave' OR m.user_id != m.sender)" + ); + let rows = self.query_one_shot(&sql, &[&&*key]).await?; + let mut interner = Interner::default(); + let mut rooms = SnapshotRooms::with_capacity(rows.len()); + for row in &rows { + let room = snapshot_from_row(row, &mut interner)?; + // Rooms of versions we don't know shouldn't go down sync, + // and their metadata may be broken. + if has_known_version(&room) { + rooms.insert(Arc::clone(&room.room_id), room); + } + } + Ok(Arc::new(rooms)) + }) + .await + .map_err(anyhow::Error::new) + } + + /// `get_sliding_sync_self_leave_rooms_after_to_token`: the rooms the + /// user left themselves after `to_token`, from their snapshots, which + /// may have been relevant before the leave. + pub async fn self_leave_rooms_after( + &self, + user_id: &str, + to_token: &RoomStreamToken, + ) -> anyhow::Result { + // The token's instance map covers more of the stream than its + // minimum: fetch from the minimum and filter. + let min_to_token_position = to_token.stream; + let sql = format!( + "SELECT {SNAPSHOT_COLUMNS}, m.room_type, m.is_encrypted + FROM sliding_sync_membership_snapshots AS m + INNER JOIN rooms AS r USING (room_id) + WHERE user_id = $1 + AND m.forgotten = 0 + AND m.membership = 'leave' + AND m.user_id = m.sender + AND (m.event_stream_ordering > $2)" + ); + let rows = self + .query_one_shot(&sql, &[&user_id, &min_to_token_position]) + .await?; + let mut interner = Interner::default(); + let mut rooms = SnapshotRooms::new(); + for row in &rows { + let room = snapshot_from_row(row, &mut interner)?; + if !has_known_version(&room) { + continue; + } + // Only the leaves after the token. + if filter_results_by_stream( + None, + Some(to_token), + Some(&room.event_pos.instance_name), + room.event_pos.stream, + )? { + continue; + } + rooms.insert(Arc::clone(&room.room_id), room); + } + Ok(rooms) + } + + /// `get_sliding_sync_room_for_user`: the user's snapshot of the room, + /// forgotten or not left, whatever its version. + pub async fn snapshot_room( + &self, + user_id: &str, + room_id: &str, + ) -> anyhow::Result> { + let sql = format!( + "SELECT {SNAPSHOT_COLUMNS}, + COALESCE(j.room_type, m.room_type), + COALESCE(j.is_encrypted, m.is_encrypted) + FROM sliding_sync_membership_snapshots AS m + INNER JOIN rooms AS r USING (room_id) + LEFT JOIN sliding_sync_joined_rooms AS j + ON (j.room_id = m.room_id AND m.membership = 'join') + WHERE user_id = $1 + AND m.forgotten = 0 + AND m.room_id = $2" + ); + let rows = self.query_one_shot(&sql, &[&user_id, &room_id]).await?; + let mut interner = Interner::default(); + rows.first() + .map(|row| snapshot_from_row(row, &mut interner)) + .transpose() + } + + /// `get_sliding_sync_room_for_user_batch`: [`Self::snapshot_room`] for + /// each of the rooms. + pub async fn snapshot_rooms( + &self, + user_id: &str, + room_ids: &[Arc], + ) -> anyhow::Result { + let mut rooms = SnapshotRooms::with_capacity(room_ids.len()); + if room_ids.is_empty() { + return Ok(rooms); + } + let sql = format!( + "SELECT {SNAPSHOT_COLUMNS}, + COALESCE(j.room_type, m.room_type), + COALESCE(j.is_encrypted, m.is_encrypted) + FROM sliding_sync_membership_snapshots AS m + INNER JOIN rooms AS r USING (room_id) + LEFT JOIN sliding_sync_joined_rooms AS j + ON (j.room_id = m.room_id AND m.membership = 'join') + WHERE m.forgotten = 0 + AND m.room_id = ANY($1) + AND user_id = $2" + ); + let ids: Vec<&str> = room_ids.iter().map(|id| &**id).collect(); + let rows = self.query_one_shot(&sql, &[&ids, &user_id]).await?; + let mut interner = Interner::default(); + for row in &rows { + let room = snapshot_from_row(row, &mut interner)?; + rooms.insert(Arc::clone(&room.room_id), room); + } + Ok(rooms) + } + + // Membership changes. + + /// `get_sliding_sync_membership_changes`: the meaningful changes to the + /// user's membership after `from_token` and up to `to_token`, by room, + /// from the snapshots and, for state resets, the current state deltas. + /// In stream order of each room's first change. + /// + /// `membership_may_have_changed` is + /// `_membership_stream_cache.has_entity_changed(user_id, from_token.stream)`. + pub async fn snapshot_membership_changes( + &self, + user_id: &str, + from_token: &RoomStreamToken, + to_token: &RoomStreamToken, + excluded_room_ids: &HashSet>, + membership_may_have_changed: bool, + ) -> anyhow::Result, MembershipChange>> { + assert_live(from_token)?; + assert_live(to_token)?; + let mut changes = IndexMap::new(); + if from_token == to_token || !membership_may_have_changed { + return Ok(changes); + } + + // The tokens' instance maps cover more of the stream than their + // minimums: fetch up to the maximum and filter. + let min_from_id = from_token.stream; + let max_to_id = to_token.get_max_stream_pos(); + // The snapshots don't include users state reset out of a room, so + // the current state deltas are read for those. + let rows = self + .query_one_shot( + "SELECT + room_id, + membership_event_id, + event_instance_name, + event_stream_ordering, + membership, + sender, + prev_membership, + room_version + FROM + ( + SELECT + s.room_id, + s.membership_event_id, + s.event_instance_name, + s.event_stream_ordering, + s.membership, + s.sender, + m_prev.membership AS prev_membership + FROM sliding_sync_membership_snapshots as s + LEFT JOIN event_edges AS e ON e.event_id = s.membership_event_id + LEFT JOIN room_memberships AS m_prev ON m_prev.event_id = e.prev_event_id + WHERE s.user_id = $1 + + UNION ALL + + SELECT + s.room_id, + e.event_id, + s.instance_name, + s.stream_id, + m.membership, + e.sender, + m_prev.membership AS prev_membership + FROM current_state_delta_stream AS s + LEFT JOIN events AS e ON e.event_id = s.event_id + LEFT JOIN room_memberships AS m ON m.event_id = s.event_id + LEFT JOIN room_memberships AS m_prev ON m_prev.event_id = s.prev_event_id + WHERE + s.type = 'm.room.member' + AND s.state_key = $1 + ) AS c + INNER JOIN rooms USING (room_id) + WHERE event_stream_ordering > $2 AND event_stream_ordering <= $3 + ORDER BY event_stream_ordering ASC", + &[&user_id, &min_from_id, &max_to_id], + ) + .await?; + + let mut interner = Interner::default(); + for row in &rows { + let room_id: String = row.try_get(0)?; + let membership_event_id: Option = row.try_get(1)?; + let instance_name: Option = row.try_get(2)?; + let stream: i64 = row.try_get(3)?; + let membership: Option = row.try_get(4)?; + let sender: Option = row.try_get(5)?; + let prev_membership: Option = row.try_get(6)?; + let room_version: Option = row.try_get(7)?; + + if excluded_room_ids.contains(room_id.as_str()) { + continue; + } + if !filter_results_by_stream( + Some(from_token), + Some(to_token), + instance_name.as_deref(), + stream, + )? { + continue; + } + // When the server leaves a room, every piece of current state + // gets a row with no event, the user's leave included. Don't + // report a leave twice. + if membership_event_id.is_none() && prev_membership.as_deref() == Some("leave") { + continue; + } + // Python's "SUSPICIOUS" case, ported as it is: a state reset + // after a change in the same window is not reported. + if membership_event_id.is_none() && changes.contains_key(room_id.as_str()) { + continue; + } + // A row without an event is the server leaving or a state + // reset, which leave the user out of the room. + let membership = membership.unwrap_or_else(|| "leave".to_owned()); + // The same membership again, such as a display name change. + if prev_membership.as_deref() == Some(membership.as_str()) { + continue; + } + let room_id: Arc = Arc::from(room_id); + changes.insert( + Arc::clone(&room_id), + MembershipChange { + room_id, + sender: sender.map(Arc::from), + membership: interner.intern(&membership), + event_id: membership_event_id.map(Arc::from), + event_pos: position(instance_name, stream, &mut interner), + room_version_id: room_version.map(|version| interner.intern(&version)), + }, + ); + } + Ok(changes) + } + + /// `get_current_state_delta_membership_changes_for_user`: the changes + /// to the user's membership in the current state of their rooms after + /// `from_token` and up to `to_token`, in stream order, with the + /// membership each replaced. + /// + /// `membership_may_have_changed` is as for + /// [`Self::snapshot_membership_changes`]. + pub async fn membership_delta_changes( + &self, + user_id: &str, + from_token: &RoomStreamToken, + to_token: &RoomStreamToken, + excluded_room_ids: &HashSet>, + membership_may_have_changed: bool, + ) -> anyhow::Result> { + assert_live(from_token)?; + assert_live(to_token)?; + if from_token == to_token || !membership_may_have_changed { + return Ok(Vec::new()); + } + + let min_from_id = from_token.stream; + let max_to_id = to_token.get_max_stream_pos(); + // The delta's own position, not its event's: see the Python for why + // that is enough between two valid tokens. + let rows = self + .query_one_shot( + "SELECT + s.room_id, + e.event_id, + s.instance_name, + s.stream_id, + m.membership, + e.sender, + s.prev_event_id, + e_prev.instance_name AS prev_instance_name, + e_prev.stream_ordering AS prev_stream_ordering, + m_prev.membership AS prev_membership, + e_prev.sender AS prev_sender + FROM current_state_delta_stream AS s + LEFT JOIN events AS e ON e.event_id = s.event_id + LEFT JOIN room_memberships AS m ON m.event_id = s.event_id + LEFT JOIN events AS e_prev ON e_prev.event_id = s.prev_event_id + LEFT JOIN room_memberships AS m_prev ON m_prev.event_id = s.prev_event_id + WHERE s.stream_id > $1 AND s.stream_id <= $2 + AND s.type = 'm.room.member' + AND s.state_key = $3 + ORDER BY s.stream_id ASC", + &[&min_from_id, &max_to_id, &user_id], + ) + .await?; + + let mut interner = Interner::default(); + let mut changes = Vec::new(); + for row in &rows { + let room_id: String = row.try_get(0)?; + let event_id: Option = row.try_get(1)?; + let instance_name: Option = row.try_get(2)?; + let stream: i64 = row.try_get(3)?; + let membership: Option = row.try_get(4)?; + let sender: Option = row.try_get(5)?; + let prev_event_id: Option = row.try_get(6)?; + let prev_instance_name: Option = row.try_get(7)?; + let prev_stream: Option = row.try_get(8)?; + let prev_membership: Option = row.try_get(9)?; + let prev_sender: Option = row.try_get(10)?; + + if !filter_results_by_stream( + Some(from_token), + Some(to_token), + instance_name.as_deref(), + stream, + )? { + continue; + } + // As in `snapshot_membership_changes`: the server's leave after + // the user's own. + if event_id.is_none() && prev_membership.as_deref() == Some("leave") { + continue; + } + if excluded_room_ids.contains(room_id.as_str()) { + continue; + } + changes.push(DeltaMembershipChange { + room_id: Arc::from(room_id), + event_id: event_id.map(Arc::from), + event_pos: position(instance_name, stream, &mut interner), + membership: interner.intern(membership.as_deref().unwrap_or("leave")), + sender: sender.map(Arc::from), + prev_event_id: prev_event_id.map(Arc::from), + prev_event_pos: prev_stream + .map(|stream| position(prev_instance_name, stream, &mut interner)), + prev_membership: prev_membership.map(|membership| interner.intern(&membership)), + prev_sender: prev_sender.map(Arc::from), + }); + } + Ok(changes) + } + + // Stream positions. + + /// `get_rooms_that_have_updates_since_sliding_sync_table`: which of the + /// rooms may have new events after `from_token`, false positives + /// allowed. + /// + /// `changed_since` is `_events_stream_cache.get_all_entities_changed` + /// at `from_token.stream`: the rooms with new events, or `None` if the + /// cache does not reach back that far, when the table answers. + pub async fn rooms_with_updates_since( + &self, + room_ids: &[Arc], + from_token: &RoomStreamToken, + changed_since: Option<&HashSet>>, + ) -> anyhow::Result>> { + if let Some(changed) = changed_since { + return Ok(room_ids + .iter() + .filter(|room_id| changed.contains(*room_id)) + .cloned() + .collect()); + } + let mut rooms = HashSet::new(); + if room_ids.is_empty() { + return Ok(rooms); + } + let borrowed = self.borrow().await?; + for batch in room_ids.chunks(SELECT_MANY_BATCH_SIZE) { + let ids: Vec<&str> = batch.iter().map(|id| &**id).collect(); + let rows = borrowed + .query( + "SELECT room_id FROM sliding_sync_joined_rooms + WHERE room_id = ANY($1) AND event_stream_ordering > $2", + &[&ids, &from_token.stream], + ) + .await?; + for row in rows { + let room_id: String = row.try_get(0)?; + rooms.insert(Arc::from(room_id)); + } + } + Ok(rooms) + } + + /// `_bulk_get_max_event_pos`: the stream position of the latest event + /// in each room, capped at `now_token`, the current room key + /// (`get_room_max_token`), or `None` for a room without a + /// `sliding_sync_joined_rooms` row. Through the mirror of + /// `_get_max_event_pos`. + /// + /// `last_change_pos` is `_events_stream_cache.get_max_pos_of_last_change` + /// for the rooms it knows. As in Python, the misses are answered from it + /// alone, uncapped, when it knows them all. Otherwise every miss goes to + /// the database, whose answer wins, and its positions stand only for the + /// rooms without a row: Python's batch loop runs over all the rooms + /// rather than the ones it didn't know. + pub async fn max_event_pos( + &self, + room_ids: &[Arc], + now_token: &RoomStreamToken, + last_change_pos: &HashMap, i64>, + ) -> anyhow::Result, Option>> { + self.store() + .caches() + .max_event_pos + .get_or_load_many(room_ids, |missing| async move { + let mut results: HashMap, Option> = missing + .iter() + .map(|room_id| (Arc::clone(room_id), last_change_pos.get(room_id).copied())) + .collect(); + if results.values().all(Option::is_some) { + return Ok(results); + } + let mut recheck = Vec::new(); + { + let borrowed = self.borrow().await?; + for batch in missing.chunks(SELECT_MANY_BATCH_SIZE) { + let ids: Vec<&str> = batch.iter().map(|id| &**id).collect(); + let rows = borrowed + .query( + "SELECT room_id, event_stream_ordering + FROM sliding_sync_joined_rooms + WHERE room_id = ANY($1) + ORDER BY event_stream_ordering DESC", + &[&ids], + ) + .await?; + for row in rows { + let room_id: String = row.try_get(0)?; + let stream: i64 = row.try_get(1)?; + let room_id: Arc = Arc::from(room_id); + if stream <= now_token.stream { + results.insert(room_id, Some(stream)); + } else { + recheck.push(room_id); + } + } + } + } + // Rooms whose latest event may be ahead of the token, which + // is rare, one by one. + for room_id in recheck { + if let Some((_, position)) = self.last_event_before(&room_id, now_token).await? + { + results.insert(room_id, Some(position.stream)); + } + } + Ok(results) + }) + .await + .map_err(anyhow::Error::new) + } + + /// `bulk_get_last_event_pos_in_room_before_stream_ordering`: the stream + /// position of the last event at or before `end_token` in each room + /// that has one. `now_token` and `last_change_pos` are as for + /// [`Self::max_event_pos`]. + pub async fn last_event_pos_before( + &self, + room_ids: &[Arc], + end_token: &RoomStreamToken, + now_token: &RoomStreamToken, + last_change_pos: &HashMap, i64>, + ) -> anyhow::Result, i64>> { + // Most rooms' latest event is before the token, and cached. + let uncapped = self + .max_event_pos(room_ids, now_token, last_change_pos) + .await?; + let mut results = HashMap::with_capacity(uncapped.len()); + let mut recheck = Vec::new(); + let min_token = end_token.stream; + for (room_id, stream) in uncapped { + // `None` is what the cache holds for a room the load didn't + // find. + let Some(stream) = stream else { continue }; + if stream <= min_token { + results.insert(room_id, stream); + } else { + recheck.push(room_id); + } + } + for room_id in recheck { + if let Some((_, position)) = self.last_event_before(&room_id, end_token).await? { + results.insert(room_id, position.stream); + } + } + Ok(results) + } + + // Room metadata. + + /// `get_member_counts`, through its mirror. + pub async fn cached_member_counts( + &self, + room_id: &str, + ) -> anyhow::Result>> { + self.store() + .caches() + .member_counts + .get_or_load(Arc::from(room_id), || async { + Ok(Arc::new(self.member_counts(room_id).await?)) + }) + .await + .map_err(anyhow::Error::new) + } + + /// `get_room_type`: the type in the room's create event, from its + /// current state. + pub async fn room_type(&self, room_id: &str) -> anyhow::Result { + self.store() + .caches() + .room_types + .get_or_load(Arc::from(room_id), || async { + let filter = StateFilter::from_types([("m.room.create", Some(""))]); + let state = self.current_state_ids(room_id, &filter).await?; + let Some(create_id) = state.get(&StateKey::new("m.room.create", "")) else { + return Ok(RoomType::Unknown); + }; + let events = self + .events( + std::slice::from_ref(create_id), + RedactBehaviour::Redact, + false, + false, + ) + .await?; + let Some(create) = events.first() else { + return Ok(RoomType::Unknown); + }; + Ok(RoomType::Known( + create + .content() + .get_field("type") + .and_then(|room_type| room_type.as_str()) + .map(Arc::from), + )) + }) + .await + .map_err(anyhow::Error::new) + } + + /// `get_room_version_id`: the room's version, or `None` where Python + /// raises `NotFoundError`. + pub async fn room_version_id(&self, room_id: &str) -> anyhow::Result>> { + /// The load's answer for a room without a version, so that it is not + /// cached. + #[derive(Debug)] + struct NoVersion; + impl std::fmt::Display for NoVersion { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("no room version") + } + } + impl std::error::Error for NoVersion {} + + let result = self + .store() + .caches() + .room_version_ids + .get_or_load(Arc::from(room_id), || async { + let rows = self + .query_one_shot( + "SELECT room_version FROM rooms WHERE room_id = $1", + &[&room_id], + ) + .await?; + let version: Option = match rows.first() { + Some(row) => row.try_get(0)?, + None => None, + }; + version + .map(Arc::from) + .ok_or_else(|| anyhow::Error::new(NoVersion)) + }) + .await; + match result { + Ok(version) => Ok(Some(version)), + Err(e) if e.inner().is::() => Ok(None), + Err(e) => Err(anyhow::Error::new(e)), + } + } + + /// The `instance_map` table: each writer's instance ID and name, as + /// stream tokens name writers by ID. Cached: see + /// [`Self::forget_instance_map`]. + pub async fn instance_map(&self) -> anyhow::Result)]>> { + self.store() + .caches() + .instance_map + .get_or_load((), || async { + 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() + }) + .await + .map_err(anyhow::Error::new) + } + + /// Drop the cached `instance_map`, for a caller that found a writer + /// missing from it. A writer added since the cache was filled is only in + /// the table. + pub fn forget_instance_map(&self) { + drop(self.store().caches().instance_map.invalidate(&())); + } + + /// `get_partial_rooms`: every room the server joined with partial + /// state, which is never many. + pub async fn partial_rooms(&self) -> anyhow::Result>>> { + self.store() + .caches() + .partial_rooms + .get_or_load((), || async { + let rows = self + .query_one_shot("SELECT room_id FROM partial_state_rooms", &[]) + .await?; + let mut rooms = HashSet::with_capacity(rows.len()); + for row in rows { + let room_id: String = row.try_get(0)?; + rooms.insert(Arc::from(room_id)); + } + Ok(Arc::new(rooms)) + }) + .await + .map_err(anyhow::Error::new) + } + + /// `get_tags_for_user`: the user's tags, by room. + pub async fn tags_for_user(&self, user_id: &str) -> anyhow::Result> { + self.store() + .caches() + .tags_for_user + .get_or_load(Arc::from(user_id), || async { + let rows = self + .query_one_shot( + "SELECT room_id, tag, content FROM room_tags WHERE user_id = $1", + &[&user_id], + ) + .await?; + let mut tags = UserTags::new(); + for row in rows { + let room_id: String = row.try_get(0)?; + let tag: String = row.try_get(1)?; + let content: String = row.try_get(2)?; + let content = serde_json::from_str(&content) + .with_context(|| format!("tag {tag} of {user_id} in {room_id}"))?; + tags.entry(Arc::from(room_id)) + .or_default() + .insert(tag, content); + } + Ok(Arc::new(tags)) + }) + .await + .map_err(anyhow::Error::new) + } + + /// `get_is_encrypted_for_room_at_token`: whether the room was encrypted + /// at `to_token`, from its current state rolled back by the deltas + /// since. `now_token` is the current room key. + pub async fn is_encrypted_at( + &self, + room_id: &str, + to_token: &RoomStreamToken, + now_token: &RoomStreamToken, + ) -> anyhow::Result { + let filter = StateFilter::from_types([("m.room.encryption", Some(""))]); + let state = self.current_state_ids(room_id, &filter).await?; + let mut encryption_event_id = state.get(&StateKey::new("m.room.encryption", "")).cloned(); + + // The first change after the token holds the state at it. + let deltas = self + .state_deltas_by_stream_id(room_id, Some(to_token), Some(now_token)) + .await?; + if let Some(delta) = deltas + .iter() + .find(|delta| &*delta.event_type == "m.room.encryption") + { + match &delta.prev_event_id { + Some(prev_event_id) => encryption_event_id = Some(Arc::clone(prev_event_id)), + // No encryption event before the change. + None => return Ok(false), + } + } + + let Some(encryption_event_id) = encryption_event_id else { + return Ok(false); + }; + let events = self + .events( + std::slice::from_ref(&encryption_event_id), + RedactBehaviour::Redact, + false, + false, + ) + .await?; + // Python's `get_event` raises for an event it can't find. + let event = events + .first() + .with_context(|| format!("encryption event {encryption_event_id} not found"))?; + Ok(event + .content() + .get_field("algorithm") + .is_some_and(|algorithm| !algorithm.is_null())) + } + + // Connection state. + + /// `get_sliding_sync_connection_lazy_members`: which of the users the + /// connection sent lazy-loaded memberships for in the room, as of + /// `connection_position`, to when each was last needed. + pub async fn connection_lazy_members( + &self, + connection_position: i64, + room_id: &str, + user_ids: &[Arc], + ) -> anyhow::Result, i64>> { + let mut members = HashMap::new(); + if user_ids.is_empty() { + return Ok(members); + } + let ids: Vec<&str> = user_ids.iter().map(|id| &**id).collect(); + // The position's connection is found through its row, as the caller + // doesn't have the connection key. Of the two `connection_position` + // columns, the members' is the one filtered on. + let rows = self + .query_one_shot( + "SELECT user_id, members.connection_position, last_seen_ts + FROM sliding_sync_connection_lazy_members AS members + INNER JOIN sliding_sync_connection_positions AS pos USING (connection_key) + WHERE pos.connection_position = $1 AND room_id = $2 AND user_id = ANY($3)", + &[&connection_position, &room_id, &ids], + ) + .await?; + for row in rows { + let user_id: String = row.try_get(0)?; + let member_position: Option = row.try_get(1)?; + let last_seen_ts: i64 = row.try_get(2)?; + // An entry with a position applies to that fork of the + // connection only. One without applies to every position. + if member_position.is_none_or(|position| position == connection_position) { + members.insert(Arc::from(user_id), last_seen_ts); + } + } + Ok(members) + } + + /// `get_and_clear_connection_positions`: the connection's state at + /// `connection_position`, or [`UnknownPosition`] if the position is not + /// one of this user's device's connection. Loading a position commits + /// the client to it: the connection's other positions are deleted. + /// + /// Memoized by its arguments, as Python's `@cached` is, so a long-poll's + /// second call and a client's retry don't run the transaction again. + /// `now_ms` is the time to record as the connection's last use. + pub async fn load_connection_state( + &self, + user_id: &str, + device_id: &str, + conn_id: &str, + connection_position: i64, + now_ms: i64, + ) -> anyhow::Result> { + let key = ConnectionPositionKey { + user_id: Arc::from(user_id), + device_id: Arc::from(device_id), + conn_id: Arc::from(conn_id), + connection_position, + }; + self.store() + .caches() + .connection_states + .get_or_load(key, || async { + self.retry_transaction(|| { + self.load_connection_state_txn( + user_id, + device_id, + conn_id, + connection_position, + now_ms, + ) + .boxed() + }) + .await + .map(Arc::new) + }) + .await + .map_err(|e| match e.inner().downcast_ref::() { + Some(unknown) => anyhow::Error::new(unknown.clone()), + None => anyhow::Error::new(e), + }) + } + + /// Run `txn` until it succeeds, fails with an error other than a + /// serialization failure or a deadlock, or has been tried + /// [`MAX_TRANSACTION_ATTEMPTS`] times. + async fn retry_transaction<'s, T>( + &'s self, + txn: impl Fn() -> BoxFuture<'s, anyhow::Result>, + ) -> anyhow::Result { + let mut attempt = 1; + loop { + match txn().await { + Err(e) if attempt < MAX_TRANSACTION_ATTEMPTS && is_retryable(&e) => { + warn!("Retrying a sliding sync transaction, attempt {attempt}: {e:#}"); + attempt += 1; + } + result => return result, + } + } + } + + async fn load_connection_state_txn( + &self, + user_id: &str, + device_id: &str, + conn_id: &str, + connection_position: i64, + now_ms: i64, + ) -> anyhow::Result { + let txn = self.begin_write().await?; + + // The position is the client's to name, so check it is theirs. + let rows = txn + .query( + "SELECT connection_key, last_used_ts + FROM sliding_sync_connection_positions + INNER JOIN sliding_sync_connections USING (connection_key) + WHERE + connection_position = $1 + AND user_id = $2 AND effective_device_id = $3 AND conn_id = $4", + &[&connection_position, &user_id, &device_id, &conn_id], + ) + .await?; + let Some(row) = rows.first() else { + return Err(UnknownPosition.into()); + }; + let connection_key: i64 = row.try_get(0)?; + let last_used_ts: Option = row.try_get(1)?; + + // Not on every use, to spare the writes. + if last_used_ts.is_none_or(|ts| now_ms - ts > UPDATE_INTERVAL_LAST_USED_TS_MS) { + txn.query( + "UPDATE sliding_sync_connections SET last_used_ts = $1 WHERE connection_key = $2", + &[&now_ms, &connection_key], + ) + .await?; + } + + // The client has the position, so the others can go. What follows + // relies on it being the only one left. + txn.query( + "DELETE FROM sliding_sync_connection_positions + WHERE connection_key = $1 AND connection_position != $2", + &[&connection_key, &connection_position], + ) + .await?; + + // With the forks gone, the position's lazy members hold for every + // position after it. + txn.query( + "UPDATE sliding_sync_connection_lazy_members SET connection_position = NULL + WHERE connection_key = $1 AND connection_position = $2", + &[&connection_key, &connection_position], + ) + .await?; + + let rows = txn + .query( + "SELECT required_state_id, required_state + FROM sliding_sync_connection_required_state + WHERE connection_key = $1", + &[&connection_key], + ) + .await?; + let mut required_states: HashMap = + HashMap::with_capacity(rows.len()); + for row in &rows { + let required_state_id: i64 = row.try_get(0)?; + let required_state: String = row.try_get(1)?; + required_states.insert(required_state_id, decode_required_state(&required_state)?); + } + + let rows = txn + .query( + "SELECT room_id, timeline_limit, required_state_id + FROM sliding_sync_connection_room_configs + WHERE connection_position = $1", + &[&connection_position], + ) + .await?; + let mut room_configs = HashMap::with_capacity(rows.len()); + let mut used_required_state_ids = HashSet::new(); + for row in &rows { + let room_id: String = row.try_get(0)?; + let timeline_limit: i64 = row.try_get(1)?; + let required_state_id: i64 = row.try_get(2)?; + let required_state_map = required_states + .get(&required_state_id) + .with_context(|| format!("required state {required_state_id} of {room_id}"))? + .clone(); + used_required_state_ids.insert(required_state_id); + room_configs.insert( + room_id, + RoomSyncConfig { + timeline_limit, + required_state_map, + }, + ); + } + + // The required states no room of the one remaining position uses. + let unused: Vec = required_states + .keys() + .filter(|id| !used_required_state_ids.contains(*id)) + .copied() + .collect(); + if !unused.is_empty() { + txn.query( + "DELETE FROM sliding_sync_connection_required_state + WHERE connection_key = $1 AND required_state_id = ANY($2)", + &[&connection_key, &unused], + ) + .await?; + } + + let rows = txn + .query( + "SELECT stream, room_id, room_status, last_token + FROM sliding_sync_connection_streams + WHERE connection_position = $1", + &[&connection_position], + ) + .await?; + let mut state = StoredConnectionState { + last_used_ts, + room_configs, + ..Default::default() + }; + for row in &rows { + let stream: String = row.try_get(0)?; + let room_id: String = row.try_get(1)?; + let room_status: String = row.try_get(2)?; + let last_token: Option = row.try_get(3)?; + let statuses = match stream.as_str() { + "rooms" => &mut state.rooms, + "receipts" => &mut state.receipts, + "account_data" => &mut state.account_data, + // Later versions may add streams. + _ => { + warn!("Unrecognized sliding sync stream in DB {stream:?}"); + continue; + } + }; + statuses.insert( + room_id, + status_from_row(&room_status, last_token) + .with_context(|| format!("stream {stream} of connection {connection_key}"))?, + ); + } + + txn.commit().await?; + Ok(state) + } + + /// `persist_per_connection_state`: store the request's `updates` to + /// the connection's state at `previous_position` as a new position, + /// and return it. Without a previous position the connection starts + /// over, dropping what was stored for it. [`UnknownPosition`] if the + /// previous position is not this user's device's connection's. + #[allow(clippy::too_many_arguments)] + pub async fn persist_connection_state( + &self, + user_id: &str, + device_id: &str, + conn_id: &str, + previous_position: Option, + updates: &ConnectionStateUpdates, + now_ms: i64, + ) -> anyhow::Result { + self.retry_transaction(|| { + self.persist_connection_state_txn( + user_id, + device_id, + conn_id, + previous_position, + updates, + now_ms, + ) + .boxed() + }) + .await + } + + async fn persist_connection_state_txn( + &self, + user_id: &str, + device_id: &str, + conn_id: &str, + previous_position: Option, + updates: &ConnectionStateUpdates, + now_ms: i64, + ) -> anyhow::Result { + let txn = self.begin_write().await?; + + let connection_key: i64 = match previous_position { + Some(previous_position) => { + // Lock the connection up front. Concurrent upserts of its lazy + // members deadlocked otherwise. The position is the client's + // to name, so check it is theirs. + let rows = txn + .query( + "SELECT connection_key + FROM sliding_sync_connection_positions + INNER JOIN sliding_sync_connections USING (connection_key) + WHERE + connection_position = $1 + AND user_id = $2 AND effective_device_id = $3 AND conn_id = $4 + FOR NO KEY UPDATE OF sliding_sync_connections", + &[&previous_position, &user_id, &device_id, &conn_id], + ) + .await?; + let Some(row) = rows.first() else { + return Err(UnknownPosition.into()); + }; + row.try_get(0)? + } + None => { + // Starting over: drop the connection's old state, which + // cascades to its positions and their rows, so one-shot + // requests don't pile up. + txn.query( + "DELETE FROM sliding_sync_connections + WHERE user_id = $1 AND effective_device_id = $2 AND conn_id = $3", + &[&user_id, &device_id, &conn_id], + ) + .await?; + let rows = txn + .query( + "INSERT INTO sliding_sync_connections + (user_id, effective_device_id, conn_id, created_ts, last_used_ts) + VALUES ($1, $2, $3, $4, $4) + RETURNING connection_key", + &[&user_id, &device_id, &conn_id, &now_ms], + ) + .await?; + rows.first() + .context("no connection key returned")? + .try_get(0)? + } + }; + + let rows = txn + .query( + "INSERT INTO sliding_sync_connection_positions (connection_key, created_ts) + VALUES ($1, $2) + RETURNING connection_position", + &[&connection_key, &now_ms], + ) + .await?; + let connection_position: i64 = rows + .first() + .context("no connection position returned")? + .try_get(0)?; + + // Rooms share required state rows by their JSON, within the + // connection. + let mut required_state_to_id: HashMap = HashMap::new(); + if previous_position.is_some() { + let rows = txn + .query( + "SELECT required_state_id, required_state + FROM sliding_sync_connection_required_state + WHERE connection_key = $1", + &[&connection_key], + ) + .await?; + for row in &rows { + required_state_to_id.insert(row.try_get(1)?, row.try_get(0)?); + } + } + let mut room_to_state_id: HashMap<&str, i64> = HashMap::new(); + let mut unique_required_state: IndexMap> = IndexMap::new(); + for (room_id, config) in &updates.room_configs { + let serialized = encode_required_state(&config.required_state_map); + match required_state_to_id.get(&serialized) { + Some(id) => { + room_to_state_id.insert(room_id, *id); + } + None => unique_required_state + .entry(serialized) + .or_default() + .push(room_id), + } + } + for (serialized, room_ids) in &unique_required_state { + let rows = txn + .query( + "INSERT INTO sliding_sync_connection_required_state + (connection_key, required_state) + VALUES ($1, $2) + RETURNING required_state_id", + &[&connection_key, serialized], + ) + .await?; + let required_state_id: i64 = rows + .first() + .context("no required state ID returned")? + .try_get(0)?; + for room_id in room_ids { + room_to_state_id.insert(room_id, required_state_id); + } + } + + // Start from the previous position's rows. The updates overwrite + // them. + if let Some(previous_position) = previous_position { + txn.query( + "INSERT INTO sliding_sync_connection_streams + (connection_position, stream, room_id, room_status, last_token) + SELECT $1, stream, room_id, room_status, last_token + FROM sliding_sync_connection_streams + WHERE connection_position = $2", + &[&connection_position, &previous_position], + ) + .await?; + txn.query( + "INSERT INTO sliding_sync_connection_room_configs + (connection_position, room_id, timeline_limit, required_state_id) + SELECT $1, room_id, timeline_limit, required_state_id + FROM sliding_sync_connection_room_configs + WHERE connection_position = $2", + &[&connection_position, &previous_position], + ) + .await?; + } + + let mut streams: Vec<&str> = Vec::new(); + let mut stream_rooms: Vec<&str> = Vec::new(); + let mut statuses: Vec<&str> = Vec::new(); + let mut tokens: Vec> = Vec::new(); + for (stream, updates) in [ + ("rooms", &updates.rooms), + ("receipts", &updates.receipts), + ("account_data", &updates.account_data), + ] { + for (room_id, status) in updates { + let (status, token) = status_to_row(status); + streams.push(stream); + stream_rooms.push(room_id); + statuses.push(status); + tokens.push(token); + } + } + if !streams.is_empty() { + txn.query( + "INSERT INTO sliding_sync_connection_streams + (connection_position, stream, room_id, room_status, last_token) + SELECT $1, * FROM UNNEST($2::text[], $3::text[], $4::text[], $5::text[]) + ON CONFLICT (connection_position, stream, room_id) + DO UPDATE SET room_status = EXCLUDED.room_status, last_token = EXCLUDED.last_token", + &[ + &connection_position, + &streams, + &stream_rooms, + &statuses, + &tokens, + ], + ) + .await?; + } + + let mut config_rooms: Vec<&str> = Vec::with_capacity(updates.room_configs.len()); + let mut timeline_limits: Vec = Vec::with_capacity(updates.room_configs.len()); + let mut required_state_ids: Vec = Vec::with_capacity(updates.room_configs.len()); + for (room_id, config) in &updates.room_configs { + config_rooms.push(room_id); + timeline_limits.push(config.timeline_limit); + required_state_ids.push(room_to_state_id[room_id.as_str()]); + } + if !config_rooms.is_empty() { + txn.query( + "INSERT INTO sliding_sync_connection_room_configs + (connection_position, room_id, timeline_limit, required_state_id) + SELECT $1, * FROM UNNEST($2::text[], $3::bigint[], $4::bigint[]) + ON CONFLICT (connection_position, room_id) + DO UPDATE SET + timeline_limit = EXCLUDED.timeline_limit, + required_state_id = EXCLUDED.required_state_id", + &[ + &connection_position, + &config_rooms, + &timeline_limits, + &required_state_ids, + ], + ) + .await?; + } + + persist_lazy_members( + &txn, + connection_key, + connection_position, + &updates.room_lazy_membership, + now_ms, + ) + .await?; + + txn.commit().await?; + Ok(connection_position) + } +} + +/// `_persist_sliding_sync_connection_lazy_members_txn`: record the lazy +/// members sent, and forget those whose membership changed unsent. +async fn persist_lazy_members( + txn: &Borrowed<'_, '_>, + connection_key: i64, + connection_position: i64, + all_changes: &HashMap, + now_ms: i64, +) -> anyhow::Result<()> { + // New entries, and those whose timestamp is old enough to rewrite. + let mut rooms: Vec<&str> = Vec::new(); + let mut users: Vec<&str> = Vec::new(); + for (room_id, changes) in all_changes { + for user_id in changes.get_returned_user_ids_to_update(now_ms) { + rooms.push(room_id); + users.push(user_id); + } + } + if !rooms.is_empty() { + // An entry of another fork of the connection is left alone. This is + // a cache, so missing a row costs a few extra members down sync. + txn.query( + "INSERT INTO sliding_sync_connection_lazy_members + (connection_key, connection_position, room_id, user_id, last_seen_ts) + SELECT $1, $2, room_id, user_id, $3 + FROM UNNEST($4::text[], $5::text[]) AS u(room_id, user_id) + ON CONFLICT (connection_key, room_id, user_id) + DO UPDATE SET last_seen_ts = EXCLUDED.last_seen_ts + WHERE sliding_sync_connection_lazy_members.connection_position IS NULL + OR sliding_sync_connection_lazy_members.connection_position = EXCLUDED.connection_position", + &[&connection_key, &connection_position, &now_ms, &rooms, &users], + ) + .await?; + } + + // Removed from every fork: expiring an entry of a cache is always safe. + let mut rooms: Vec<&str> = Vec::new(); + let mut users: Vec<&str> = Vec::new(); + for (room_id, changes) in all_changes { + for user_id in &changes.invalidated_user_ids { + rooms.push(room_id); + users.push(user_id); + } + } + if !rooms.is_empty() { + txn.query( + "DELETE FROM sliding_sync_connection_lazy_members + WHERE connection_key = $1 + AND (room_id, user_id) IN (SELECT * FROM UNNEST($2::text[], $3::text[]))", + &[&connection_key, &rooms, &users], + ) + .await?; + } + Ok(()) +} + +// Test-only entry points, so trial can compare the Rust store against the +// Python one on a real schema. + +/// A [`RoomForUser`] as tests read it: `(room_id, sender, membership, +/// event_id, room_version_id, instance_name, stream, has_known_state, +/// room_type, is_encrypted)`. +type PySnapshotRoom = ( + String, + Option, + String, + Option, + Option, + String, + i64, + bool, + Option, + bool, +); + +fn opt_string(value: &Option>) -> Option { + value.as_deref().map(str::to_owned) +} + +fn snapshot_to_py(room: &RoomForUser) -> PySnapshotRoom { + ( + room.room_id.to_string(), + opt_string(&room.sender), + room.membership.to_string(), + opt_string(&room.event_id), + opt_string(&room.room_version_id), + room.event_pos.instance_name.to_string(), + room.event_pos.stream, + room.has_known_state, + opt_string(&room.room_type), + room.is_encrypted, + ) +} + +fn snapshots_to_py(rooms: &SnapshotRooms) -> HashMap { + rooms + .iter() + .map(|(room_id, room)| (room_id.to_string(), snapshot_to_py(room))) + .collect() +} + +fn position_to_py(position: &PersistedEventPosition) -> (String, i64) { + (position.instance_name.to_string(), position.stream) +} + +/// A status as tests write it: `(flag, token)`. +type PyStatus = (String, Option); + +fn status_to_py(status: &HaveSentRoom) -> PyStatus { + let (flag, token) = status_to_row(status); + (flag.to_owned(), token.map(str::to_owned)) +} + +fn status_from_py((flag, token): PyStatus) -> PyResult> { + status_from_row(&flag, token).map_err(anyhow_to_py) +} + +fn statuses_from_py( + statuses: HashMap, +) -> PyResult>> { + statuses + .into_iter() + .map(|(room_id, status)| Ok((room_id, status_from_py(status)?))) + .collect() +} + +fn statuses_to_py(statuses: &HashMap>) -> HashMap { + statuses + .iter() + .map(|(room_id, status)| (room_id.clone(), status_to_py(status))) + .collect() +} + +/// A room config as tests write it: `(timeline_limit, {type: [state_key]})`. +type PyRoomConfig = (i64, BTreeMap>); + +fn config_to_py(config: &RoomSyncConfig) -> PyRoomConfig { + ( + config.timeline_limit, + config + .required_state_map + .iter() + .map(|(event_type, state_keys)| { + (event_type.clone(), state_keys.iter().cloned().collect()) + }) + .collect(), + ) +} + +fn config_from_py((timeline_limit, map): PyRoomConfig) -> RoomSyncConfig { + RoomSyncConfig { + timeline_limit, + required_state_map: map + .into_iter() + .map(|(event_type, state_keys)| (event_type, state_keys.into_iter().collect())) + .collect(), + } +} + +/// A stored connection state as tests read it: `(last_used_ts, rooms, +/// receipts, account_data, room_configs)`. +type PyStoredState = ( + Option, + HashMap, + HashMap, + HashMap, + HashMap, +); + +/// A room's lazy membership changes as tests write them: +/// `({user_id: last_seen_ts}, invalidated_user_ids)`. +type PyLazyChanges = (HashMap>, HashSet); + +/// Raise [`UnknownPosition`] as Python's `SlidingSyncUnknownPosition`. +fn store_error_to_py(py: Python<'_>, e: anyhow::Error) -> PyErr { + if e.downcast_ref::().is_some() { + return 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, + }; + } + anyhow_to_py(e) +} + +fn arcs(ids: Vec) -> Vec> { + ids.into_iter().map(Arc::from).collect() +} + +fn arc_keys(map: HashMap) -> HashMap, V> { + map.into_iter().map(|(k, v)| (Arc::from(k), v)).collect() +} + +/// The sliding sync store's queries, each run to completion on the tokio +/// runtime with the GIL released, over the native driver's pool and the +/// Rust sliding sync handler's caches. For tests only. +#[pyclass(name = "SlidingSyncStoreForTest")] +pub struct PySlidingSyncStoreForTest { + store: Arc, + pool: ConnectionPool, + handle: Handle, +} + +impl PySlidingSyncStoreForTest { + /// Run `f` over a reader for one request, blocking until it is done. + fn run<'a, T, F>(&'a self, py: Python<'_>, f: F) -> PyResult + where + T: Send + 'a, + F: FnOnce(SlidingSyncReader<'a>) -> BoxFuture<'a, anyhow::Result> + Send, + { + let result = py.detach(|| { + self.handle + .block_on(f(self.store.sliding_sync_reader(self.pool.clone()))) + }); + result.map_err(|e| store_error_to_py(py, e)) + } +} + +#[pymethods] +impl PySlidingSyncStoreForTest { + #[new] + fn new(pool: &Bound<'_, PyConnectionPool>, handler: &SlidingSyncHandler) -> Self { + let pool = pool.get(); + PySlidingSyncStoreForTest { + store: Arc::clone(handler.store()), + pool: pool.pool.clone(), + handle: pool.handle.clone(), + } + } + + fn rooms_from_snapshots( + &self, + py: Python<'_>, + user_id: String, + ) -> PyResult> { + let rooms = self.run(py, |reader| { + async move { reader.rooms_from_snapshots(&user_id).await }.boxed() + })?; + Ok(snapshots_to_py(&rooms)) + } + + fn self_leave_rooms_after( + &self, + py: Python<'_>, + user_id: String, + to_token: RoomStreamToken, + ) -> PyResult> { + let rooms = self.run(py, |reader| { + async move { reader.self_leave_rooms_after(&user_id, &to_token).await }.boxed() + })?; + Ok(snapshots_to_py(&rooms)) + } + + fn snapshot_room( + &self, + py: Python<'_>, + user_id: String, + room_id: String, + ) -> PyResult> { + let room = self.run(py, |reader| { + async move { reader.snapshot_room(&user_id, &room_id).await }.boxed() + })?; + Ok(room.as_ref().map(snapshot_to_py)) + } + + fn snapshot_rooms( + &self, + py: Python<'_>, + user_id: String, + room_ids: Vec, + ) -> PyResult> { + let room_ids = arcs(room_ids); + let rooms = self.run(py, |reader| { + async move { reader.snapshot_rooms(&user_id, &room_ids).await }.boxed() + })?; + Ok(snapshots_to_py(&rooms)) + } + + /// The changes in order, as `(room_id, sender, membership, event_id, + /// (instance_name, stream), room_version_id)`. + #[allow(clippy::type_complexity)] + fn snapshot_membership_changes( + &self, + py: Python<'_>, + user_id: String, + from_token: RoomStreamToken, + to_token: RoomStreamToken, + excluded_room_ids: Vec, + membership_may_have_changed: bool, + ) -> PyResult< + Vec<( + String, + Option, + String, + Option, + (String, i64), + Option, + )>, + > { + let excluded: HashSet> = arcs(excluded_room_ids).into_iter().collect(); + let changes = self.run(py, |reader| { + async move { + reader + .snapshot_membership_changes( + &user_id, + &from_token, + &to_token, + &excluded, + membership_may_have_changed, + ) + .await + } + .boxed() + })?; + Ok(changes + .values() + .map(|change| { + ( + change.room_id.to_string(), + opt_string(&change.sender), + change.membership.to_string(), + opt_string(&change.event_id), + position_to_py(&change.event_pos), + opt_string(&change.room_version_id), + ) + }) + .collect()) + } + + /// The changes in order, as `(room_id, event_id, (instance_name, + /// stream), membership, sender, prev_event_id, prev_event_pos, + /// prev_membership, prev_sender)`. + #[allow(clippy::type_complexity)] + fn membership_delta_changes( + &self, + py: Python<'_>, + user_id: String, + from_token: RoomStreamToken, + to_token: RoomStreamToken, + excluded_room_ids: Vec, + membership_may_have_changed: bool, + ) -> PyResult< + Vec<( + String, + Option, + (String, i64), + String, + Option, + Option, + Option<(String, i64)>, + Option, + Option, + )>, + > { + let excluded: HashSet> = arcs(excluded_room_ids).into_iter().collect(); + let changes = self.run(py, |reader| { + async move { + reader + .membership_delta_changes( + &user_id, + &from_token, + &to_token, + &excluded, + membership_may_have_changed, + ) + .await + } + .boxed() + })?; + Ok(changes + .iter() + .map(|change| { + ( + change.room_id.to_string(), + opt_string(&change.event_id), + position_to_py(&change.event_pos), + change.membership.to_string(), + opt_string(&change.sender), + opt_string(&change.prev_event_id), + change.prev_event_pos.as_ref().map(position_to_py), + opt_string(&change.prev_membership), + opt_string(&change.prev_sender), + ) + }) + .collect()) + } + + #[pyo3(signature = (room_ids, from_token, changed_since))] + fn rooms_with_updates_since( + &self, + py: Python<'_>, + room_ids: Vec, + from_token: RoomStreamToken, + changed_since: Option>, + ) -> PyResult> { + let room_ids = arcs(room_ids); + let changed: Option>> = + changed_since.map(|changed| changed.into_iter().map(Arc::from).collect()); + let rooms = self.run(py, |reader| { + async move { + reader + .rooms_with_updates_since(&room_ids, &from_token, changed.as_ref()) + .await + } + .boxed() + })?; + Ok(rooms.iter().map(|id| id.to_string()).collect()) + } + + fn max_event_pos( + &self, + py: Python<'_>, + room_ids: Vec, + now_token: RoomStreamToken, + last_change_pos: HashMap, + ) -> PyResult>> { + let room_ids = arcs(room_ids); + let last_change_pos = arc_keys(last_change_pos); + let positions = self.run(py, |reader| { + async move { + reader + .max_event_pos(&room_ids, &now_token, &last_change_pos) + .await + } + .boxed() + })?; + Ok(positions + .into_iter() + .map(|(room_id, stream)| (room_id.to_string(), stream)) + .collect()) + } + + fn last_event_pos_before( + &self, + py: Python<'_>, + room_ids: Vec, + end_token: RoomStreamToken, + now_token: RoomStreamToken, + last_change_pos: HashMap, + ) -> PyResult> { + let room_ids = arcs(room_ids); + let last_change_pos = arc_keys(last_change_pos); + let positions = self.run(py, |reader| { + async move { + reader + .last_event_pos_before(&room_ids, &end_token, &now_token, &last_change_pos) + .await + } + .boxed() + })?; + Ok(positions + .into_iter() + .map(|(room_id, stream)| (room_id.to_string(), stream)) + .collect()) + } + + fn cached_member_counts( + &self, + py: Python<'_>, + room_id: String, + ) -> PyResult> { + let counts = self.run(py, |reader| { + async move { reader.cached_member_counts(&room_id).await }.boxed() + })?; + Ok((*counts).clone()) + } + + /// `(known, room_type)`: `known` is false for `ROOM_UNKNOWN_SENTINEL`. + fn room_type(&self, py: Python<'_>, room_id: String) -> PyResult<(bool, Option)> { + let room_type = self.run(py, |reader| { + async move { reader.room_type(&room_id).await }.boxed() + })?; + Ok(match room_type { + RoomType::Unknown => (false, None), + RoomType::Known(room_type) => (true, opt_string(&room_type)), + }) + } + + fn room_version_id(&self, py: Python<'_>, room_id: String) -> PyResult> { + let version = self.run(py, |reader| { + async move { reader.room_version_id(&room_id).await }.boxed() + })?; + Ok(opt_string(&version)) + } + + fn partial_rooms(&self, py: Python<'_>) -> PyResult> { + let rooms = self.run(py, |reader| { + async move { reader.partial_rooms().await }.boxed() + })?; + Ok(rooms.iter().map(|id| id.to_string()).collect()) + } + + fn tags_for_user<'py>( + &self, + py: Python<'py>, + user_id: String, + ) -> PyResult>>> { + let tags = self.run(py, |reader| { + async move { reader.tags_for_user(&user_id).await }.boxed() + })?; + tags.iter() + .map(|(room_id, tags)| { + let tags = tags + .iter() + .map(|(tag, content)| { + Ok(( + tag.clone(), + crate::storage::sync_store::json_to_py(py, content)?, + )) + }) + .collect::>()?; + Ok((room_id.to_string(), tags)) + }) + .collect() + } + + fn is_encrypted_at( + &self, + py: Python<'_>, + room_id: String, + to_token: RoomStreamToken, + now_token: RoomStreamToken, + ) -> PyResult { + self.run(py, |reader| { + async move { + reader + .is_encrypted_at(&room_id, &to_token, &now_token) + .await + } + .boxed() + }) + } + + fn connection_lazy_members( + &self, + py: Python<'_>, + connection_position: i64, + room_id: String, + user_ids: Vec, + ) -> PyResult> { + let user_ids = arcs(user_ids); + let members = self.run(py, |reader| { + async move { + reader + .connection_lazy_members(connection_position, &room_id, &user_ids) + .await + } + .boxed() + })?; + Ok(members + .into_iter() + .map(|(user_id, ts)| (user_id.to_string(), ts)) + .collect()) + } + + fn load_connection_state( + &self, + py: Python<'_>, + user_id: String, + device_id: String, + conn_id: String, + connection_position: i64, + now_ms: i64, + ) -> PyResult { + let state = self.run(py, |reader| { + async move { + reader + .load_connection_state( + &user_id, + &device_id, + &conn_id, + connection_position, + now_ms, + ) + .await + } + .boxed() + })?; + Ok(( + state.last_used_ts, + statuses_to_py(&state.rooms), + statuses_to_py(&state.receipts), + statuses_to_py(&state.account_data), + state + .room_configs + .iter() + .map(|(room_id, config)| (room_id.clone(), config_to_py(config))) + .collect(), + )) + } + + #[pyo3(signature = (user_id, device_id, conn_id, previous_position, *, rooms, receipts, account_data, room_configs, room_lazy_membership, now_ms))] + #[allow(clippy::too_many_arguments)] + fn persist_connection_state( + &self, + py: Python<'_>, + user_id: String, + device_id: String, + conn_id: String, + previous_position: Option, + rooms: HashMap, + receipts: HashMap, + account_data: HashMap, + room_configs: HashMap, + room_lazy_membership: HashMap, + now_ms: i64, + ) -> PyResult { + let updates = ConnectionStateUpdates { + rooms: statuses_from_py(rooms)?, + receipts: statuses_from_py(receipts)?, + account_data: statuses_from_py(account_data)?, + room_configs: room_configs + .into_iter() + .map(|(room_id, config)| (room_id, config_from_py(config))) + .collect(), + room_lazy_membership: room_lazy_membership + .into_iter() + .map(|(room_id, (returned, invalidated))| { + ( + room_id, + RoomLazyMembershipChanges { + returned_user_id_to_last_seen_ts_map: returned, + invalidated_user_ids: invalidated, + }, + ) + }) + .collect(), + }; + self.run(py, |reader| { + async move { + reader + .persist_connection_state( + &user_id, + &device_id, + &conn_id, + previous_position, + &updates, + now_ms, + ) + .await + } + .boxed() + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn required_state_round_trips_as_python_encodes_it() { + let mut map = RequiredStateMap::new(); + map.entry("m.room.member".to_owned()) + .or_default() + .extend(["$ME".to_owned(), "@é:test".to_owned()]); + map.entry("m.room.create".to_owned()) + .or_default() + .insert(String::new()); + let json = encode_required_state(&map); + // `json.dumps(sorted(...), separators=(",", ":"))`, ASCII only. + assert_eq!( + json, + r#"[["m.room.create",""],["m.room.member","$ME"],["m.room.member","@\u00e9:test"]]"# + ); + assert_eq!(decode_required_state(&json).unwrap(), map); + assert_eq!(encode_required_state(&RequiredStateMap::new()), "[]"); + } + + #[test] + fn statuses_round_trip() { + for status in [ + HaveSentRoom::Never, + HaveSentRoom::Live, + HaveSentRoom::Previously("s5".to_owned()), + ] { + let (flag, token) = status_to_row(&status); + assert_eq!( + status_from_row(flag, token.map(str::to_owned)).unwrap(), + status + ); + } + assert!(status_from_row("previously", None).is_err()); + assert!(status_from_row("sent", None).is_err()); + } +} diff --git a/rust/src/storage/sync_store.rs b/rust/src/storage/sync_store.rs index ca43c406f8..4cebf441b8 100644 --- a/rust/src/storage/sync_store.rs +++ b/rust/src/storage/sync_store.rs @@ -196,6 +196,8 @@ pub struct SyncDbStats { pub checkouts: u64, /// Time spent waiting for a connection. pub checkout_time: Duration, + /// Transactions begun. + pub transactions: u64, /// Statements executed. pub queries: u64, /// Rows returned. @@ -238,7 +240,7 @@ impl<'a> SyncReader<'a> { *lock_ignoring_poison(&self.stats) } - fn update_stats(&self, f: impl FnOnce(&mut SyncDbStats)) { + pub(crate) fn update_stats(&self, f: impl FnOnce(&mut SyncDbStats)) { f(&mut lock_ignoring_poison(&self.stats)); } @@ -278,6 +280,25 @@ impl<'a> SyncReader<'a> { borrowed .execute("BEGIN ISOLATION LEVEL REPEATABLE READ READ ONLY") .await?; + self.update_stats(|stats| stats.transactions += 1); + Ok(borrowed) + } + + /// Borrow a connection and open a `REPEATABLE READ` transaction on it + /// that may write, as a Python `runInteraction` does. Unlike + /// [`SyncReader::begin`]'s, such a transaction can fail with a + /// serialization error or a deadlock, which the caller retries. Finish + /// with [`Borrowed::commit`]. + pub(crate) async fn begin_write(&self) -> anyhow::Result> { + let mut borrowed = self.borrow().await?; + // Set before awaiting: if we are dropped after BEGIN is sent but + // before it returns, the transaction is open and must be rolled back. + // A ROLLBACK outside one only warns. + borrowed.in_transaction = true; + borrowed + .execute("BEGIN ISOLATION LEVEL REPEATABLE READ") + .await?; + self.update_stats(|stats| stats.transactions += 1); Ok(borrowed) } @@ -382,39 +403,54 @@ impl<'a> SyncReader<'a> { } /// `get_references_for_events`: the `(event_id, sender)` of each event - /// that references one of `event_ids`, oldest first. + /// that references one of `event_ids`, oldest first. Events with no + /// references are left out. pub async fn references_for_events( &self, event_ids: &[Arc], - ) -> anyhow::Result, Vec>> { - let mut references: HashMap, Vec> = HashMap::new(); - if event_ids.is_empty() { - return Ok(references); - } - let borrowed = self.borrow().await?; - for batch in event_ids.chunks(SELECT_MANY_BATCH_SIZE) { - let ids: Vec<&str> = batch.iter().map(|id| &**id).collect(); - let rows = borrowed - .query( - "SELECT relates_to_id, ref.event_id, ref.sender - FROM events AS ref - INNER JOIN event_relations USING (event_id) - INNER JOIN events AS parent ON - parent.event_id = relates_to_id - AND parent.room_id = ref.room_id - WHERE relates_to_id = ANY($1) AND relation_type = 'm.reference' - ORDER BY ref.topological_ordering, ref.stream_ordering", - &[&ids], - ) - .await?; - for row in rows { - let parent: &str = row.try_get(0)?; - references.entry(Arc::from(parent)).or_default().push(( - Arc::from(row.try_get::<_, &str>(1)?), - Arc::from(row.try_get::<_, &str>(2)?), - )); - } - } + ) -> anyhow::Result, Arc<[RelatedEvent]>>> { + let mut references = self + .store + .caches + .references + .get_or_load_many(event_ids, |missing| async move { + let mut found: HashMap, Vec> = HashMap::new(); + let borrowed = self.borrow().await?; + for batch in missing.chunks(SELECT_MANY_BATCH_SIZE) { + let ids: Vec<&str> = batch.iter().map(|id| &**id).collect(); + let rows = borrowed + .query( + "SELECT relates_to_id, ref.event_id, ref.sender + FROM events AS ref + INNER JOIN event_relations USING (event_id) + INNER JOIN events AS parent ON + parent.event_id = relates_to_id + AND parent.room_id = ref.room_id + WHERE relates_to_id = ANY($1) AND relation_type = 'm.reference' + ORDER BY ref.topological_ordering, ref.stream_ordering", + &[&ids], + ) + .await?; + for row in rows { + let parent: &str = row.try_get(0)?; + found.entry(Arc::from(parent)).or_default().push(( + Arc::from(row.try_get::<_, &str>(1)?), + Arc::from(row.try_get::<_, &str>(2)?), + )); + } + } + // Python caches the empty list for events with none. + Ok(missing + .into_iter() + .map(|event_id| { + let refs = found.remove(&event_id).unwrap_or_default(); + (event_id, Arc::from(refs)) + }) + .collect()) + }) + .await + .map_err(anyhow::Error::new)?; + references.retain(|_, refs| !refs.is_empty()); Ok(references) } @@ -424,38 +460,53 @@ impl<'a> SyncReader<'a> { &self, event_ids: &[Arc], ) -> anyhow::Result, Arc>> { - let mut edits = HashMap::new(); - if event_ids.is_empty() { - return Ok(edits); - } - let borrowed = self.borrow().await?; - for batch in event_ids.chunks(SELECT_MANY_BATCH_SIZE) { - let ids: Vec<&str> = batch.iter().map(|id| &**id).collect(); - // `DISTINCT ON` picks the first row of each original, so the - // ordering picks the latest edit. - let rows = borrowed - .query( - "SELECT DISTINCT ON (original.event_id) original.event_id, edit.event_id - FROM events AS edit - INNER JOIN event_relations USING (event_id) - INNER JOIN events AS original ON - original.event_id = relates_to_id - AND edit.type = original.type - AND edit.sender = original.sender - AND edit.room_id = original.room_id - WHERE relates_to_id = ANY($1) AND relation_type = 'm.replace' - ORDER BY original.event_id DESC, edit.origin_server_ts DESC, edit.event_id DESC", - &[&ids], - ) - .await?; - for row in rows { - edits.insert( - Arc::from(row.try_get::<_, &str>(0)?), - Arc::from(row.try_get::<_, &str>(1)?), - ); - } - } - Ok(edits) + let edits = self + .store + .caches + .applicable_edits + .get_or_load_many(event_ids, |missing| async move { + let mut found: HashMap, Arc> = HashMap::new(); + let borrowed = self.borrow().await?; + for batch in missing.chunks(SELECT_MANY_BATCH_SIZE) { + let ids: Vec<&str> = batch.iter().map(|id| &**id).collect(); + // `DISTINCT ON` picks the first row of each original, so + // the ordering picks the latest edit. + let rows = borrowed + .query( + "SELECT DISTINCT ON (original.event_id) original.event_id, edit.event_id + FROM events AS edit + INNER JOIN event_relations USING (event_id) + INNER JOIN events AS original ON + original.event_id = relates_to_id + AND edit.type = original.type + AND edit.sender = original.sender + AND edit.room_id = original.room_id + WHERE relates_to_id = ANY($1) AND relation_type = 'm.replace' + ORDER BY original.event_id DESC, edit.origin_server_ts DESC, edit.event_id DESC", + &[&ids], + ) + .await?; + for row in rows { + found.insert( + Arc::from(row.try_get::<_, &str>(0)?), + Arc::from(row.try_get::<_, &str>(1)?), + ); + } + } + Ok(missing + .into_iter() + .map(|event_id| { + let edit = found.remove(&event_id); + (event_id, edit) + }) + .collect()) + }) + .await + .map_err(anyhow::Error::new)?; + Ok(edits + .into_iter() + .filter_map(|(event_id, edit)| Some((event_id, edit?))) + .collect()) } /// The IDs and counts of `get_thread_summaries`: for each of @@ -465,64 +516,55 @@ impl<'a> SyncReader<'a> { &self, event_ids: &[Arc], ) -> anyhow::Result, (Arc, i64)>> { - let mut summaries = HashMap::new(); - if event_ids.is_empty() { - return Ok(summaries); - } - // One snapshot, as Python's one transaction, so that a count never - // misses its latest reply. - let borrowed = self.begin().await?; - for batch in event_ids.chunks(SELECT_MANY_BATCH_SIZE) { - let ids: Vec<&str> = batch.iter().map(|id| &**id).collect(); - let rows = borrowed - .query( - "SELECT DISTINCT ON (parent.event_id) parent.event_id, child.event_id - FROM events AS child - INNER JOIN event_relations USING (event_id) - INNER JOIN events AS parent ON - parent.event_id = relates_to_id - AND parent.room_id = child.room_id - WHERE relates_to_id = ANY($1) AND relation_type = 'm.thread' - ORDER BY parent.event_id, child.topological_ordering DESC, - child.stream_ordering DESC", - &[&ids], - ) - .await?; - let mut latest: HashMap, Arc> = HashMap::new(); - for row in rows { - latest.insert( - Arc::from(row.try_get::<_, &str>(0)?), - Arc::from(row.try_get::<_, &str>(1)?), - ); - } - if latest.is_empty() { - continue; - } - let parents: Vec<&str> = latest.keys().map(|id| &**id).collect(); - let rows = borrowed - .query( - "SELECT parent.event_id, COUNT(child.event_id) - FROM events AS child - INNER JOIN event_relations USING (event_id) - INNER JOIN events AS parent ON - parent.event_id = relates_to_id - AND parent.room_id = child.room_id - WHERE relates_to_id = ANY($1) AND relation_type = 'm.thread' - GROUP BY parent.event_id", - &[&parents], - ) - .await?; - let mut counts: HashMap<&str, i64> = HashMap::new(); - for row in &rows { - counts.insert(row.try_get(0)?, row.try_get(1)?); - } - for (parent, latest_id) in &latest { - let count = counts.get(&**parent).copied().unwrap_or(0); - summaries.insert(Arc::clone(parent), (Arc::clone(latest_id), count)); - } - } - borrowed.commit().await?; - Ok(summaries) + let summaries = self + .store + .caches + .thread_summaries + .get_or_load_many(event_ids, |missing| async move { + let mut found: HashMap, (Arc, i64)> = HashMap::new(); + let borrowed = self.borrow().await?; + for batch in missing.chunks(SELECT_MANY_BATCH_SIZE) { + let ids: Vec<&str> = batch.iter().map(|id| &**id).collect(); + // Python counts in a second statement of the same + // transaction. The window runs before `DISTINCT ON`, so + // it counts every reply, and one statement is one + // snapshot. + let rows = borrowed + .query( + "SELECT DISTINCT ON (parent.event_id) parent.event_id, child.event_id, + COUNT(*) OVER (PARTITION BY parent.event_id) + FROM events AS child + INNER JOIN event_relations USING (event_id) + INNER JOIN events AS parent ON + parent.event_id = relates_to_id + AND parent.room_id = child.room_id + WHERE relates_to_id = ANY($1) AND relation_type = 'm.thread' + ORDER BY parent.event_id, child.topological_ordering DESC, + child.stream_ordering DESC", + &[&ids], + ) + .await?; + for row in rows { + found.insert( + Arc::from(row.try_get::<_, &str>(0)?), + (Arc::from(row.try_get::<_, &str>(1)?), row.try_get(2)?), + ); + } + } + Ok(missing + .into_iter() + .map(|event_id| { + let summary = found.remove(&event_id); + (event_id, summary) + }) + .collect()) + }) + .await + .map_err(anyhow::Error::new)?; + Ok(summaries + .into_iter() + .filter_map(|(event_id, summary)| Some((event_id, summary?))) + .collect()) } /// `get_threads_participated`: which of the threads of `event_ids` the @@ -809,46 +851,60 @@ impl<'a> SyncReader<'a> { } /// `get_global_account_data_by_type_for_user`: the user's account data - /// of the type, if any. + /// of the type, if any, through its mirror. pub async fn global_account_data_by_type( &self, user_id: &str, data_type: &str, - ) -> anyhow::Result> { - let rows = self - .query_one_shot( - "SELECT content FROM account_data WHERE user_id = $1 AND account_data_type = $2", - &[&user_id, &data_type], - ) - .await?; - let Some(row) = rows.first() else { - return Ok(None); - }; - let content: Option = row.try_get(0)?; - match content.filter(|content| !content.is_empty()) { - Some(content) => { - Ok(Some(serde_json::from_str(&content).with_context(|| { - format!("account data {data_type} of {user_id}") - })?)) - } - None => Ok(None), - } + ) -> anyhow::Result>> { + self.store + .caches + .global_account_data + .get_or_load((Arc::from(user_id), Arc::from(data_type)), || async { + let rows = self + .query_one_shot( + "SELECT content FROM account_data + WHERE user_id = $1 AND account_data_type = $2", + &[&user_id, &data_type], + ) + .await?; + let Some(row) = rows.first() else { + return Ok(None); + }; + let content: Option = row.try_get(0)?; + match content.filter(|content| !content.is_empty()) { + Some(content) => Ok(Some(Arc::new( + serde_json::from_str(&content) + .with_context(|| format!("account data {data_type} of {user_id}"))?, + ))), + None => Ok(None), + } + }) + .await + .map_err(anyhow::Error::new) } /// `ignored_users`: the users the user ignores. - pub async fn ignored_users(&self, user_id: &str) -> anyhow::Result>> { - let rows = self - .query_one_shot( - "SELECT ignored_user_id FROM ignored_users WHERE ignorer_user_id = $1", - &[&user_id], - ) - .await?; - let mut users = HashSet::with_capacity(rows.len()); - for row in rows { - let user_id: String = row.try_get(0)?; - users.insert(Arc::from(user_id)); - } - Ok(users) + pub async fn ignored_users(&self, user_id: &str) -> anyhow::Result>>> { + self.store + .caches + .ignored_users + .get_or_load(Arc::from(user_id), || async { + let rows = self + .query_one_shot( + "SELECT ignored_user_id FROM ignored_users WHERE ignorer_user_id = $1", + &[&user_id], + ) + .await?; + let mut users = HashSet::with_capacity(rows.len()); + for row in rows { + let user_id: String = row.try_get(0)?; + users.insert(Arc::from(user_id)); + } + Ok(Arc::new(users)) + }) + .await + .map_err(anyhow::Error::new) } /// `are_users_erased`: which of the users have asked for their events @@ -965,7 +1021,7 @@ impl<'a> SyncReader<'a> { } /// [`SyncReader::events`] on a connection already borrowed. - async fn load_events( + pub(crate) async fn load_events( &self, borrowed: &Borrowed<'_, '_>, event_ids: &[Arc], @@ -1271,6 +1327,19 @@ impl<'a> SyncReader<'a> { &self, room_id: &str, end_token: &RoomStreamToken, + ) -> anyhow::Result, PersistedEventPosition)>> { + self.last_event_of_types_before(room_id, end_token, &[]) + .await + } + + /// `get_last_event_pos_in_room_before_stream_ordering` with + /// `event_types`: as [`Self::last_event_before`], counting only events + /// of one of `event_types`. No types means any type, as in Python. + pub async fn last_event_of_types_before( + &self, + room_id: &str, + end_token: &RoomStreamToken, + event_types: &[String], ) -> anyhow::Result, PersistedEventPosition)>> { // The first subquery is the events within the token's instance map, // between its minimum and maximum, to be filtered against the map; @@ -1278,34 +1347,43 @@ impl<'a> SyncReader<'a> { // in stream order is the answer. let min_stream = end_token.stream; let max_stream = end_token.get_max_stream_pos(); - let rows = self - .query_one_shot( - "SELECT * FROM ( - SELECT instance_name, stream_ordering, topological_ordering, event_id - FROM events - LEFT JOIN rejections USING (event_id) - WHERE room_id = $1 - AND $2 < stream_ordering AND stream_ordering <= $3 - AND NOT outlier - AND rejections.event_id IS NULL - ORDER BY stream_ordering DESC - ) AS a - UNION ALL - SELECT * FROM ( - SELECT instance_name, stream_ordering, topological_ordering, event_id - FROM events - LEFT JOIN rejections USING (event_id) - WHERE room_id = $1 - AND stream_ordering <= $2 - AND NOT outlier - AND rejections.event_id IS NULL - ORDER BY stream_ordering DESC - LIMIT 1 - ) AS b - ORDER BY stream_ordering DESC", - &[&room_id, &min_stream, &max_stream], - ) - .await?; + let type_clause = if event_types.is_empty() { + "" + } else { + "AND type = ANY($4)" + }; + let sql = format!( + "SELECT * FROM ( + SELECT instance_name, stream_ordering, topological_ordering, event_id + FROM events + LEFT JOIN rejections USING (event_id) + WHERE room_id = $1 + {type_clause} + AND $2 < stream_ordering AND stream_ordering <= $3 + AND NOT outlier + AND rejections.event_id IS NULL + ORDER BY stream_ordering DESC + ) AS a + UNION ALL + SELECT * FROM ( + SELECT instance_name, stream_ordering, topological_ordering, event_id + FROM events + LEFT JOIN rejections USING (event_id) + WHERE room_id = $1 + {type_clause} + AND stream_ordering <= $2 + AND NOT outlier + AND rejections.event_id IS NULL + ORDER BY stream_ordering DESC + LIMIT 1 + ) AS b + ORDER BY stream_ordering DESC" + ); + let mut params: Vec<&(dyn ToSql + Sync)> = vec![&room_id, &min_stream, &max_stream]; + if !event_types.is_empty() { + params.push(&event_types); + } + let rows = self.query_one_shot(&sql, ¶ms).await?; for row in rows { let instance_name: Option = row.try_get(0)?; let stream_ordering: i64 = row.try_get(1)?; @@ -1544,6 +1622,115 @@ impl<'a> SyncReader<'a> { Ok(deltas) } + /// `get_current_state_deltas_for_room`: the changes to the room's + /// current state after `from_token` and up to `to_token`, in `stream_id` + /// order, bounded on the delta rows' own `stream_id`. Sliding sync reads + /// these. A token that falls inside a persist batch misses that batch's + /// deltas, as the Python explains: [`Self::state_deltas`] is the variant + /// that pairs with a timeline. + /// + /// Python first asks its delta stream cache whether the room changed + /// after `from_token`. Rust always reads the table, which gives the same + /// answer. + pub async fn state_deltas_by_stream_id( + &self, + room_id: &str, + from_token: Option<&RoomStreamToken>, + to_token: Option<&RoomStreamToken>, + ) -> anyhow::Result> { + if let (Some(from), Some(to)) = (from_token, to_token) { + if to.is_before_or_eq(from) { + return Ok(Vec::new()); + } + } + let lower_value = from_token.map(|from| from.stream); + let upper_value = to_token.map(|to| to.get_max_stream_pos()); + let lower_clause = if lower_value.is_some() { + "AND $2 < stream_id" + } else { + "" + }; + let upper_clause = match (lower_value, upper_value) { + (_, None) => "", + (Some(_), Some(_)) => "AND stream_id <= $3", + (None, Some(_)) => "AND stream_id <= $2", + }; + let sql = format!( + "SELECT instance_name, stream_id, type, state_key, event_id, prev_event_id + FROM current_state_delta_stream + WHERE room_id = $1 {lower_clause} {upper_clause} + ORDER BY stream_id ASC" + ); + let mut params: Vec<&(dyn ToSql + Sync)> = vec![&room_id]; + if let Some(value) = &lower_value { + params.push(value); + } + if let Some(value) = &upper_value { + params.push(value); + } + let rows = self.query_one_shot(&sql, ¶ms).await?; + + let mut interner = Interner::default(); + let mut deltas = Vec::with_capacity(rows.len()); + for row in rows { + let instance_name: Option = row.try_get(0)?; + let stream_id: i64 = row.try_get(1)?; + if !filter_results_by_stream(from_token, to_token, instance_name.as_deref(), stream_id)? + { + continue; + } + let event_type: String = row.try_get(2)?; + let state_key: String = row.try_get(3)?; + let event_id: Option = row.try_get(4)?; + let prev_event_id: Option = row.try_get(5)?; + deltas.push(StateDelta { + stream_id, + event_type: interner.intern(&event_type), + state_key: Arc::from(state_key), + event_id: event_id.map(Arc::from), + prev_event_id: prev_event_id.map(Arc::from), + }); + } + Ok(deltas) + } + + /// `get_member_counts`: the number of members of the room by + /// membership, from its current state. A rejected membership event has + /// no membership and is not counted. + pub async fn member_counts(&self, room_id: &str) -> anyhow::Result> { + let rows = self + .query_one_shot( + "SELECT count(*), membership FROM current_state_events + WHERE type = 'm.room.member' AND room_id = $1 + AND membership IS NOT NULL + GROUP BY membership", + &[&room_id], + ) + .await?; + let mut counts = HashMap::with_capacity(rows.len()); + for row in rows { + let count: i64 = row.try_get(0)?; + let membership: String = row.try_get(1)?; + counts.insert(membership, count); + } + Ok(counts) + } + + /// `get_latest_bump_stamp_for_room`: the room's `bump_stamp` in + /// `sliding_sync_joined_rooms`, or `None` if it has none or no row. + pub async fn latest_bump_stamp(&self, room_id: &str) -> anyhow::Result> { + let rows = self + .query_one_shot( + "SELECT bump_stamp FROM sliding_sync_joined_rooms WHERE room_id = $1", + &[&room_id], + ) + .await?; + match rows.first() { + Some(row) => Ok(row.try_get(0)?), + None => Ok(None), + } + } + /// `_get_state_group_for_events`: the state group of each event, which /// is the state after it. An event with none, an outlier or an event /// that is not known, is left out, as it is when the Python is read @@ -2010,10 +2197,10 @@ const STATE_GROUP_BATCH_SIZE: usize = 100; /// Interns strings, so the event types of one result share an allocation. #[derive(Default)] -struct Interner(HashSet>); +pub(crate) struct Interner(HashSet>); impl Interner { - fn intern(&mut self, s: &str) -> Arc { + pub(crate) fn intern(&mut self, s: &str) -> Arc { if let Some(existing) = self.0.get(s) { return Arc::clone(existing); } @@ -2111,7 +2298,7 @@ impl Borrowed<'_, '_> { } /// Run statements that return nothing. - async fn execute(&self, sql: &str) -> anyhow::Result<()> { + pub(crate) async fn execute(&self, sql: &str) -> anyhow::Result<()> { let start = Instant::now(); self.conn() .batch_execute(sql) @@ -2175,11 +2362,14 @@ impl Querier for Borrowed<'_, '_> { // Test-only entry points, so trial can compare the Rust store against the // Python one on a real schema. -fn anyhow_to_py(e: anyhow::Error) -> PyErr { +pub(crate) fn anyhow_to_py(e: anyhow::Error) -> PyErr { pyo3::exceptions::PyRuntimeError::new_err(format!("{e:#}")) } -fn json_to_py<'py>(py: Python<'py>, value: &serde_json::Value) -> PyResult> { +pub(crate) fn json_to_py<'py>( + py: Python<'py>, + value: &serde_json::Value, +) -> PyResult> { pythonize::pythonize(py, value).map_err(Into::into) } @@ -2364,14 +2554,14 @@ impl PySyncStoreForTest { } .boxed() })?; - data.as_ref().map(|data| json_to_py(py, data)).transpose() + data.as_deref().map(|data| json_to_py(py, data)).transpose() } fn ignored_users(&self, py: Python<'_>, user_id: String) -> PyResult> { self.run(py, |reader| { async move { reader.ignored_users(&user_id).await }.boxed() }) - .map(arc_set) + .map(|users| users.iter().map(|id| id.to_string()).collect()) } fn are_users_erased( @@ -2539,6 +2729,46 @@ impl PySyncStoreForTest { })) } + /// `get_last_event_pos_in_room_before_stream_ordering` with + /// `event_types`, as `(event_id, instance_name, stream_ordering)`. + fn last_event_of_types_before( + &self, + py: Python<'_>, + room_id: String, + end_token: RoomStreamToken, + event_types: Vec, + ) -> PyResult> { + let last = self.run(py, |reader| { + async move { + reader + .last_event_of_types_before(&room_id, &end_token, &event_types) + .await + } + .boxed() + })?; + Ok(last.map(|(event_id, position)| { + ( + event_id.to_string(), + position.instance_name.to_string(), + position.stream, + ) + })) + } + + /// `get_member_counts`. + fn member_counts(&self, py: Python<'_>, room_id: String) -> PyResult> { + self.run(py, |reader| { + async move { reader.member_counts(&room_id).await }.boxed() + }) + } + + /// `get_latest_bump_stamp_for_room`. + fn latest_bump_stamp(&self, py: Python<'_>, room_id: String) -> PyResult> { + self.run(py, |reader| { + async move { reader.latest_bump_stamp(&room_id).await }.boxed() + }) + } + /// `get_unread_event_push_actions_by_room_for_user`, with each timeline's /// counts as `(notify_count, unread_count, highlight_count)`: the main /// timeline's, and the threads' by thread. @@ -2609,6 +2839,39 @@ impl PySyncStoreForTest { .collect()) } + /// `get_current_state_deltas_for_room`, with each delta as + /// `(stream_id, type, state_key, event_id, prev_event_id)`. + #[pyo3(signature = (room_id, from_token, to_token))] + #[allow(clippy::type_complexity)] + fn state_deltas_by_stream_id( + &self, + py: Python<'_>, + room_id: String, + from_token: Option, + to_token: Option, + ) -> PyResult, Option)>> { + let deltas = self.run(py, |reader| { + async move { + reader + .state_deltas_by_stream_id(&room_id, from_token.as_ref(), to_token.as_ref()) + .await + } + .boxed() + })?; + Ok(deltas + .into_iter() + .map(|delta| { + ( + delta.stream_id, + delta.event_type.to_string(), + delta.state_key.to_string(), + delta.event_id.map(|id| id.to_string()), + delta.prev_event_id.map(|id| id.to_string()), + ) + }) + .collect()) + } + /// `_get_state_group_for_events`. fn state_group_for_events( &self, diff --git a/rust/src/sync/client_events.rs b/rust/src/sync/client_events.rs index 7f8826c0bb..7a99d57d42 100644 --- a/rust/src/sync/client_events.rs +++ b/rust/src/sync/client_events.rs @@ -47,9 +47,10 @@ pub struct ClientFilterContext { pub msc4354_enabled: bool, } -/// `filter_and_transform_events_for_client` with `filter_send_to_client` -/// and not peeking: the events the user may see, in order, each pruned if -/// its sender was erased and with the user's membership at it. +/// `filter_and_transform_events_for_client` with `filter_send_to_client`: +/// the events the user may see, in order, each pruned if its sender was +/// erased and with the user's membership at it. `is_peeking` is for a user +/// who is not in the room. /// /// `now_ms` is the time the sticky and retention checks compare with. pub async fn filter_events_for_client( @@ -57,6 +58,7 @@ pub async fn filter_events_for_client( context: &ClientFilterContext, events: &[Arc], always_include_ids: &HashSet>, + is_peeking: bool, now_ms: i64, ) -> anyhow::Result> { if events.is_empty() { @@ -165,7 +167,7 @@ pub async fn filter_events_for_client( now_ms, msc4354_enabled: context.msc4354_enabled, filter_send_to_client: true, - is_peeking: false, + is_peeking, always_include_ids, sender_ignored: context.ignored_users.contains(event.sender()), retention_max_lifetime: retention[event.room_id()], diff --git a/rust/src/sync/relations.rs b/rust/src/sync/relations.rs index b26299163f..d48a75c2e7 100644 --- a/rust/src/sync/relations.rs +++ b/rust/src/sync/relations.rs @@ -211,9 +211,9 @@ pub async fn bundled_aggregations( let references = reader.references_for_events(&event_ids).await?; for (event_id, refs) in references { let refs: Vec> = refs - .into_iter() + .iter() .filter(|(_, sender)| !ignored_users.contains(sender)) - .map(|(id, _)| id) + .map(|(id, _)| Arc::clone(id)) .collect(); if !refs.is_empty() { bundles.entry(event_id).or_default().references = Some(refs); diff --git a/rust/src/sync/room.rs b/rust/src/sync/room.rs index 07d38e01e0..ab8acd58f5 100644 --- a/rust/src/sync/room.rs +++ b/rust/src/sync/room.rs @@ -297,6 +297,7 @@ pub async fn generate_room_entry( &context.client_filter, &events, &HashSet::new(), + false, context.now_ms(), ) .await?; diff --git a/rust/src/sync/rooms.rs b/rust/src/sync/rooms.rs index 5f268c85fd..2fe6d16644 100644 --- a/rust/src/sync/rooms.rs +++ b/rust/src/sync/rooms.rs @@ -651,7 +651,7 @@ async fn room_entry( } /// The message of a panic's payload. -fn panic_message(panic: &(dyn std::any::Any + Send)) -> &str { +pub(crate) fn panic_message(panic: &(dyn std::any::Any + Send)) -> &str { if let Some(message) = panic.downcast_ref::<&str>() { message } else if let Some(message) = panic.downcast_ref::() { diff --git a/rust/src/sync/timeline.rs b/rust/src/sync/timeline.rs index 961bd0cd03..95cd585ac8 100644 --- a/rust/src/sync/timeline.rs +++ b/rust/src/sync/timeline.rs @@ -103,6 +103,7 @@ async fn filter_recents( &context.client_filter, &recents, ¤t_state_ids, + false, context.now_ms(), ) .await diff --git a/scripts-dev/complement.sh b/scripts-dev/complement.sh index ccb0764ad0..81787ed5e6 100755 --- a/scripts-dev/complement.sh +++ b/scripts-dev/complement.sh @@ -386,6 +386,11 @@ main() { fi export PASS_SYNAPSE_RUST_PERSIST="$RUST_PERSIST" fi + if [[ -n "$RUST_SLIDING_SYNC" ]]; then + # Set `experimental_features.rust_sliding_sync` (off or on). The Rust + # sliding sync handler needs RUST_DRIVER=1. + export PASS_SYNAPSE_RUST_SLIDING_SYNC="$RUST_SLIDING_SYNC" + fi if [[ -n "$UNIX_SOCKETS" ]]; then # Enable full on Unix socket mode for Synapse, Redis and Postgresql diff --git a/synapse/config/experimental.py b/synapse/config/experimental.py index caecdc6967..e138aed4ea 100644 --- a/synapse/config/experimental.py +++ b/synapse/config/experimental.py @@ -381,6 +381,24 @@ class ExperimentalConfig(Config): ) self.rust_sync: Literal["off", "shadow", "on"] = rust_sync + # Whether to compute simplified sliding sync (MSC4186) responses in + # Rust. Needs the native Rust Postgres driver. + # - "off": the Python implementation. + # - "on": Rust, falling back to Python per room or per request where + # Rust cannot compute it. + # There is no shadow mode. + rust_sliding_sync = experimental.get("rust_sliding_sync", "off") + if rust_sliding_sync is True: + rust_sliding_sync = "on" + elif rust_sliding_sync is False: + rust_sliding_sync = "off" + if rust_sliding_sync not in ("off", "on"): + raise ConfigError( + "experimental_features.rust_sliding_sync must be 'off' or 'on'", + ("experimental", "rust_sliding_sync"), + ) + self.rust_sliding_sync: Literal["off", "on"] = rust_sliding_sync + # How much of outbound federation runs in Rust. One ordered level, each # implying the ones before it: # - "off": all Python. diff --git a/synapse/handlers/sliding_sync/__init__.py b/synapse/handlers/sliding_sync/__init__.py index afb93e42b4..944bf9bdbb 100644 --- a/synapse/handlers/sliding_sync/__init__.py +++ b/synapse/handlers/sliding_sync/__init__.py @@ -15,16 +15,20 @@ import itertools import logging from itertools import chain -from typing import TYPE_CHECKING, AbstractSet, Mapping +from typing import TYPE_CHECKING, AbstractSet, Iterable, Mapping, Sequence import attr -from prometheus_client import Histogram +from prometheus_client import Counter, Histogram 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, @@ -45,6 +49,15 @@ 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 ( + ConnectionState as RustConnectionState, + InterestedRooms as RustInterestedRooms, + RustSlidingSyncFallback, + SlidingSyncHandler as RustSlidingSyncHandler, + SlidingSyncRequest as RustSlidingSyncRequest, + SlidingSyncStats as RustSlidingSyncStats, ) from synapse.types import ( JsonDict, @@ -61,10 +74,14 @@ from synapse.types import ( ) from synapse.types.handlers import SLIDING_SYNC_DEFAULT_BUMP_EVENT_TYPES from synapse.types.handlers.sliding_sync import ( + AnySlidingSyncRoomResult, + HaveSentRoom, HaveSentRoomFlag, MutablePerConnectionState, + OperationType, PerConnectionState, RoomLazyMembershipChanges, + RoomStatusMap, RoomSyncConfig, SlidingSyncConfig, SlidingSyncResult, @@ -86,6 +103,96 @@ sync_processing_time = Histogram( labelnames=["initial", SERVER_NAME_LABEL], ) +# Room entries computed in `rust_sliding_sync: on` mode. `result` is "rust", +# "fallback" for a room left to Python, or "error" for a room Python computed +# after the Rust handler failed on it. +rust_sliding_sync_rooms_counter = Counter( + "synapse_sliding_sync_rust_rooms_total", + "Room entries of sliding sync responses computed in `rust_sliding_sync: on` " + "mode, by whether Rust computed them or fell back to Python", + labelnames=["result", SERVER_NAME_LABEL], +) + +# Requests served in `rust_sliding_sync: on` mode. `result` is "rust" for a +# 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 " + "served them or Python computed the whole response", + labelnames=["result", SERVER_NAME_LABEL], +) + +# What the Rust implementation did per request, from its `SlidingSyncStats`. +rust_sliding_sync_db_time = Histogram( + "synapse_sliding_sync_rust_db_seconds", + "Time a Rust sliding sync request spent waiting on database statements", + labelnames=[SERVER_NAME_LABEL], +) +rust_sliding_sync_checkout_time = Histogram( + "synapse_sliding_sync_rust_db_checkout_seconds", + "Time a Rust sliding sync request spent waiting for database connections", + labelnames=[SERVER_NAME_LABEL], +) +_RUST_COUNT_BUCKETS = (0, 1, 2, 5, 10, 20, 50, 100, 200, 500, 1000) +rust_sliding_sync_transactions = Histogram( + "synapse_sliding_sync_rust_db_transactions", + "Database transactions begun per Rust sliding sync request", + labelnames=[SERVER_NAME_LABEL], + buckets=_RUST_COUNT_BUCKETS, +) +rust_sliding_sync_queries = Histogram( + "synapse_sliding_sync_rust_db_queries", + "Database statements executed per Rust sliding sync request", + labelnames=[SERVER_NAME_LABEL], + buckets=_RUST_COUNT_BUCKETS, +) +rust_sliding_sync_rooms_computed = Histogram( + "synapse_sliding_sync_rust_rooms_computed", + "Room entries the Rust implementation computed per sliding sync request", + labelnames=[SERVER_NAME_LABEL], + buckets=_RUST_COUNT_BUCKETS, +) +rust_sliding_sync_lazy_member_lookups = Histogram( + "synapse_sliding_sync_rust_lazy_member_lookups", + "Rooms whose lazy-loaded members the Rust implementation looked up in the " + "connection state, per sliding sync request", + labelnames=[SERVER_NAME_LABEL], + buckets=_RUST_COUNT_BUCKETS, +) + + +def observe_rust_sliding_sync_stats( + stats: RustSlidingSyncStats, server_name: str +) -> None: + """Record one request's `SlidingSyncStats` in the histograms.""" + labels = {SERVER_NAME_LABEL: server_name} + rust_sliding_sync_db_time.labels(**labels).observe(stats.db_time) + rust_sliding_sync_checkout_time.labels(**labels).observe(stats.checkout_time) + rust_sliding_sync_transactions.labels(**labels).observe(stats.transactions) + rust_sliding_sync_queries.labels(**labels).observe(stats.queries) + rust_sliding_sync_rooms_computed.labels(**labels).observe(stats.rooms_computed) + rust_sliding_sync_lazy_member_lookups.labels(**labels).observe( + stats.lazy_member_lookups + ) + + +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. # @@ -108,11 +215,81 @@ class SlidingSyncHandler: self.relations_handler = hs.get_relations_handler() self.rooms_to_exclude_globally = hs.config.server.rooms_to_exclude_from_sync self.is_mine_id = hs.is_mine_id + self._hs_config = hs.config self.connection_store = SlidingSyncConnectionStore(self.clock, self.store) self.extensions = SlidingSyncExtensionHandler(hs) self.room_lists = SlidingSyncRoomLists(hs) + # The Rust sliding sync handler, used when the `rust_sliding_sync` + # experimental option is "on". It needs the native Rust Postgres driver. + self._rust_sliding_sync_mode = hs.config.experimental.rust_sliding_sync + self._rust_sliding_sync: RustSlidingSyncHandler | None = None + if self._rust_sliding_sync_mode != "off": + self._rust_sliding_sync = hs.get_rust_handlers().sliding_sync + if self._rust_sliding_sync is None: + logger.warning( + "rust_sliding_sync is %r but the main database does not use " + "the native Rust driver: using the Python sliding sync handler", + self._rust_sliding_sync_mode, + ) + + # Whether a Rust error raises rather than falling back to Python. Tests + # set it. + self._rust_sliding_sync_strict = False + + @property + def parses_request_body(self) -> bool: + """Whether the servlet should keep the request body's bytes in the + `SlidingSyncConfig`, for the Rust implementation to parse.""" + return self._rust_sliding_sync is not None + + async def _build_rust_request( + self, + sync_config: SlidingSyncConfig, + from_token: SlidingSyncStreamToken | None, + to_token: StreamToken, + ) -> RustSlidingSyncRequest: + """What the Rust implementation reads about the request: the body, and + what Python knows about the user, the tokens and the server. + + Raises: + ValueError if Rust cannot parse the body. Python validated it, so + this is a body Rust does not cover (see `rust/src/sliding_sync/ + request.rs`), and Python computes the request. + """ + user_id = sync_config.user.to_string() + body = sync_config.body_bytes + if body is None: + # Not from a request. Only unset fields can hold `Absent`. + body = sync_config.model_dump_json( + by_alias=True, + exclude_unset=True, + include={"conn_id", "lists", "room_subscriptions", "extensions"}, + ).encode() + is_server_admin = await self.store.is_server_admin(user_id) + client_config = await self.store.get_admin_client_config_for_user(user_id) + # A room key never holds the separator. + to_suffix = (await to_token.to_string(self.store)).split("_", 1)[1] + config = self._hs_config + return RustSlidingSyncRequest( + body, + user_id=user_id, + device_id=sync_config.requester.device_id, + from_room_key=from_token.stream_token.room_key if from_token else None, + connection_position=from_token.connection_position if from_token else None, + to_room_key=to_token.room_key, + to_suffix=to_suffix, + is_server_admin=is_server_admin, + return_soft_failed_events=client_config.return_soft_failed_events, + return_policy_server_spammy_events=client_config.return_policy_server_spammy_events, + have_finished_bg_jobs=await self.store.have_finished_sliding_sync_background_jobs(), + rooms_to_exclude=set(self.rooms_to_exclude_globally), + msc4306_enabled=config.experimental.msc4306_enabled, + msc4354_enabled=config.experimental.msc4354_enabled, + profiles_enabled=config.server.include_profile_updates_in_sync, + ) + async def wait_for_sync_for_user( self, requester: Requester, @@ -255,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 @@ -292,7 +485,7 @@ class SlidingSyncHandler: relevant_rooms_to_send_map = interested_rooms.relevant_rooms_to_send_map # Fetch room data - rooms: dict[str, SlidingSyncResult.RoomResult] = {} + rooms: dict[str, AnySlidingSyncRoomResult] = {} new_connection_state = previous_connection_state.get_mutable() @@ -321,7 +514,9 @@ class SlidingSyncHandler: if relevant_rooms_to_send_map: with start_active_span("sliding_sync.generate_room_entries"): - await concurrently_execute(handle_room, relevant_rooms_to_send_map, 20) + await concurrently_execute( + handle_room, list(relevant_rooms_to_send_map), 20 + ) extensions = await self.extensions.get_extensions_response( sync_config=sync_config, @@ -415,6 +610,321 @@ class SlidingSyncHandler: return sliding_sync_result + async def _current_sync_for_user_in_rust( + self, + rust: RustSlidingSyncHandler, + sync_config: SlidingSyncConfig, + to_token: StreamToken, + 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: + 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. + """ + 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 request in Python: %s", e) + return None + + # 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: + tokens += [ + from_token.stream_token.room_key, + from_token.stream_token.receipt_key, + ] + instance_ids = { + name: await store.get_id_for_instance(name) + for token in tokens + for name in token.instance_map + } + + stats: list[RustSlidingSyncStats] = [] + try: + 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 it in Python", user_id + ) + return None + + entries = result.take_entries() + for entry in entries: + rooms[entry.room_id] = entry + rust_sliding_sync_rooms_counter.labels(result="rust", **labels).inc( + len(entries) + len(result.empty) + ) + + python_room_ids = [] + for room_id, reason in result.fallbacks: + rust_sliding_sync_rooms_counter.labels(result="fallback", **labels).inc() + logger.debug("Computing %s in Python: %s", room_id, reason) + python_room_ids.append(room_id) + for room_id, error in result.errors: + rust_sliding_sync_rooms_counter.labels(result="error", **labels).inc() + if self._rust_sliding_sync_strict: + raise RuntimeError(f"Rust sliding sync failed for {room_id}: {error}") + # Rust only reads, and its changes to the connection state are + # for the rooms it finished, so Python can compute the room from + # scratch. + logger.error( + "Rust sliding sync failed for %s: %s. Computing it in Python", + room_id, + error, + ) + python_room_ids.append(room_id) + 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, SlidingSyncUnknownPosition): + # A request racing on the same connection cleared our position. + # Python would hit the same error, so there is nothing to retry. + 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) + ) + # Python's `get_entities_changed` does not trust an empty cache. + # Leaving this as None makes Rust ask the database instead. + events_cache = store._events_stream_cache + changed = events_cache.get_all_entities_changed(from_stream) + if changed.hit and events_cache._cache: + 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) + } + + @trace + @tag_args + 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( self, @@ -664,35 +1174,13 @@ class SlidingSyncHandler: # or `limited` mean for clients that interpret them correctly. In future this # behavior is almost certainly going to change. # - from_bound = None - initial = True - ignore_timeline_bound = False - if from_token and not newly_joined and not state_reset_out_of_room: - room_status = previous_connection_state.rooms.have_sent_room(room_id) - if room_status.status == HaveSentRoomFlag.LIVE: - from_bound = from_token.stream_token.room_key - initial = False - elif room_status.status == HaveSentRoomFlag.PREVIOUSLY: - assert room_status.last_token is not None - from_bound = room_status.last_token - initial = False - elif room_status.status == HaveSentRoomFlag.NEVER: - from_bound = None - initial = True - else: - assert_never(room_status.status) - - log_kv({"sliding_sync.room_status": room_status}) - - if prev_room_sync_config is not None: - # Check if the timeline limit has increased, if so ignore the - # timeline bound and record the change (see "XXX: Odd behavior" - # above). - if ( - prev_room_sync_config.timeline_limit - < room_sync_config.timeline_limit - ): - ignore_timeline_bound = True + from_bound, initial, ignore_timeline_bound = _timeline_bounds( + previous_connection_state.rooms.have_sent_room(room_id), + from_token=from_token, + skip_previous_bounds=newly_joined or state_reset_out_of_room, + prev_room_sync_config=prev_room_sync_config, + room_sync_config=room_sync_config, + ) log_kv( { @@ -859,24 +1347,7 @@ class SlidingSyncHandler: invite_or_knock_event = await self.store.get_event( room_membership_for_user_at_to_token.event_id ) - - stripped_state = [] - if invite_or_knock_event.membership == Membership.INVITE: - invite_state = invite_or_knock_event.unsigned.get( - "invite_room_state", [] - ) - if not isinstance(invite_state, list): - invite_state = [] - - stripped_state.extend(invite_state) - elif invite_or_knock_event.membership == Membership.KNOCK: - knock_state = invite_or_knock_event.unsigned.get("knock_room_state", []) - if not isinstance(knock_state, list): - knock_state = [] - - stripped_state.extend(knock_state) - - stripped_state.append(strip_event(invite_or_knock_event)) + stripped_state = _stripped_state(invite_or_knock_event) # Get the changes to current state in the token range from the # `current_state_delta_stream` table. @@ -1031,186 +1502,30 @@ class SlidingSyncHandler: # https://github.com/matrix-org/matrix-spec-proposals/pull/3575#discussion_r1653045932 # # Calculate the `StateFilter` based on the `required_state` for the room - required_state_filter = StateFilter.none() - - # Keep track of which users' state we may need to fetch. We split this - # into explicit users and lazy loaded users. - explicit_user_state = set() - lazy_load_user_ids = set() - - # Whether lazy-loading of room members is enabled. - lazy_load_room_members = False - - if room_membership_for_user_at_to_token.membership not in ( - Membership.INVITE, - Membership.KNOCK, - ): - # If we have a double wildcard ("*", "*") in the `required_state`, we need - # to fetch all state for the room - # - # Note: MSC3575 describes different behavior to how we're handling things - # here but since it's not wrong to return more state than requested - # (`required_state` is just the minimum requested), it doesn't matter if we - # include more than client wanted. This complexity is also under scrutiny, - # see - # https://github.com/matrix-org/matrix-spec-proposals/pull/3575#discussion_r1185109050 - # - # > One unique exception is when you request all state events via ["*", "*"]. When used, - # > all state events are returned by default, and additional entries FILTER OUT the returned set - # > of state events. These additional entries cannot use '*' themselves. - # > For example, ["*", "*"], ["m.room.member", "@alice:example.com"] will _exclude_ every m.room.member - # > event _except_ for @alice:example.com, and include every other state event. - # > In addition, ["*", "*"], ["m.space.child", "*"] is an error, the m.space.child filter is not - # > required as it would have been returned anyway. - # > - # > -- MSC3575 (https://github.com/matrix-org/matrix-spec-proposals/pull/3575) - if StateValues.WILDCARD in room_sync_config.required_state_map.get( - StateValues.WILDCARD, set() - ): - set_tag( - SynapseTags.FUNC_ARG_PREFIX + "required_state_wildcard", - True, - ) - required_state_filter = StateFilter.all() - # TODO: `StateFilter` currently doesn't support wildcard event types. We're - # currently working around this by returning all state to the client but it - # would be nice to fetch less from the database and return just what the - # client wanted. - elif ( - room_sync_config.required_state_map.get(StateValues.WILDCARD) - is not None - ): - set_tag( - SynapseTags.FUNC_ARG_PREFIX + "required_state_wildcard_event_type", - True, - ) - required_state_filter = StateFilter.all() - else: - required_state_types: list[tuple[str, str | None]] = [] - num_wild_state_keys = 0 - num_others = 0 - for ( - state_type, - state_key_set, - ) in room_sync_config.required_state_map.items(): - for state_key in state_key_set: - if state_key == StateValues.WILDCARD: - num_wild_state_keys += 1 - # `None` is a wildcard in the `StateFilter` - required_state_types.append((state_type, None)) - # We need to fetch all relevant people when we're lazy-loading membership - elif ( - state_type == EventTypes.Member - and state_key == StateValues.LAZY - ): - lazy_load_room_members = True - - # Everyone in the timeline is relevant - timeline_membership: set[str] = set() - if timeline_events is not None: - for timeline_event in timeline_events: - # Anyone who sent a message is relevant - timeline_membership.add(timeline_event.event.sender) - - # We also care about invite, ban, kick, targets, - # etc. - if timeline_event.event.type == EventTypes.Member: - timeline_membership.add( - timeline_event.event.state_key - ) - - # The client needs to know the membership of everyone in - # the timeline we're returning. - lazy_load_user_ids.update(timeline_membership) - - # Update the required state filter so we pick up the new - # membership - if limited or initial: - # If the timeline is limited, we only need to - # return the membership changes for people in - # the timeline. - for user_id in timeline_membership: - required_state_types.append( - (EventTypes.Member, user_id) - ) - else: - # For non-limited timelines we always return all - # membership changes. This is so that clients - # who have fetched the full membership list - # already can continue to maintain it for - # non-limited syncs. - # - # This assumes that for non-limited syncs there - # won't be many membership changes that wouldn't - # have been included already (this can only - # happen if membership state was rolled back due - # to state resolution anyway). - # - # `None` is a wildcard in the `StateFilter` - required_state_types.append((EventTypes.Member, None)) - - # Record the extra members we're returning. - lazy_load_user_ids.update( - state_key - for event_type, state_key in room_state_delta_id_map - if event_type == EventTypes.Member - ) - else: - num_others += 1 - - # Replace `$ME` with the user's ID so we can deduplicate - # when someone requests the same state with `$ME` or with - # their user ID. - normalized_state_key = state_key - if state_key == StateValues.ME: - normalized_state_key = user.to_string() - - if state_type == EventTypes.Member: - # Also track explicitly requested member state for - # lazy membership tracking. - explicit_user_state.add(normalized_state_key) - - required_state_types.append( - (state_type, normalized_state_key) - ) - - set_tag( - SynapseTags.FUNC_ARG_PREFIX - + "required_state_wildcard_state_key_count", - num_wild_state_keys, - ) - set_tag( - SynapseTags.FUNC_ARG_PREFIX + "required_state_lazy", - lazy_load_room_members, - ) - set_tag( - SynapseTags.FUNC_ARG_PREFIX + "required_state_other_count", - num_others, - ) - - required_state_filter = StateFilter.from_types(required_state_types) - - # Remove any explicitly requested user state from the lazy-loaded set, - # as we track them separately. - lazy_load_user_ids -= explicit_user_state + required_state_request = _required_state_request( + user.to_string(), + membership=room_membership_for_user_at_to_token.membership, + required_state_map=room_sync_config.required_state_map, + timeline=[timeline_event.event for timeline_event in timeline_events], + limited=limited, + initial=initial, + state_delta_keys=room_state_delta_id_map.keys(), + ) + required_state_filter = required_state_request.required_state_filter + explicit_user_state = required_state_request.explicit_user_state + lazy_load_user_ids = required_state_request.lazy_load_user_ids + lazy_load_room_members = required_state_request.lazy_load_room_members # We need this base set of info for the response so let's just fetch it along # with the `required_state` for the room hero_room_state = [ (EventTypes.Member, hero_user_id) for hero_user_id in hero_user_ids ] - meta_room_state = list(hero_room_state) - if initial or avatar_changed: - meta_room_state.append((EventTypes.RoomAvatar, "")) - - state_filter = StateFilter.all() - if required_state_filter != StateFilter.all(): - state_filter = StateFilter( - types=StateFilter.from_types( - chain(meta_room_state, required_state_filter.to_types()) - ).types, - include_others=required_state_filter.include_others, - ) + state_filter = _state_filter_with_meta( + required_state_filter, + hero_user_ids, + include_avatar=initial or avatar_changed, + ) # The required state map to store in the room sync config, if it has # changed. @@ -1436,66 +1751,22 @@ class SlidingSyncHandler: # sensible order again. bump_stamp = 0 - room_sync_required_state_map_to_persist: Mapping[str, AbstractSet[str]] = ( - room_sync_config.required_state_map - ) - if changed_required_state_map: - room_sync_required_state_map_to_persist = changed_required_state_map - # Record the `room_sync_config` if we're `ignore_timeline_bound` (which means - # that the `timeline_limit` has increased) - unstable_expanded_timeline = False - if ignore_timeline_bound: - # FIXME: We signal the fact that we're sending down more events to - # the client by setting `unstable_expanded_timeline` to true (see - # "XXX: Odd behavior" above). - unstable_expanded_timeline = True - - new_connection_state.room_configs[room_id] = RoomSyncConfig( - timeline_limit=room_sync_config.timeline_limit, - required_state_map=room_sync_required_state_map_to_persist, - ) - elif prev_room_sync_config is not None: - # If the result is `limited` then we need to record that the - # `timeline_limit` has been reduced, as when/if the client later requests - # more timeline then we have more data to send. - # - # Otherwise (when not `limited`) we don't need to record that the - # `timeline_limit` has been reduced, as the *effective* `timeline_limit` - # (i.e. the amount of timeline we have previously sent to the client) is at - # least the previous `timeline_limit`. - # - # This is to handle the case where the `timeline_limit` e.g. goes from 10 to - # 5 to 10 again (without any timeline gaps), where there's no point sending - # down the initial historical chunk events when the `timeline_limit` is - # increased as the client already has the 10 previous events. However, if - # client has a gap in the timeline (i.e. `limited` is True), then we *do* - # need to record the reduced timeline. - # - # TODO: Handle timeline gaps (`get_timeline_gaps()`) - This is separate from - # the gaps we might see on the client because a response was `limited` we're - # talking about above. - if ( - limited - and prev_room_sync_config.timeline_limit - > room_sync_config.timeline_limit - ): - new_connection_state.room_configs[room_id] = RoomSyncConfig( - timeline_limit=room_sync_config.timeline_limit, - required_state_map=room_sync_required_state_map_to_persist, - ) - - elif changed_required_state_map is not None: - new_connection_state.room_configs[room_id] = RoomSyncConfig( - timeline_limit=room_sync_config.timeline_limit, - required_state_map=room_sync_required_state_map_to_persist, - ) - - else: - new_connection_state.room_configs[room_id] = RoomSyncConfig( - timeline_limit=room_sync_config.timeline_limit, - required_state_map=room_sync_required_state_map_to_persist, - ) + # that the `timeline_limit` has increased). + # + # FIXME: We signal the fact that we're sending down more events to the + # client by setting `unstable_expanded_timeline` to true (see "XXX: Odd + # behavior" above). + unstable_expanded_timeline = ignore_timeline_bound + room_config_to_persist = _room_config_to_persist( + room_sync_config, + prev_room_sync_config=prev_room_sync_config, + changed_required_state_map=changed_required_state_map, + ignore_timeline_bound=ignore_timeline_bound, + limited=limited, + ) + if room_config_to_persist is not None: + new_connection_state.room_configs[room_id] = room_config_to_persist set_tag(SynapseTags.RESULT_PREFIX + "initial", initial) @@ -1627,6 +1898,365 @@ class SlidingSyncHandler: return None +def _timeline_bounds( + room_status: HaveSentRoom[RoomStreamToken], + *, + from_token: SlidingSyncStreamToken | None, + skip_previous_bounds: bool, + prev_room_sync_config: RoomSyncConfig | None, + room_sync_config: RoomSyncConfig, +) -> tuple[RoomStreamToken | None, bool, bool]: + """Where a room's updates start, from what the connection sent before. + + Args: + room_status: Whether the connection sent the room before. + from_token: The request's `pos`. + skip_previous_bounds: Whether to send the room as if never sent: when + newly joined or state reset out of the room. + prev_room_sync_config: The room's config the connection stored. + room_sync_config: The room's config in this request. + + Returns: + `(from_bound, initial, ignore_timeline_bound)`: the lower bound of + the room's updates (`None` for the whole room), whether the room is + sent in full, and whether the timeline ignores the bound because the + timeline limit grew. + """ + from_bound = None + initial = True + ignore_timeline_bound = False + if from_token and not skip_previous_bounds: + if room_status.status == HaveSentRoomFlag.LIVE: + from_bound = from_token.stream_token.room_key + initial = False + elif room_status.status == HaveSentRoomFlag.PREVIOUSLY: + assert room_status.last_token is not None + from_bound = room_status.last_token + initial = False + elif room_status.status == HaveSentRoomFlag.NEVER: + from_bound = None + initial = True + else: + assert_never(room_status.status) + + log_kv({"sliding_sync.room_status": room_status}) + + if prev_room_sync_config is not None: + # Check if the timeline limit has increased, if so ignore the + # timeline bound and record the change (see "XXX: Odd behavior" + # in `get_room_sync_data`). + if prev_room_sync_config.timeline_limit < room_sync_config.timeline_limit: + ignore_timeline_bound = True + + return from_bound, initial, ignore_timeline_bound + + +def _stripped_state(invite_or_knock_event: EventBase) -> list[JsonDict]: + """The stripped state of an invite or knock: what the event's `unsigned` + carries, then the event itself, stripped.""" + stripped_state: list[JsonDict] = [] + if invite_or_knock_event.membership == Membership.INVITE: + invite_state = invite_or_knock_event.unsigned.get("invite_room_state", []) + if not isinstance(invite_state, list): + invite_state = [] + + stripped_state.extend(invite_state) + elif invite_or_knock_event.membership == Membership.KNOCK: + knock_state = invite_or_knock_event.unsigned.get("knock_room_state", []) + if not isinstance(knock_state, list): + knock_state = [] + + stripped_state.extend(knock_state) + + stripped_state.append(strip_event(invite_or_knock_event)) + return stripped_state + + +def _room_config_to_persist( + room_sync_config: RoomSyncConfig, + *, + prev_room_sync_config: RoomSyncConfig | None, + changed_required_state_map: Mapping[str, AbstractSet[str]] | None, + ignore_timeline_bound: bool, + limited: bool | None, +) -> RoomSyncConfig | None: + """The room config to record in the connection state after sending the + room, or None to keep the stored one. + + Note: an empty `changed_required_state_map` stores the request's map, + but still counts as a change. + """ + room_sync_required_state_map_to_persist: Mapping[str, AbstractSet[str]] = ( + room_sync_config.required_state_map + ) + if changed_required_state_map: + room_sync_required_state_map_to_persist = changed_required_state_map + + config_to_persist = RoomSyncConfig( + timeline_limit=room_sync_config.timeline_limit, + required_state_map=room_sync_required_state_map_to_persist, + ) + + if ignore_timeline_bound: + return config_to_persist + elif prev_room_sync_config is not None: + # If the result is `limited` then we need to record that the + # `timeline_limit` has been reduced, as when/if the client later requests + # more timeline then we have more data to send. + # + # Otherwise (when not `limited`) we don't need to record that the + # `timeline_limit` has been reduced, as the *effective* `timeline_limit` + # (i.e. the amount of timeline we have previously sent to the client) is at + # least the previous `timeline_limit`. + # + # This is to handle the case where the `timeline_limit` e.g. goes from 10 to + # 5 to 10 again (without any timeline gaps), where there's no point sending + # down the initial historical chunk events when the `timeline_limit` is + # increased as the client already has the 10 previous events. However, if + # client has a gap in the timeline (i.e. `limited` is True), then we *do* + # need to record the reduced timeline. + # + # TODO: Handle timeline gaps (`get_timeline_gaps()`) - This is separate from + # the gaps we might see on the client because a response was `limited` we're + # talking about above. + if ( + limited + and prev_room_sync_config.timeline_limit > room_sync_config.timeline_limit + ): + return config_to_persist + elif changed_required_state_map is not None: + return config_to_persist + return None + else: + return config_to_persist + + +@attr.s(slots=True, frozen=True, auto_attribs=True) +class _RequiredStateRequest: + """What a room's `required_state` asks for, given its timeline. Returned + by `_required_state_request`.""" + + required_state_filter: StateFilter + """The state the request asks for.""" + + lazy_load_room_members: bool + """Whether lazy-loading of room members is enabled.""" + + lazy_load_user_ids: set[str] + """The lazy-loaded members, less the explicitly requested ones.""" + + explicit_user_state: set[str] + """The members requested explicitly, by user ID or `$ME`.""" + + +def _required_state_request( + user_id: str, + *, + membership: str, + required_state_map: Mapping[str, AbstractSet[str]], + timeline: Sequence[EventBase], + limited: bool | None, + initial: bool, + state_delta_keys: Iterable[tuple[str, str]], +) -> _RequiredStateRequest: + """Derive the state filter and member sets of a room's `required_state`. + + Args: + user_id: The user making the request. + membership: The user's membership in the room at the `to_token`. + required_state_map: The room's `RoomSyncConfig.required_state_map`. + timeline: The timeline events being returned. + limited: Whether the timeline is limited, or None with no timeline. + initial: Whether the room is sent down in full. + state_delta_keys: The keys of the room's state deltas in the token + range. + """ + required_state_filter = StateFilter.none() + + # Keep track of which users' state we may need to fetch. We split this + # into explicit users and lazy loaded users. + explicit_user_state = set() + lazy_load_user_ids = set() + + # Whether lazy-loading of room members is enabled. + lazy_load_room_members = False + + if membership not in ( + Membership.INVITE, + Membership.KNOCK, + ): + # If we have a double wildcard ("*", "*") in the `required_state`, we need + # to fetch all state for the room + # + # Note: MSC3575 describes different behavior to how we're handling things + # here but since it's not wrong to return more state than requested + # (`required_state` is just the minimum requested), it doesn't matter if we + # include more than client wanted. This complexity is also under scrutiny, + # see + # https://github.com/matrix-org/matrix-spec-proposals/pull/3575#discussion_r1185109050 + # + # > One unique exception is when you request all state events via ["*", "*"]. When used, + # > all state events are returned by default, and additional entries FILTER OUT the returned set + # > of state events. These additional entries cannot use '*' themselves. + # > For example, ["*", "*"], ["m.room.member", "@alice:example.com"] will _exclude_ every m.room.member + # > event _except_ for @alice:example.com, and include every other state event. + # > In addition, ["*", "*"], ["m.space.child", "*"] is an error, the m.space.child filter is not + # > required as it would have been returned anyway. + # > + # > -- MSC3575 (https://github.com/matrix-org/matrix-spec-proposals/pull/3575) + if StateValues.WILDCARD in required_state_map.get(StateValues.WILDCARD, set()): + set_tag( + SynapseTags.FUNC_ARG_PREFIX + "required_state_wildcard", + True, + ) + required_state_filter = StateFilter.all() + # TODO: `StateFilter` currently doesn't support wildcard event types. We're + # currently working around this by returning all state to the client but it + # would be nice to fetch less from the database and return just what the + # client wanted. + elif required_state_map.get(StateValues.WILDCARD) is not None: + set_tag( + SynapseTags.FUNC_ARG_PREFIX + "required_state_wildcard_event_type", + True, + ) + required_state_filter = StateFilter.all() + else: + required_state_types: list[tuple[str, str | None]] = [] + num_wild_state_keys = 0 + num_others = 0 + for ( + state_type, + state_key_set, + ) in required_state_map.items(): + for state_key in state_key_set: + if state_key == StateValues.WILDCARD: + num_wild_state_keys += 1 + # `None` is a wildcard in the `StateFilter` + required_state_types.append((state_type, None)) + # We need to fetch all relevant people when we're lazy-loading membership + elif ( + state_type == EventTypes.Member + and state_key == StateValues.LAZY + ): + lazy_load_room_members = True + + # Everyone in the timeline is relevant + timeline_membership: set[str] = set() + if timeline is not None: + for timeline_event in timeline: + # Anyone who sent a message is relevant + timeline_membership.add(timeline_event.sender) + + # We also care about invite, ban, kick, targets, + # etc. + if timeline_event.type == EventTypes.Member: + timeline_membership.add(timeline_event.state_key) + + # The client needs to know the membership of everyone in + # the timeline we're returning. + lazy_load_user_ids.update(timeline_membership) + + # Update the required state filter so we pick up the new + # membership + if limited or initial: + # If the timeline is limited, we only need to + # return the membership changes for people in + # the timeline. + for member_user_id in timeline_membership: + required_state_types.append( + (EventTypes.Member, member_user_id) + ) + else: + # For non-limited timelines we always return all + # membership changes. This is so that clients + # who have fetched the full membership list + # already can continue to maintain it for + # non-limited syncs. + # + # This assumes that for non-limited syncs there + # won't be many membership changes that wouldn't + # have been included already (this can only + # happen if membership state was rolled back due + # to state resolution anyway). + # + # `None` is a wildcard in the `StateFilter` + required_state_types.append((EventTypes.Member, None)) + + # Record the extra members we're returning. + lazy_load_user_ids.update( + state_key + for event_type, state_key in state_delta_keys + if event_type == EventTypes.Member + ) + else: + num_others += 1 + + # Replace `$ME` with the user's ID so we can deduplicate + # when someone requests the same state with `$ME` or with + # their user ID. + normalized_state_key = state_key + if state_key == StateValues.ME: + normalized_state_key = user_id + + if state_type == EventTypes.Member: + # Also track explicitly requested member state for + # lazy membership tracking. + explicit_user_state.add(normalized_state_key) + + required_state_types.append((state_type, normalized_state_key)) + + set_tag( + SynapseTags.FUNC_ARG_PREFIX + "required_state_wildcard_state_key_count", + num_wild_state_keys, + ) + set_tag( + SynapseTags.FUNC_ARG_PREFIX + "required_state_lazy", + lazy_load_room_members, + ) + set_tag( + SynapseTags.FUNC_ARG_PREFIX + "required_state_other_count", + num_others, + ) + + required_state_filter = StateFilter.from_types(required_state_types) + + # Remove any explicitly requested user state from the lazy-loaded set, + # as we track them separately. + lazy_load_user_ids -= explicit_user_state + + return _RequiredStateRequest( + required_state_filter=required_state_filter, + lazy_load_room_members=lazy_load_room_members, + lazy_load_user_ids=lazy_load_user_ids, + explicit_user_state=explicit_user_state, + ) + + +def _state_filter_with_meta( + required_state_filter: StateFilter, + hero_user_ids: Iterable[str], + *, + include_avatar: bool, +) -> StateFilter: + """The state to fetch for a room: `required_state_filter`, plus the + heroes' memberships and, if `include_avatar`, the room avatar.""" + meta_room_state = [ + (EventTypes.Member, hero_user_id) for hero_user_id in hero_user_ids + ] + if include_avatar: + meta_room_state.append((EventTypes.RoomAvatar, "")) + + state_filter = StateFilter.all() + if required_state_filter != StateFilter.all(): + state_filter = StateFilter( + types=StateFilter.from_types( + chain(meta_room_state, required_state_filter.to_types()) + ).types, + include_others=required_state_filter.include_others, + ) + return state_filter + + @attr.s(auto_attribs=True) class _RequiredStateChangesReturn: """Return type for _required_state_changes.""" diff --git a/synapse/handlers/sliding_sync/extensions.py b/synapse/handlers/sliding_sync/extensions.py index 0348f96115..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, @@ -52,6 +53,7 @@ from synapse.types import ( UserID, ) from synapse.types.handlers.sliding_sync import ( + AnySlidingSyncRoomResult, HaveSentRoomFlag, MutablePerConnectionState, OperationType, @@ -77,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__) @@ -98,12 +108,12 @@ 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], - actual_room_response_map: Mapping[str, SlidingSyncResult.RoomResult], + actual_room_response_map: Mapping[str, AnySlidingSyncRoomResult], to_token: StreamToken, from_token: SlidingSyncStreamToken | None, ) -> SlidingSyncResult.Extensions: @@ -454,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, @@ -686,11 +696,11 @@ 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, SlidingSyncResult.RoomResult], + actual_room_response_map: Mapping[str, AnySlidingSyncRoomResult], receipts_request: SlidingSyncConfig.Extensions.ReceiptsExtension, to_token: StreamToken, from_token: SlidingSyncStreamToken | None, @@ -815,10 +825,10 @@ class SlidingSyncExtensionHandler: # in the timeline to avoid bloating and blowing up the sync response # as the number of users in the room increases. (this behavior is part of the spec) initial_rooms_and_event_ids = [ - (room_id, event.event.event_id) + (room_id, event_id) for room_id in initial_rooms if room_id in actual_room_response_map - for event in actual_room_response_map[room_id].timeline_events + for event_id in actual_room_response_map[room_id].timeline_event_ids ] initial_receipts = await self.store.get_linearized_receipts_for_events( room_and_event_ids=initial_rooms_and_event_ids, @@ -894,7 +904,7 @@ class SlidingSyncExtensionHandler: sync_config: SlidingSyncConfig, actual_lists: Mapping[str, SlidingSyncResult.SlidingWindowList], actual_room_ids: set[str], - actual_room_response_map: Mapping[str, SlidingSyncResult.RoomResult], + actual_room_response_map: Mapping[str, AnySlidingSyncRoomResult], typing_request: SlidingSyncConfig.Extensions.TypingExtension, to_token: StreamToken, from_token: SlidingSyncStreamToken | None, @@ -1129,7 +1139,7 @@ class SlidingSyncExtensionHandler: user_id: str, actual_room_ids: set[str], sync_config: SlidingSyncConfig, - actual_room_response_map: Mapping[str, SlidingSyncResult.RoomResult], + actual_room_response_map: Mapping[str, AnySlidingSyncRoomResult], actual_lists: Mapping[str, SlidingSyncResult.SlidingWindowList], ) -> tuple[set[str], set[str]]: """ @@ -1186,16 +1196,11 @@ class SlidingSyncExtensionHandler: if room_id not in lazy_rooms: continue # Include users from timeline events - for timeline_event in room_data.timeline_events: - lazy_profile_user_ids.add(timeline_event.event.sender) + lazy_profile_user_ids.update(room_data.timeline_senders) # Include users from required state - for state_event in room_data.required_state: - if state_event.type == EventTypes.Member: - lazy_profile_user_ids.add(state_event.state_key) + lazy_profile_user_ids.update(room_data.required_state_member_keys) # Include heroes - if room_data.heroes: - for hero in room_data.heroes: - lazy_profile_user_ids.add(hero.user_id) + lazy_profile_user_ids.update(room_data.hero_user_ids) non_lazy_rooms = actual_room_ids.difference(lazy_rooms) # If we still have non-lazy rooms, get their members. @@ -1290,7 +1295,7 @@ class SlidingSyncExtensionHandler: actual_room_ids: set[str], to_token: StreamToken, from_token: SlidingSyncStreamToken | None, - actual_room_response_map: Mapping[str, SlidingSyncResult.RoomResult], + actual_room_response_map: Mapping[str, AnySlidingSyncRoomResult], actual_lists: Mapping[str, SlidingSyncResult.SlidingWindowList], ) -> SlidingSyncResult.Extensions.ProfilesExtension | None: """ diff --git a/synapse/handlers/sliding_sync/room_lists.py b/synapse/handlers/sliding_sync/room_lists.py index f4d1f484e7..3e8dfe0d36 100644 --- a/synapse/handlers/sliding_sync/room_lists.py +++ b/synapse/handlers/sliding_sync/room_lists.py @@ -21,6 +21,7 @@ from typing import ( Literal, Mapping, MutableMapping, + Sequence, cast, ) @@ -215,6 +216,57 @@ def filter_membership_for_sync( ) +def _sync_list_ops( + ranges: Sequence[tuple[int, int]], + sorted_room_ids: Sequence[str], + room_sync_config: RoomSyncConfig, + relevant_room_map: dict[str, RoomSyncConfig], +) -> list[SlidingSyncResult.SlidingWindowList.Operation]: + """The `SYNC` operations of a list's `ranges` over its sorted rooms. + Adds each room in a range to `relevant_room_map` with the list's + `room_sync_config`, combined with the room's existing config. + + Note: once a room combines with an existing config, the rooms after it + in the list take the combined config too. + """ + ops: list[SlidingSyncResult.SlidingWindowList.Operation] = [] + for range in ranges: + room_ids_in_list: list[str] = [] + + # We're going to loop through the sorted list of rooms starting + # at the range start index and keep adding rooms until we fill + # up the range or run out of rooms. + # + # Both sides of range are inclusive so we `+ 1` + max_num_rooms = range[1] - range[0] + 1 + for room_id in sorted_room_ids[range[0] :]: + if len(room_ids_in_list) >= max_num_rooms: + break + + # Take the superset of the `RoomSyncConfig` for each room. + # + # Update our `relevant_room_map` with the room we're going + # to display and need to fetch more info about. + existing_room_sync_config = relevant_room_map.get(room_id) + if existing_room_sync_config is not None: + room_sync_config = existing_room_sync_config.combine_room_sync_config( + room_sync_config + ) + + relevant_room_map[room_id] = room_sync_config + + room_ids_in_list.append(room_id) + + ops.append( + SlidingSyncResult.SlidingWindowList.Operation( + op=OperationType.SYNC, + range=range, + room_ids=room_ids_in_list, + ) + ) + return ops + + class SlidingSyncRoomLists: """Handles calculating the room lists from sliding sync requests""" @@ -547,44 +599,12 @@ class SlidingSyncRoomLists: limit=max(range[1] + 1 for range in list_config.ranges), ) - for range in list_config.ranges: - room_ids_in_list: list[str] = [] - - # We're going to loop through the sorted list of rooms starting - # at the range start index and keep adding rooms until we fill - # up the range or run out of rooms. - # - # Both sides of range are inclusive so we `+ 1` - max_num_rooms = range[1] - range[0] + 1 - for room_membership in sorted_room_info[range[0] :]: - room_id = room_membership.room_id - - if len(room_ids_in_list) >= max_num_rooms: - break - - # Take the superset of the `RoomSyncConfig` for each room. - # - # Update our `relevant_room_map` with the room we're going - # to display and need to fetch more info about. - existing_room_sync_config = relevant_room_map.get( - room_id - ) - if existing_room_sync_config is not None: - room_sync_config = existing_room_sync_config.combine_room_sync_config( - room_sync_config - ) - - relevant_room_map[room_id] = room_sync_config - - room_ids_in_list.append(room_id) - - ops.append( - SlidingSyncResult.SlidingWindowList.Operation( - op=OperationType.SYNC, - range=range, - room_ids=room_ids_in_list, - ) - ) + ops = _sync_list_ops( + list_config.ranges, + [room.room_id for room in sorted_room_info], + room_sync_config, + relevant_room_map, + ) lists[list_key] = SlidingSyncResult.SlidingWindowList( count=len(filtered_sync_room_map), @@ -735,44 +755,12 @@ class SlidingSyncRoomLists: ops: list[SlidingSyncResult.SlidingWindowList.Operation] = [] if list_config.ranges: - for range in list_config.ranges: - room_ids_in_list: list[str] = [] - - # We're going to loop through the sorted list of rooms starting - # at the range start index and keep adding rooms until we fill - # up the range or run out of rooms. - # - # Both sides of range are inclusive so we `+ 1` - max_num_rooms = range[1] - range[0] + 1 - for room_membership in sorted_room_info[range[0] :]: - room_id = room_membership.room_id - - if len(room_ids_in_list) >= max_num_rooms: - break - - # Take the superset of the `RoomSyncConfig` for each room. - # - # Update our `relevant_room_map` with the room we're going - # to display and need to fetch more info about. - existing_room_sync_config = relevant_room_map.get( - room_id - ) - if existing_room_sync_config is not None: - room_sync_config = existing_room_sync_config.combine_room_sync_config( - room_sync_config - ) - - relevant_room_map[room_id] = room_sync_config - - room_ids_in_list.append(room_id) - - ops.append( - SlidingSyncResult.SlidingWindowList.Operation( - op=OperationType.SYNC, - range=range, - room_ids=room_ids_in_list, - ) - ) + ops = _sync_list_ops( + list_config.ranges, + [room.room_id for room in sorted_room_info], + room_sync_config, + relevant_room_map, + ) lists[list_key] = SlidingSyncResult.SlidingWindowList( count=len(sorted_room_info), diff --git a/synapse/rest/client/sync.py b/synapse/rest/client/sync.py index 6915b51521..1af7d06486 100644 --- a/synapse/rest/client/sync.py +++ b/synapse/rest/client/sync.py @@ -21,7 +21,7 @@ import itertools import logging from collections import defaultdict -from typing import TYPE_CHECKING, Any, Literal, Mapping, Sequence +from typing import TYPE_CHECKING, Any, Callable, Literal, Mapping, Sequence import attr @@ -52,7 +52,7 @@ from synapse.handlers.sync import ( SyncConfig, SyncResult, ) -from synapse.http.server import HttpServer +from synapse.http.server import HttpServer, respond_with_json_bytes from synapse.http.servlet import ( RestServlet, parse_and_validate_json_object_from_request, @@ -61,8 +61,18 @@ from synapse.http.servlet import ( parse_string, ) from synapse.http.site import SynapseRequest -from synapse.logging.opentracing import log_kv, set_tag, trace_with_opname +from synapse.logging.context import defer_to_thread +from synapse.logging.opentracing import ( + log_kv, + set_tag, + start_active_span, + trace_with_opname, +) from synapse.rest.admin.experimental_features import ExperimentalFeature +from synapse.synapse_rust.sliding_sync import ( + RoomResult as RustSlidingSyncRoomResult, + encode_rooms as encode_rust_sliding_sync_rooms, +) from synapse.synapse_rust.sync import RoomsEncoder as RustRoomsEncoder from synapse.types import ( JsonDict, @@ -71,10 +81,11 @@ from synapse.types import ( SlidingSyncStreamToken, StreamToken, ) +from synapse.types.handlers.sliding_sync import AnySlidingSyncRoomResult from synapse.types.rest.client import SlidingSyncBody from synapse.util.caches.lrucache import LruCache from synapse.util.cancellation import cancellable -from synapse.util.json import json_decoder +from synapse.util.json import json_decoder, json_encoder from ._base import client_patterns, set_timeline_upper_limit @@ -967,8 +978,9 @@ class SlidingSyncRestServlet(RestServlet): self.filtering = hs.get_filtering() self.sliding_sync_handler = hs.get_sliding_sync_handler() self.event_serializer = hs.get_event_client_serializer() + self.reactor = hs.get_reactor() - async def on_POST(self, request: SynapseRequest) -> tuple[int, JsonDict]: + async def on_POST(self, request: SynapseRequest) -> tuple[int, JsonDict] | None: requester = await self.auth.get_user_by_req_experimental_feature( request, allow_guest=True, feature=ExperimentalFeature.MSC3575 ) @@ -990,6 +1002,11 @@ class SlidingSyncRestServlet(RestServlet): # by filter ID. For now, we will just prototype with always passing everything # in. body = parse_and_validate_json_object_from_request(request, SlidingSyncBody) + # The Rust implementation parses the body itself, from the same bytes. + body_bytes = None + if self.sliding_sync_handler.parses_request_body: + request.content.seek(0) # type: ignore[union-attr] + body_bytes = request.content.read() # type: ignore[union-attr] # Tag and log useful data to differentiate requests. set_tag( @@ -1028,6 +1045,7 @@ class SlidingSyncRestServlet(RestServlet): lists=body.lists, room_subscriptions=body.room_subscriptions, extensions=body.extensions, + body_bytes=body_bytes, ) ( @@ -1049,30 +1067,34 @@ class SlidingSyncRestServlet(RestServlet): logger.info("Client has disconnected; not serializing response.") return 200, {} - response_content = await self.encode_response(requester, sliding_sync_results) + response_bytes = await self.encode_response(requester, sliding_sync_results) - return 200, response_content + respond_with_json_bytes(request, 200, response_bytes, send_cors=True) + return None @trace_with_opname("sliding_sync.encode_response") async def encode_response( self, requester: Requester, sliding_sync_result: SlidingSyncResult, - ) -> JsonDict: - response: JsonDict = defaultdict(dict) + ) -> bytes: + """The response body. The JSON encoding runs on a thread, as + `respond_with_json` does.""" + response: JsonDict = {} response["pos"] = await sliding_sync_result.next_pos.to_string(self.store) serialized_lists = self.encode_lists(sliding_sync_result.lists) if serialized_lists: response["lists"] = serialized_lists - response["rooms"] = await self.encode_rooms( - requester, sliding_sync_result.rooms - ) + rooms = await self.encode_rooms(requester, sliding_sync_result.rooms) response["extensions"] = await self.encode_extensions( requester, sliding_sync_result.extensions, sliding_sync_result.rooms ) - return response + with start_active_span("encode_json_response"): + return await defer_to_thread( + self.reactor, _encode_sliding_sync_response, response, rooms + ) @trace_with_opname("sliding_sync.encode_lists") def encode_lists( @@ -1100,8 +1122,8 @@ class SlidingSyncRestServlet(RestServlet): async def encode_rooms( self, requester: Requester, - rooms: dict[str, SlidingSyncResult.RoomResult], - ) -> JsonDict: + rooms: Mapping[str, AnySlidingSyncRoomResult], + ) -> "_EncodedRooms": time_now = self.clock.time_msec() serialize_options = await self.event_serializer.create_config( @@ -1109,8 +1131,35 @@ class SlidingSyncRestServlet(RestServlet): requester=requester, ) + # The rooms the Rust handler computed are encoded in Rust. + rust_rooms = [ + room + for room in rooms.values() + if isinstance(room, RustSlidingSyncRoomResult) + ] + encode_in_rust = None + if rust_rooms: + # Module callbacks take Python events, so only then do the rooms' + # events become Python objects. + unsigned_additions: Mapping[str, JsonDict] = {} + if self.event_serializer.has_unsigned_callbacks: + unsigned_additions = await self.event_serializer.unsigned_additions( + event for room in rust_rooms for event in room.events_to_serialize() + ) + + def encode_in_rust() -> bytes: + return encode_rust_sliding_sync_rooms( + rust_rooms, + time_now, + serialize_options, + unsigned_additions=unsigned_additions, + ) + serialized_rooms: dict[str, JsonDict] = {} for room_id, room_result in rooms.items(): + if isinstance(room_result, RustSlidingSyncRoomResult): + continue + serialized_rooms[room_id] = { "notification_count": room_result.notification_count, "highlight_count": room_result.highlight_count, @@ -1215,14 +1264,14 @@ class SlidingSyncRestServlet(RestServlet): # https://github.com/matrix-org/matrix-spec-proposals/pull/3575#discussion_r1117629919 serialized_rooms[room_id]["invite_state"] = room_result.stripped_state - return serialized_rooms + return _EncodedRooms(python=serialized_rooms, encode_in_rust=encode_in_rust) @trace_with_opname("sliding_sync.encode_extensions") async def encode_extensions( self, requester: Requester, extensions: SlidingSyncResult.Extensions, - ref_rooms_results: Mapping[str, SlidingSyncResult.RoomResult], + ref_rooms_results: Mapping[str, AnySlidingSyncRoomResult], ) -> JsonDict: """ Args: @@ -1337,7 +1386,7 @@ class SlidingSyncRestServlet(RestServlet): self, requester: Requester, sticky_events: SlidingSyncResult.Extensions.StickyEventsExtension, - ref_rooms_results: Mapping[str, SlidingSyncResult.RoomResult], + ref_rooms_results: Mapping[str, AnySlidingSyncRoomResult], ) -> JsonDict: """ Serialise the sticky events extension response. @@ -1379,9 +1428,7 @@ class SlidingSyncRestServlet(RestServlet): # Nothing to deduplicate sticky_events_to_write = possibly_duplicated_sticky_events else: - sent_event_ids_in_room_section = { - ev.event.event_id for ev in room_result.timeline_events - } + sent_event_ids_in_room_section = set(room_result.timeline_event_ids) sticky_events_to_write = [ ev for ev in possibly_duplicated_sticky_events @@ -1432,6 +1479,33 @@ def _serialise_thread_subscriptions( return out +@attr.s(slots=True, frozen=True, auto_attribs=True) +class _EncodedRooms: + """The rooms section of a sliding sync response.""" + + python: JsonDict + """The rooms Python computed, encoded as dicts, by room ID.""" + encode_in_rust: Callable[[], bytes] | None + """Encodes the rooms Rust computed as JSON object members, or `None` when + there are none.""" + + +def _encode_sliding_sync_response(response: JsonDict, rooms: _EncodedRooms) -> bytes: + """Encodes `response` with `rooms` as its rooms section. Rust's rooms are + spliced in as bytes, so they never become Python objects.""" + rooms_json = json_encoder.encode(rooms.python).encode("utf-8") + if rooms.encode_in_rust is not None: + rust_json = rooms.encode_in_rust() + if rooms.python: + rooms_json = b"{" + rust_json + b"," + rooms_json[1:] + else: + rooms_json = b"{" + rust_json + b"}" + + # `response` always holds `pos`, so its object is never empty. + response_json = json_encoder.encode(response).encode("utf-8") + return response_json[:-1] + b',"rooms":' + rooms_json + b"}" + + def register_servlets(hs: "HomeServer", http_server: HttpServer) -> None: SyncRestServlet(hs).register(http_server) diff --git a/synapse/synapse_rust/handlers.pyi b/synapse/synapse_rust/handlers.pyi index d07b3bafd8..8ad7e1a758 100644 --- a/synapse/synapse_rust/handlers.pyi +++ b/synapse/synapse_rust/handlers.pyi @@ -15,6 +15,7 @@ from typing import TYPE_CHECKING, Optional from twisted.internet.defer import Deferred from synapse.synapse_rust.persist import Persister +from synapse.synapse_rust.sliding_sync import SlidingSyncHandler from synapse.synapse_rust.state import StateResolutionHandler from synapse.synapse_rust.sync import SyncHandler from synapse.types import JsonDict @@ -51,3 +52,8 @@ class RustHandlers: """The Rust event persister, or None if `rust_event_persistence` is "off", the main database does not use the native Rust Postgres driver, or this instance does not write events.""" + + @property + def sliding_sync(self) -> SlidingSyncHandler | None: + """The Rust sliding sync handler, or None if the main database does not + use the native Rust Postgres driver.""" diff --git a/synapse/synapse_rust/sliding_sync.pyi b/synapse/synapse_rust/sliding_sync.pyi new file mode 100644 index 0000000000..2606c1b599 --- /dev/null +++ b/synapse/synapse_rust/sliding_sync.pyi @@ -0,0 +1,622 @@ +# 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: +# . + +"""Simplified sliding sync (MSC4186) in Rust. See `rust/src/sliding_sync/`.""" + +from collections.abc import Mapping +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, + 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, + state: ConnectionState, + interested: InterestedRooms, + ) -> Deferred[RoomsResult]: + """`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, + state: ConnectionState, + interested: InterestedRooms, + ) -> RoomsResult: ... + def finish_request( + self, + 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 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 + extensions read of a `SlidingSyncResult.RoomResult`.""" + + def __bool__(self) -> bool: ... + @property + def room_id(self) -> str: ... + @property + def initial(self) -> bool: ... + @property + def unstable_expanded_timeline(self) -> bool: ... + @property + def timeline_event_ids(self) -> list[str]: ... + @property + def timeline_senders(self) -> list[str]: ... + @property + def required_state_member_keys(self) -> list[str]: ... + @property + def hero_user_ids(self) -> list[str]: ... + def events_to_serialize(self) -> list[EventBase]: + """Every event the room's encoding serializes, for module callbacks.""" + +class RoomsResult: + """What `generate_room_entries` made of the rooms.""" + + def take_entries(self) -> list[RoomResult]: + """The entries to send. A second call returns an empty list.""" + @property + def empty(self) -> list[str]: + """The rooms computed with nothing to send.""" + @property + def fallbacks(self) -> list[tuple[str, str]]: ... + @property + def errors(self) -> list[tuple[str, str]]: ... + @property + def stats(self) -> SlidingSyncStats: ... + +def encode_rooms( + rooms: list[RoomResult], + time_now_ms: int, + config: SerializeEventConfig, + *, + unsigned_additions: Mapping[str, JsonDict] | None = None, +) -> bytes: + """The servlet's `encode_rooms` for rooms computed in Rust, as JSON object + members (`"room_id":{...}` joined by commas, no braces). Encodes without + the GIL.""" + +class SlidingSyncStats: + """What the Rust implementation did for one request.""" + + def __init__( + self, + *, + db_time: float = 0.0, + checkout_time: float = 0.0, + transactions: int = 0, + queries: int = 0, + rooms_computed: int = 0, + lazy_member_lookups: int = 0, + ) -> None: ... + @property + def db_time(self) -> float: + """Time spent waiting on statements, in seconds.""" + @property + def checkout_time(self) -> float: + """Time spent waiting for a database connection, in seconds.""" + @property + def transactions(self) -> int: ... + @property + def queries(self) -> int: ... + @property + def rooms_computed(self) -> int: ... + @property + def lazy_member_lookups(self) -> int: ... + +class SlidingSyncRequest: + """One sliding sync request as the Rust implementation sees it: the body, + parsed from the bytes pydantic validated, plus what Python adds.""" + + def __init__( + self, + body: bytes, + *, + user_id: str, + device_id: str | None, + from_room_key: RoomStreamToken | None, + connection_position: int | None, + to_room_key: RoomStreamToken, + to_suffix: str, + is_server_admin: bool, + return_soft_failed_events: bool, + return_policy_server_spammy_events: bool, + have_finished_bg_jobs: bool, + rooms_to_exclude: set[str], + msc4306_enabled: bool, + msc4354_enabled: bool, + profiles_enabled: bool, + ) -> None: + """Raises `ValueError` if the body does not parse.""" + def body_for_test(self) -> JsonDict: + """The parsed body with every default filled in, and the lists and + room subscriptions as lists of `[name, value]` pairs.""" + @property + def user_id(self) -> str: ... + @property + def device_id(self) -> str | None: ... + @property + def from_room_key(self) -> RoomStreamToken | None: ... + @property + def connection_position(self) -> int | None: ... + @property + def to_room_key(self) -> RoomStreamToken: ... + @property + def to_suffix(self) -> str: ... + @property + def is_server_admin(self) -> bool: ... + @property + def have_finished_bg_jobs(self) -> bool: ... + @property + def rooms_to_exclude(self) -> set[str]: ... + +# Test exports of the pure logic, for `tests/handlers/test_sliding_sync_rust_pure.py`. +# A `RoomSyncConfig` crosses as `(timeline_limit, {type: state keys})`. + +def room_sync_config_from_room_config_for_test( + timeline_limit: int, required_state: list[tuple[str, str]] +) -> tuple[int, dict[str, list[str]]]: ... +def combine_room_sync_config_for_test( + a: tuple[int, dict[str, list[str]]], b: tuple[int, dict[str, list[str]]] +) -> tuple[int, dict[str, list[str]]]: ... +def must_await_full_state_for_test( + config: tuple[int, dict[str, list[str]]], server_name: str +) -> bool: + """Raises `ValueError` where Python raises `IndexError`.""" + +def required_state_changes_for_test( + user_id: str, + prev_required_state_map: dict[str, list[str]], + request_required_state_map: dict[str, list[str]], + previously_returned_lazy_user_ids: set[str], + request_lazy_load_user_ids: set[str], + state_deltas: list[tuple[str, str]], +) -> tuple[ + dict[str, list[str]] | None, + tuple[dict[str, list[str] | None], bool], + list[str], + list[str], +]: + """`(changed_required_state_map, added_state_filter, + extra_users_to_add_to_lazy_cache, lazy_members_invalidated)`, with a + state filter as `({type: state keys or None}, include_others)`.""" + +def required_state_request_for_test( + user_id: str, + membership: str, + required_state_map: dict[str, list[str]], + timeline: list[tuple[str, str, str | None]], + limited: bool | None, + initial: bool, + state_delta_keys: list[tuple[str, str]], +) -> tuple[tuple[dict[str, list[str] | None], bool], bool, list[str], list[str]]: + """`(required_state_filter, lazy_load_room_members, lazy_load_user_ids, + explicit_user_state)`, the timeline as `(sender, type, state_key)`.""" + +def state_filter_with_meta_for_test( + required_state_filter: StateFilter, hero_user_ids: list[str], include_avatar: bool +) -> tuple[dict[str, list[str] | None], bool]: ... +def room_status_map_ops_for_test( + statuses: dict[str, tuple[str, int | None]], + ops: list[tuple[str, list[str], int | None]], +) -> tuple[dict[str, tuple[str, int | None]], dict[str, tuple[str, int | None]], int]: + """`(updates, flattened, len)` of a `MutableRoomStatusMap[int]` after + `ops` of `("sent", room_ids, None)` or `("unsent", room_ids, token)`.""" + +def lazy_membership_changes_for_test( + returned_user_id_to_last_seen_ts_map: dict[str, int | None], + invalidated_user_ids: set[str], + now_ms: int, +) -> tuple[list[str], bool]: + """`(get_returned_user_ids_to_update, has_updates)`.""" + +def per_connection_state_for_test( + rooms: dict[str, tuple[str, int | None]], + rooms_sent: list[str], + receipts: dict[str, tuple[str, int | None]], + receipts_sent: list[str], + account_data: dict[str, tuple[str, int | None]], + account_data_sent: list[str], + room_configs: dict[str, tuple[int, dict[str, list[str]]]], + room_config_updates: dict[str, tuple[int, dict[str, list[str]]]], + lazy: dict[str, tuple[dict[str, int | None], set[str]]], + now_ms: int, +) -> tuple[bool, int, int]: + """`(has_updates, len(state), len(mutable state))`.""" + +# A room for user in the test form `(room_id, sender, membership, event_id, +# stream, has_known_state, room_type, is_encrypted)`. +_PyRoom = tuple[str, str | None, str, str | None, int, bool, str | None, bool] + +def filter_membership_for_sync_for_test( + user_id: str, room: _PyRoom, newly_left: bool +) -> bool: ... +def filter_rooms_for_test( + rooms: list[_PyRoom], + filters: str, + dm_room_ids: set[str], + room_tags: dict[str, set[str]], + sent_before: set[str], +) -> list[str]: + """`filters` is JSON. Raises `NotImplementedError` as Python does.""" + +def sort_rooms_for_test( + rooms: list[_PyRoom], + earliest_cache_position: int, + cache: dict[str, int], + to_stream: int, + limit: int | None, + recheck: dict[str, int | None], + joined_positions: dict[str, int], +) -> list[str]: + """Raises `KeyError` as Python does.""" + +def sync_list_ops_for_test( + ranges: list[tuple[int, int]], + sorted_room_ids: list[str], + room_sync_config: tuple[int, dict[str, list[str]]], + relevant_room_map: list[tuple[str, tuple[int, dict[str, list[str]]]]], +) -> tuple[ + list[tuple[tuple[int, int], list[str]]], + list[tuple[str, tuple[int, dict[str, list[str]]]]], +]: ... +def rooms_to_send_for_test( + relevant_room_map: list[tuple[str, tuple[int, dict[str, list[str]]]]], + room_configs: dict[str, tuple[int, dict[str, list[str]]]], + statuses: dict[str, tuple[str, int | None]], + last_used_ts: int | None, + rooms_that_have_updates: set[str], + now_ms: int, +) -> list[str]: + """Raises `ValueError` where Python raises `SlidingSyncUnknownPosition`.""" + +def timeline_bounds_for_test( + room_status: tuple[str, int | None], + from_stream: int | None, + skip_previous_bounds: bool, + prev_room_sync_config: tuple[int, dict[str, list[str]]] | None, + room_sync_config: tuple[int, dict[str, list[str]]], +) -> tuple[int | None, bool, bool]: + """`(from_bound, initial, ignore_timeline_bound)`, tokens as streams.""" + +def stripped_state_for_test(event: str) -> str: + """The event and the stripped state as JSON.""" + +def get_bump_stamp_for_test( + timeline: list[tuple[str, int]], + check_outside_timeline: bool, + latest_room_bump_stamp: int | None, + have_finished_bg_jobs: bool, + to_stream: int, + last_event_stream: int | None, +) -> int | None: ... +def room_config_to_persist_for_test( + room_sync_config: tuple[int, dict[str, list[str]]], + prev_room_sync_config: tuple[int, dict[str, list[str]]] | None, + changed_required_state_map: dict[str, list[str]] | None, + ignore_timeline_bound: bool, + limited: bool | None, +) -> tuple[int, dict[str, list[str]]] | None: ... +def find_relevant_room_ids_for_test( + requested_lists: list[str] | None, + requested_room_ids: list[str] | None, + actual_lists: list[tuple[str, list[list[str]]]], + actual_room_ids: set[str], +) -> list[str]: + """Each list as its operations' room IDs. Raises `AssertionError` as + Python does.""" + +# A membership snapshot as `(room_id, sender, membership, event_id, +# room_version_id, instance_name, stream, has_known_state, room_type, +# is_encrypted)`. +_PySnapshot = tuple[ + str, str | None, str, str | None, str | None, str, int, bool, str | None, bool +] +# A room status as `(flag, token)`, and a room config as `(timeline_limit, +# {type: state keys})`. +_PyStatus = tuple[str, str | None] +_PyConfig = tuple[int, dict[str, list[str]]] + +class SlidingSyncStoreForTest: + """The sliding sync store's queries (`rust/src/storage/sliding_sync_store.rs`), + each run to completion on the tokio runtime, over the native driver's pool + and the Rust store's caches. For `tests/storage/test_sliding_sync_store.py`. + Raises `SlidingSyncUnknownPosition` as the Python does.""" + + def __init__(self, pool: ConnectionPool, handler: SlidingSyncHandler) -> None: ... + def rooms_from_snapshots(self, user_id: str) -> dict[str, _PySnapshot]: ... + def self_leave_rooms_after( + self, user_id: str, to_token: RoomStreamToken + ) -> dict[str, _PySnapshot]: ... + def snapshot_room(self, user_id: str, room_id: str) -> _PySnapshot | None: ... + def snapshot_rooms( + self, user_id: str, room_ids: list[str] + ) -> dict[str, _PySnapshot]: ... + def snapshot_membership_changes( + self, + user_id: str, + from_token: RoomStreamToken, + to_token: RoomStreamToken, + excluded_room_ids: list[str], + membership_may_have_changed: bool, + ) -> list[tuple[str, str | None, str, str | None, tuple[str, int], str | None]]: ... + def membership_delta_changes( + self, + user_id: str, + from_token: RoomStreamToken, + to_token: RoomStreamToken, + excluded_room_ids: list[str], + membership_may_have_changed: bool, + ) -> list[ + tuple[ + str, + str | None, + tuple[str, int], + str, + str | None, + str | None, + tuple[str, int] | None, + str | None, + str | None, + ] + ]: ... + def rooms_with_updates_since( + self, + room_ids: list[str], + from_token: RoomStreamToken, + changed_since: set[str] | None, + ) -> set[str]: ... + def max_event_pos( + self, + room_ids: list[str], + now_token: RoomStreamToken, + last_change_pos: dict[str, int], + ) -> dict[str, int | None]: ... + def last_event_pos_before( + self, + room_ids: list[str], + end_token: RoomStreamToken, + now_token: RoomStreamToken, + last_change_pos: dict[str, int], + ) -> dict[str, int]: ... + def cached_member_counts(self, room_id: str) -> dict[str, int]: ... + def room_type(self, room_id: str) -> tuple[bool, str | None]: + """`(known, room_type)`: not known is `ROOM_UNKNOWN_SENTINEL`.""" + def room_version_id(self, room_id: str) -> str | None: ... + def partial_rooms(self) -> set[str]: ... + def tags_for_user(self, user_id: str) -> dict[str, dict[str, JsonMapping]]: ... + def is_encrypted_at( + self, room_id: str, to_token: RoomStreamToken, now_token: RoomStreamToken + ) -> bool: ... + def connection_lazy_members( + self, connection_position: int, room_id: str, user_ids: list[str] + ) -> dict[str, int]: ... + def load_connection_state( + self, + user_id: str, + device_id: str, + conn_id: str, + connection_position: int, + now_ms: int, + ) -> tuple[ + int | None, + dict[str, _PyStatus], + dict[str, _PyStatus], + dict[str, _PyStatus], + dict[str, _PyConfig], + ]: + """`(last_used_ts, rooms, receipts, account_data, room_configs)`.""" + def persist_connection_state( + self, + user_id: str, + device_id: str, + conn_id: str, + previous_position: int | None, + *, + rooms: dict[str, _PyStatus], + receipts: dict[str, _PyStatus], + account_data: dict[str, _PyStatus], + room_configs: dict[str, _PyConfig], + 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/synapse/synapse_rust/sync.pyi b/synapse/synapse_rust/sync.pyi index 01c03fa402..0f7cb32586 100644 --- a/synapse/synapse_rust/sync.pyi +++ b/synapse/synapse_rust/sync.pyi @@ -473,6 +473,18 @@ class SyncStoreForTest: """`get_last_event_pos_in_room_before_stream_ordering`, as `(event_id, instance_name, stream_ordering)`.""" + def last_event_of_types_before( + self, room_id: str, end_token: RoomStreamToken, event_types: list[str] + ) -> tuple[str, str, int] | None: + """`get_last_event_pos_in_room_before_stream_ordering` with + `event_types`, as `(event_id, instance_name, stream_ordering)`.""" + + def member_counts(self, room_id: str) -> dict[str, int]: + """`get_member_counts`.""" + + def latest_bump_stamp(self, room_id: str) -> int | None: + """`get_latest_bump_stamp_for_room`.""" + def unread_counts( self, room_id: str, user_id: str ) -> tuple[tuple[int, int, int], dict[str, tuple[int, int, int]]]: @@ -494,6 +506,15 @@ class SyncStoreForTest: """`get_current_state_deltas_for_room_by_event_position`, with each delta as `(stream_id, type, state_key, event_id, prev_event_id)`.""" + def state_deltas_by_stream_id( + self, + room_id: str, + from_token: RoomStreamToken | None, + to_token: RoomStreamToken | None, + ) -> list[tuple[int, str, str, str | None, str | None]]: + """`get_current_state_deltas_for_room`, with each delta as + `(stream_id, type, state_key, event_id, prev_event_id)`.""" + def state_group_for_events(self, event_ids: list[str]) -> dict[str, int]: """`_get_state_group_for_events`.""" diff --git a/synapse/types/handlers/sliding_sync.py b/synapse/types/handlers/sliding_sync.py index f6d30460b5..b139dce376 100644 --- a/synapse/types/handlers/sliding_sync.py +++ b/synapse/types/handlers/sliding_sync.py @@ -25,6 +25,7 @@ from typing import ( Mapping, MutableMapping, Sequence, + TypeAlias, TypeVar, cast, ) @@ -35,6 +36,7 @@ from pydantic import ConfigDict from synapse.api.constants import EventTypes from synapse.events import EventBase from synapse.events.utils import FilteredEvent +from synapse.synapse_rust.sliding_sync import RoomResult as RustRoomResult from synapse.types import ( DeviceListUpdates, JsonDict, @@ -73,6 +75,9 @@ class SlidingSyncConfig(SlidingSyncBody): user: UserID requester: Requester + # The request body as the client sent it, which the Rust implementation + # parses itself. `None` when the config was not built from a request. + body_bytes: bytes | None = None model_config = ConfigDict( extra="ignore", frozen=True, @@ -202,6 +207,31 @@ class SlidingSyncResult: notification_count: int highlight_count: int + @property + def timeline_event_ids(self) -> list[str]: + """The IDs of the timeline's events. The extensions read this, + which a room computed in Rust has too.""" + return [event.event.event_id for event in self.timeline_events] + + @property + def timeline_senders(self) -> list[str]: + """The senders of the timeline's events.""" + return [event.event.sender for event in self.timeline_events] + + @property + def required_state_member_keys(self) -> list[str]: + """The state keys of the membership events in `required_state`.""" + return [ + event.state_key + for event in self.required_state + if event.type == EventTypes.Member + ] + + @property + def hero_user_ids(self) -> list[str]: + """The user IDs of the heroes.""" + return [hero.user_id for hero in self.heroes or ()] + def __bool__(self) -> bool: """Are there any updates that should be returned immediately to the client? @@ -480,7 +510,8 @@ class SlidingSyncResult: next_pos: SlidingSyncStreamToken lists: Mapping[str, SlidingWindowList] - rooms: dict[str, RoomResult] + # A room computed in Rust is a `synapse_rust.sliding_sync.RoomResult`. + rooms: dict[str, "AnySlidingSyncRoomResult"] extensions: Extensions def __bool__(self) -> bool: @@ -510,6 +541,11 @@ class SlidingSyncResult: ) +# A room's entry in a response: computed in Python, or in Rust, which has the +# properties of the Python one that the extensions read. +AnySlidingSyncRoomResult: TypeAlias = SlidingSyncResult.RoomResult | RustRoomResult + + class StateValues: """ Understood values of the (type, state_key) tuple in `required_state`. diff --git a/tests/config/test_experimental.py b/tests/config/test_experimental.py index b6188859ee..c9a4045ef8 100644 --- a/tests/config/test_experimental.py +++ b/tests/config/test_experimental.py @@ -74,6 +74,25 @@ class ExperimentalConfigTestCase(unittest.TestCase): with self.assertRaises(ConfigError): _read_config({"experimental_features": {"rust_state_res": "maybe"}}) + @parameterized.expand( + [ + # YAML reads these as booleans. + ("rust_sliding_sync: on", "on"), + ("rust_sliding_sync: off", "off"), + ("rust_sliding_sync: 'on'", "on"), + ("{}", "off"), + ] + ) + def test_rust_sliding_sync(self, features_yaml: str, expected: str) -> None: + config = _read_config({"experimental_features": yaml.safe_load(features_yaml)}) + self.assertEqual(config.rust_sliding_sync, expected) + + @parameterized.expand([("shadow",), ("maybe",)]) + def test_rust_sliding_sync_rejects_unknown_modes(self, mode: str) -> None: + """There is no shadow mode.""" + with self.assertRaises(ConfigError): + _read_config({"experimental_features": {"rust_sliding_sync": mode}}) + def _read_config(config_values: JsonDict) -> ExperimentalConfig: config = ExperimentalConfig(RootConfig()) diff --git a/tests/handlers/test_sliding_sync_rust.py b/tests/handlers/test_sliding_sync_rust.py new file mode 100644 index 0000000000..20bfe04004 --- /dev/null +++ b/tests/handlers/test_sliding_sync_rust.py @@ -0,0 +1,79 @@ +# +# 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 wiring of the Rust sliding sync handler (`rust/src/sliding_sync/`).""" + +from unittest import SkipTest + +from prometheus_client import REGISTRY + +from synapse.handlers.sliding_sync import observe_rust_sliding_sync_stats +from synapse.synapse_rust.sliding_sync import SlidingSyncStats + +import tests.utils +from tests import unittest +from tests.unittest import override_config + + +class RustSlidingSyncHandlerTestCase(unittest.HomeserverTestCase): + """The Rust sliding sync handler is only wired up on the native Rust + driver.""" + + @override_config({"experimental_features": {"rust_sliding_sync": "on"}}) + def test_needs_the_rust_driver(self) -> None: + if tests.utils.RUST_SLIDING_SYNC_FOR_TESTS: + raise SkipTest("SYNAPSE_TEST_RUST_SLIDING_SYNC overrides the config") + handler = self.hs.get_sliding_sync_handler() + self.assertEqual(handler._rust_sliding_sync_mode, "on") + if tests.utils.USE_RUST_DRIVER_FOR_TESTS: + self.assertIsNotNone(handler._rust_sliding_sync) + self.assertIsNotNone(self.hs.get_rust_handlers().sliding_sync) + else: + # Falls back to the Python handler, with a warning. + self.assertIsNone(handler._rust_sliding_sync) + self.assertIsNone(self.hs.get_rust_handlers().sliding_sync) + + def test_off_by_default(self) -> None: + handler = self.hs.get_sliding_sync_handler() + # Unless the test run sets it for every homeserver. + expected = tests.utils.RUST_SLIDING_SYNC_FOR_TESTS or "off" + self.assertEqual(handler._rust_sliding_sync_mode, expected) + if expected == "off": + self.assertIsNone(handler._rust_sliding_sync) + + def test_observe_stats(self) -> None: + def sample(name: str) -> float: + value = REGISTRY.get_sample_value(name, {"server_name": "stats.test"}) + return value or 0.0 + + before_db = sample("synapse_sliding_sync_rust_db_seconds_sum") + before_queries = sample("synapse_sliding_sync_rust_db_queries_sum") + before_rooms = sample("synapse_sliding_sync_rust_rooms_computed_count") + + stats = SlidingSyncStats( + db_time=0.5, queries=12, transactions=2, rooms_computed=3 + ) + self.assertEqual(stats.queries, 12) + self.assertEqual(stats.lazy_member_lookups, 0) + observe_rust_sliding_sync_stats(stats, "stats.test") + + self.assertAlmostEqual( + sample("synapse_sliding_sync_rust_db_seconds_sum") - before_db, 0.5 + ) + self.assertEqual( + sample("synapse_sliding_sync_rust_db_queries_sum") - before_queries, 12 + ) + self.assertEqual( + sample("synapse_sliding_sync_rust_rooms_computed_count") - before_rooms, 1 + ) diff --git a/tests/handlers/test_sliding_sync_rust_pure.py b/tests/handlers/test_sliding_sync_rust_pure.py new file mode 100644 index 0000000000..852e0d7803 --- /dev/null +++ b/tests/handlers/test_sliding_sync_rust_pure.py @@ -0,0 +1,1281 @@ +# +# 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 pure logic of the Rust sliding sync implementation +(`rust/src/sliding_sync/`) against the Python it ports, over generated +inputs.""" + +import json +import random +from types import SimpleNamespace +from typing import AbstractSet, Any, Coroutine, Mapping, TypeVar, cast + +from synapse.api.constants import EventTypes, Membership +from synapse.api.errors import SlidingSyncUnknownPosition +from synapse.events import EventBase +from synapse.events.utils import FilteredEvent +from synapse.handlers.sliding_sync import ( + MAX_NUMBER_PREVIOUS_STATE_KEYS_TO_REMEMBER, + SlidingSyncHandler, + _required_state_changes, + _required_state_request, + _room_config_to_persist, + _state_filter_with_meta, + _stripped_state, + _timeline_bounds, +) +from synapse.handlers.sliding_sync.extensions import SlidingSyncExtensionHandler +from synapse.handlers.sliding_sync.room_lists import ( + SlidingSyncRoomLists, + _sync_list_ops, + filter_membership_for_sync, +) +from synapse.storage.roommember import RoomsForUserSlidingSync +from synapse.synapse_rust import sliding_sync as rust +from synapse.types import ( + PersistedEventPosition, + RoomStreamToken, + SlidingSyncStreamToken, + StrCollection, + StreamToken, +) +from synapse.types.handlers import SLIDING_SYNC_DEFAULT_BUMP_EVENT_TYPES +from synapse.types.handlers.sliding_sync import ( + HaveSentRoom, + HaveSentRoomFlag, + OperationType, + PerConnectionState, + RoomLazyMembershipChanges, + RoomStatusMap, + RoomSyncConfig, + SlidingSyncConfig, + SlidingSyncResult, + StateValues, +) +from synapse.types.state import StateFilter +from synapse.util.clock import Clock + +from tests import unittest + +T = TypeVar("T") + +# The values the generators draw state types and keys from: the wildcards, +# local and remote users, and plain strings. +_TYPES = [ + StateValues.WILDCARD, + EventTypes.Member, + EventTypes.Name, + EventTypes.Topic, + "m.space.child", +] +_KEYS = [ + StateValues.WILDCARD, + StateValues.LAZY, + StateValues.ME, + "", + "@alice:test", + "@bob:test", + "@carol:remote", + "@dave:remote", + "alice", + "@nohost", +] +_SERVER_NAME = "test" + +RustConfig = tuple[int, dict[str, list[str]]] + + +def _is_mine_id(user_id: str) -> bool: + parts = user_id.split(":", 1) + return len(parts) == 2 and parts[1] == _SERVER_NAME + + +def _random_required_state(rng: random.Random) -> list[tuple[str, str]]: + return [(rng.choice(_TYPES), rng.choice(_KEYS)) for _ in range(rng.randint(0, 6))] + + +def _python_config( + timeline_limit: int, required_state: list[tuple[str, str]] +) -> RoomSyncConfig: + return RoomSyncConfig.from_room_config( + SlidingSyncConfig.RoomSubscription( + timeline_limit=timeline_limit, required_state=required_state + ) + ) + + +def _to_rust(config: RoomSyncConfig) -> RustConfig: + return ( + config.timeline_limit, + {t: sorted(keys) for t, keys in config.required_state_map.items()}, + ) + + +def _as_sets( + required_state_map: Mapping[str, AbstractSet[str] | list[str]], +) -> dict[str, set[str]]: + return {t: set(keys) for t, keys in required_state_map.items()} + + +class RoomSyncConfigTestCase(unittest.TestCase): + """`rust/src/sliding_sync/config.rs`""" + + def test_from_room_config(self) -> None: + rng = random.Random(1) + for _ in range(2000): + limit = rng.randint(0, 20) + required_state = _random_required_state(rng) + python = _python_config(limit, required_state) + timeline_limit, required_state_map = ( + rust.room_sync_config_from_room_config_for_test(limit, required_state) + ) + self.assertEqual(timeline_limit, python.timeline_limit) + self.assertEqual( + _as_sets(required_state_map), + _as_sets(python.required_state_map), + required_state, + ) + + def test_combine(self) -> None: + rng = random.Random(2) + for _ in range(2000): + a = _python_config(rng.randint(0, 20), _random_required_state(rng)) + b = _python_config(rng.randint(0, 20), _random_required_state(rng)) + python = a.combine_room_sync_config(b) + timeline_limit, required_state_map = rust.combine_room_sync_config_for_test( + _to_rust(a), _to_rust(b) + ) + self.assertEqual(timeline_limit, python.timeline_limit) + self.assertEqual( + _as_sets(required_state_map), + _as_sets(python.required_state_map), + (a, b), + ) + + def test_must_await_full_state(self) -> None: + rng = random.Random(3) + raised = 0 + for _ in range(2000): + config = _python_config(1, _random_required_state(rng)) + try: + python: bool | None = config.must_await_full_state(_is_mine_id) + except IndexError: + # An empty state key under the wildcard type. + python = None + raised += 1 + try: + result: bool | None = rust.must_await_full_state_for_test( + _to_rust(config), _SERVER_NAME + ) + except ValueError: + result = None + if python is None and result is True: + # Python raises on the empty state key before it reaches a + # remote user ID in its set order. Rust takes the order + # with the user IDs first. + self.assertTrue( + any( + key.startswith("@") and not _is_mine_id(key) + for key in config.required_state_map[StateValues.WILDCARD] + ) + ) + continue + self.assertEqual(result, python, config) + # The generator reaches the case Python raises on. + self.assertGreater(raised, 0) + + +_USER_ID = "@alice:test" +_USERS = ["@alice:test", "@bob:test", "@carol:remote", "@dave:remote", "@erin:test"] + + +def _filter_to_tuple( + state_filter: StateFilter, +) -> tuple[dict[str, list[str] | None], bool]: + """A Python `StateFilter` in the form the Rust test exports return.""" + return ( + { + t: sorted(keys) if keys is not None else None + for t, keys in state_filter.types.items() + }, + state_filter.include_others, + ) + + +def _random_map(rng: random.Random) -> dict[str, set[str]]: + """A required state map, now and then with more state keys of one type + than the connection remembers.""" + config = _python_config(1, _random_required_state(rng)) + required_state_map = {t: set(keys) for t, keys in config.required_state_map.items()} + if rng.random() < 0.1: + event_type = rng.choice([EventTypes.Member, EventTypes.Name]) + required_state_map[event_type] = { + f"@many{i}:test" for i in rng.sample(range(150), rng.randint(60, 120)) + } + return required_state_map + + +def _random_deltas(rng: random.Random) -> dict[tuple[str, str], str]: + deltas = {} + for _ in range(rng.randint(0, 5)): + event_type = rng.choice(_TYPES[1:]) + state_key = rng.choice(_USERS + ["", "@many7:test"]) + deltas[(event_type, state_key)] = "$event" + return deltas + + +def _inheritable( + prev: AbstractSet[str], request: AbstractSet[str], changed: AbstractSet[str] +) -> set[str]: + """The previous state keys `_required_state_changes` may keep.""" + invalidated = (prev - request) & changed + return set(prev - {StateValues.WILDCARD, StateValues.LAZY}) - invalidated + + +class RequiredStateChangesTestCase(unittest.TestCase): + """`rust/src/sliding_sync/required_state.rs`: `_required_state_changes`.""" + + def test_against_python(self) -> None: + rng = random.Random(4) + capped = 0 + for _ in range(3000): + prev = _random_map(rng) + # Often a small edit of the previous map. + if rng.random() < 0.5: + request = {t: set(keys) for t, keys in prev.items()} + for _ in range(rng.randint(1, 3)): + event_type = rng.choice(_TYPES) + keys = request.setdefault(event_type, set()) + key = rng.choice(_KEYS) + if key in keys and len(keys) > 1: + keys.discard(key) + else: + keys.add(key) + else: + request = _random_map(rng) + previously_returned = {u for u in _USERS if rng.random() < 0.4} + request_lazy = {u for u in _USERS if rng.random() < 0.4} + deltas = _random_deltas(rng) + + python = _required_state_changes( + _USER_ID, + prev_required_state_map=prev, + request_required_state_map=request, + previously_returned_lazy_user_ids=previously_returned, + request_lazy_load_user_ids=request_lazy, + state_deltas=deltas, + ) + changed, added, extra, invalidated = rust.required_state_changes_for_test( + _USER_ID, + {t: sorted(keys) for t, keys in prev.items()}, + {t: sorted(keys) for t, keys in request.items()}, + previously_returned, + request_lazy, + list(deltas), + ) + context = (prev, request, previously_returned, request_lazy, deltas) + + self.assertEqual( + added, _filter_to_tuple(python.added_state_filter), context + ) + self.assertEqual(set(invalidated), python.lazy_members_invalidated, context) + if deltas: + self.assertEqual( + set(extra), python.extra_users_to_add_to_lazy_cache, context + ) + else: + # Without deltas Python reads the state keys of whichever + # type a set union iterated last: Rust takes the members'. + # One of the two finds no users. + if set(extra) != python.extra_users_to_add_to_lazy_cache: + self.assertIn( + set(), + (set(extra), python.extra_users_to_add_to_lazy_cache), + context, + ) + + if python.changed_required_state_map is None: + self.assertIsNone(changed, context) + continue + assert changed is not None + python_map = _as_sets(python.changed_required_state_map) + rust_map = _as_sets(changed) + self.assertEqual(rust_map.keys(), python_map.keys(), context) + changed_keys: dict[str, set[str]] = {} + for event_type, state_key in deltas: + changed_keys.setdefault(event_type, set()).add(state_key) + for event_type, python_keys in python_map.items(): + rust_keys = rust_map[event_type] + if rust_keys == python_keys: + continue + # Over the cap, each keeps the requested state keys and an + # arbitrary slice of the previous ones. + requested = request.get(event_type, set()) + inheritable = _inheritable( + prev.get(event_type, set()), + requested, + changed_keys.get(event_type, set()), + ) + self.assertGreater( + len(requested | inheritable), + MAX_NUMBER_PREVIOUS_STATE_KEYS_TO_REMEMBER, + context, + ) + for keys in (rust_keys, python_keys): + self.assertLessEqual(requested, keys) + self.assertLessEqual(keys, requested | inheritable) + self.assertLessEqual( + len(keys), + max(len(requested), MAX_NUMBER_PREVIOUS_STATE_KEYS_TO_REMEMBER), + ) + capped += 1 + self.assertGreater(capped, 0) + + +def _event(sender: str, event_type: str, state_key: str | None) -> EventBase: + """A stand-in for an event, with the fields `_required_state_request` + reads.""" + return cast( + EventBase, + SimpleNamespace(sender=sender, type=event_type, state_key=state_key), + ) + + +class RequiredStateRequestTestCase(unittest.TestCase): + """`rust/src/sliding_sync/required_state.rs`: `_required_state_request` + and `_state_filter_with_meta`.""" + + def test_against_python(self) -> None: + rng = random.Random(5) + for _ in range(3000): + membership = rng.choice( + [Membership.JOIN, Membership.LEAVE, Membership.INVITE, Membership.KNOCK] + ) + required_state_map = _python_config( + 1, _random_required_state(rng) + ).required_state_map + timeline: list[tuple[str, str, str | None]] = [] + for _ in range(rng.randint(0, 4)): + if rng.random() < 0.3: + timeline.append( + (rng.choice(_USERS), EventTypes.Member, rng.choice(_USERS)) + ) + else: + timeline.append((rng.choice(_USERS), "m.room.message", None)) + limited = rng.choice([None, True, False]) + initial = rng.random() < 0.5 + deltas = list(_random_deltas(rng)) + + python = _required_state_request( + _USER_ID, + membership=membership, + required_state_map=required_state_map, + timeline=[_event(*event) for event in timeline], + limited=limited, + initial=initial, + state_delta_keys=deltas, + ) + state_filter, lazy, lazy_user_ids, explicit = ( + rust.required_state_request_for_test( + _USER_ID, + membership, + {t: sorted(keys) for t, keys in required_state_map.items()}, + timeline, + limited, + initial, + deltas, + ) + ) + context = (membership, required_state_map, timeline, limited, initial) + self.assertEqual( + state_filter, _filter_to_tuple(python.required_state_filter), context + ) + self.assertEqual(lazy, python.lazy_load_room_members, context) + self.assertEqual(set(lazy_user_ids), python.lazy_load_user_ids, context) + self.assertEqual(set(explicit), python.explicit_user_state, context) + + heroes = rng.sample(_USERS, rng.randint(0, 3)) + include_avatar = rng.random() < 0.5 + self.assertEqual( + rust.state_filter_with_meta_for_test( + python.required_state_filter, heroes, include_avatar + ), + _filter_to_tuple( + _state_filter_with_meta( + python.required_state_filter, + heroes, + include_avatar=include_avatar, + ) + ), + context, + ) + + +_ROOMS = [f"!r{i}:test" for i in range(8)] +_HOUR_MS = 60 * 60 * 1000 + +PyStatus = tuple[str, int | None] + + +def _random_statuses(rng: random.Random) -> dict[str, PyStatus]: + statuses: dict[str, PyStatus] = {} + for room_id in _ROOMS: + choice = rng.random() + if choice < 0.3: + continue + elif choice < 0.4: + statuses[room_id] = ("never", None) + elif choice < 0.7: + statuses[room_id] = ("live", None) + else: + statuses[room_id] = ("previously", rng.randint(1, 50)) + return statuses + + +def _to_status(status: PyStatus) -> HaveSentRoom[int]: + flag, token = status + if flag == "never": + return HaveSentRoom.never() + if flag == "live": + return HaveSentRoom.live() + assert token is not None + return HaveSentRoom.previously(token) + + +def _from_status(status: HaveSentRoom[int]) -> PyStatus: + return (status.status.value, status.last_token) + + +class _FakeClock: + def __init__(self, now_ms: int) -> None: + self.now_ms = now_ms + + def time_msec(self) -> int: + return self.now_ms + + +class ConnectionStateTestCase(unittest.TestCase): + """`rust/src/sliding_sync/connection.rs`""" + + def test_room_status_map(self) -> None: + rng = random.Random(6) + for _ in range(2000): + statuses = _random_statuses(rng) + ops: list[tuple[str, list[str], int | None]] = [] + for _ in range(rng.randint(0, 4)): + room_ids = rng.sample(_ROOMS, rng.randint(0, 4)) + if rng.random() < 0.5: + ops.append(("sent", room_ids, None)) + else: + ops.append(("unsent", room_ids, rng.randint(1, 50))) + + python = RoomStatusMap( + statuses={k: _to_status(v) for k, v in statuses.items()} + ).get_mutable() + for op, room_ids, token in ops: + if op == "sent": + python.record_sent_rooms(room_ids) + else: + assert token is not None + python.record_unsent_rooms(room_ids, token) + + updates, flattened, length = rust.room_status_map_ops_for_test( + statuses, ops + ) + context = (statuses, ops) + self.assertEqual( + updates, + {k: _from_status(v) for k, v in python.get_updates().items()}, + context, + ) + copied = python.copy() + self.assertEqual( + flattened, + { + room_id: _from_status(copied.have_sent_room(room_id)) + for room_id in _ROOMS + if room_id in statuses or room_id in python.get_updates() + }, + context, + ) + self.assertEqual(length, len(python), context) + + def test_lazy_membership_changes(self) -> None: + rng = random.Random(7) + for _ in range(2000): + now_ms = 10 * _HOUR_MS + returned = { + user_id: rng.choice( + [None, now_ms, now_ms - _HOUR_MS, now_ms - _HOUR_MS + 1, 0] + ) + for user_id in rng.sample(_USERS, rng.randint(0, 4)) + } + invalidated = set(rng.sample(_USERS, rng.choice([0, 0, 1]))) + python = RoomLazyMembershipChanges( + returned_user_id_to_last_seen_ts_map=returned, + invalidated_user_ids=invalidated, + ) + clock = cast(Clock, _FakeClock(now_ms)) + to_update, has_updates = rust.lazy_membership_changes_for_test( + returned, invalidated, now_ms + ) + self.assertEqual( + to_update, sorted(python.get_returned_user_ids_to_update(clock)) + ) + self.assertEqual(has_updates, python.has_updates(clock)) + + def test_per_connection_state(self) -> None: + rng = random.Random(8) + now_ms = 10 * _HOUR_MS + for _ in range(1000): + rooms = _random_statuses(rng) + receipts = _random_statuses(rng) + account_data = _random_statuses(rng) + sent = [rng.sample(_ROOMS, rng.choice([0, 0, 1, 2])) for _ in range(3)] + room_configs = { + room_id: _to_rust(_python_config(1, _random_required_state(rng))) + for room_id in rng.sample(_ROOMS, rng.randint(0, 3)) + } + room_config_updates = { + room_id: _to_rust(_python_config(2, _random_required_state(rng))) + for room_id in rng.sample(_ROOMS, rng.choice([0, 0, 1])) + } + lazy = { + room_id: ( + { + user_id: rng.choice([None, now_ms]) + for user_id in rng.sample(_USERS, rng.choice([0, 1])) + }, + set(rng.sample(_USERS, rng.choice([0, 0, 0, 1]))), + ) + for room_id in rng.sample(_ROOMS, rng.choice([0, 1])) + } + + def config(rust_config: RustConfig) -> RoomSyncConfig: + limit, required_state_map = rust_config + return RoomSyncConfig( + limit, {t: set(keys) for t, keys in required_state_map.items()} + ) + + state = PerConnectionState( + rooms=RoomStatusMap( + statuses={ + k: HaveSentRoom( + HaveSentRoomFlag(flag), + RoomStreamToken(stream=token) if token else None, + ) + for k, (flag, token) in rooms.items() + } + ), + # The receipts stream's tokens are ints here. + receipts=RoomStatusMap( + statuses=cast(Any, {k: _to_status(v) for k, v in receipts.items()}) + ), + account_data=RoomStatusMap( + statuses={k: _to_status(v) for k, v in account_data.items()} + ), + room_configs={k: config(v) for k, v in room_configs.items()}, + ) + python = state.get_mutable() + python.rooms.record_sent_rooms(sent[0]) + python.receipts.record_sent_rooms(sent[1]) + python.account_data.record_sent_rooms(sent[2]) + for room_id, rust_config in room_config_updates.items(): + python.room_configs[room_id] = config(rust_config) + for room_id, (returned, invalidated) in lazy.items(): + python.room_lazy_membership[room_id] = RoomLazyMembershipChanges( + returned_user_id_to_last_seen_ts_map=returned, + invalidated_user_ids=invalidated, + ) + + self.assertEqual( + rust.per_connection_state_for_test( + rooms, + sent[0], + receipts, + sent[1], + account_data, + sent[2], + room_configs, + room_config_updates, + lazy, + now_ms, + ), + ( + python.has_updates(cast(Clock, _FakeClock(now_ms))), + len(state), + len(python), + ), + ) + + +PyRoom = tuple[str, str | None, str, str | None, int, bool, str | None, bool] + + +def _run(coroutine: Coroutine[Any, Any, T]) -> T: + """Run a coroutine that never waits.""" + try: + coroutine.send(None) + except StopIteration as e: + return e.value + raise AssertionError("the coroutine waited") + + +def _random_room(rng: random.Random, room_id: str) -> PyRoom: + membership = rng.choice( + [Membership.JOIN, Membership.INVITE, Membership.LEAVE, Membership.BAN] + ) + state_reset = rng.random() < 0.1 + return ( + room_id, + None if state_reset else rng.choice([_USER_ID, "@bob:test"]), + membership, + None if state_reset else f"$event_{room_id}", + rng.randint(1, 100), + rng.random() < 0.8, + rng.choice([None, None, "m.space", "org.example.type"]), + rng.random() < 0.5, + ) + + +def _python_room(room: PyRoom) -> RoomsForUserSlidingSync: + ( + room_id, + sender, + membership, + event_id, + stream, + has_known_state, + room_type, + is_encrypted, + ) = room + return RoomsForUserSlidingSync( + room_id=room_id, + sender=sender, + membership=membership, + event_id=event_id, + event_pos=PersistedEventPosition("master", stream), + room_version_id="10", + has_known_state=has_known_state, + room_type=room_type, + is_encrypted=is_encrypted, + ) + + +def _random_filters(rng: random.Random) -> dict[str, Any]: + filters: dict[str, Any] = {} + types = [None, "m.space", "org.example.type", "org.other"] + tags = ["m.favourite", "m.lowpriority", "u.work"] + if rng.random() < 0.4: + filters["is_dm"] = rng.random() < 0.5 + if rng.random() < 0.4: + filters["is_encrypted"] = rng.random() < 0.5 + if rng.random() < 0.4: + filters["is_invite"] = rng.random() < 0.5 + if rng.random() < 0.3: + filters["room_types"] = rng.sample(types, rng.randint(0, 2)) + if rng.random() < 0.3: + filters["not_room_types"] = rng.sample(types, rng.randint(0, 2)) + if rng.random() < 0.3: + filters["tags"] = rng.sample(tags, rng.randint(0, 2)) + if rng.random() < 0.3: + filters["not_tags"] = rng.sample(tags, rng.randint(0, 2)) + if rng.random() < 0.03: + filters["spaces"] = ["!space:test"] + if rng.random() < 0.03: + filters["room_name_like"] = "name" + return filters + + +class _FakeTagStore: + def __init__(self, room_tags: dict[str, set[str]]) -> None: + self.room_tags = room_tags + + async def get_tags_for_user(self, user_id: str) -> dict[str, dict[str, Any]]: + return { + room_id: {tag: {} for tag in tags} + for room_id, tags in self.room_tags.items() + } + + +class RoomListFilterTestCase(unittest.TestCase): + """`rust/src/sliding_sync/room_lists/filter.rs`""" + + def test_filter_membership_for_sync(self) -> None: + rng = random.Random(9) + for _ in range(500): + room = _random_room(rng, "!r:test") + newly_left = rng.random() < 0.3 + self.assertEqual( + rust.filter_membership_for_sync_for_test(_USER_ID, room, newly_left), + filter_membership_for_sync( + user_id=_USER_ID, + room_membership_for_user=_python_room(room), + newly_left=newly_left, + ), + (room, newly_left), + ) + + def test_filter_rooms_using_tables(self) -> None: + rng = random.Random(10) + not_implemented = 0 + for _ in range(2000): + rooms = [_random_room(rng, room_id) for room_id in _ROOMS] + filters = _random_filters(rng) + dm_room_ids = {room_id for room_id in _ROOMS if rng.random() < 0.3} + room_tags = { + room_id: set(rng.sample(["m.favourite", "m.lowpriority", "u.work"], 1)) + for room_id in _ROOMS + if rng.random() < 0.4 + } + sent_before = {room_id for room_id in _ROOMS if rng.random() < 0.5} + + fake_room_lists = cast( + SlidingSyncRoomLists, + SimpleNamespace(store=_FakeTagStore(room_tags)), + ) + previous_state = PerConnectionState( + rooms=RoomStatusMap( + statuses={room_id: HaveSentRoom.live() for room_id in sent_before} + ) + ) + try: + python: set[str] | None = set( + _run( + SlidingSyncRoomLists.filter_rooms_using_tables( + fake_room_lists, + _USER_ID, + {room[0]: _python_room(room) for room in rooms}, + previous_state, + SlidingSyncConfig.SlidingSyncList.Filters(**filters), + cast(StreamToken, None), + dm_room_ids, + ) + ) + ) + except NotImplementedError: + python = None + not_implemented += 1 + try: + result: set[str] | None = set( + rust.filter_rooms_for_test( + rooms, json.dumps(filters), dm_room_ids, room_tags, sent_before + ) + ) + except NotImplementedError: + result = None + self.assertEqual(result, python, (rooms, filters)) + self.assertGreater(not_implemented, 0) + + +class _FakeStreamCache: + def __init__(self, earliest: int, positions: dict[str, int]) -> None: + self.earliest = earliest + self.positions = positions + + def get_earliest_known_position(self) -> int: + return self.earliest + + def get_max_pos_of_last_change(self, room_id: str) -> int | None: + return self.positions.get(room_id) + + +class _FakeSortStore: + def __init__( + self, + stream_cache: _FakeStreamCache, + recheck: dict[str, int | None], + joined_positions: dict[str, int], + ) -> None: + self._events_stream_cache = stream_cache + self.recheck = recheck + self.joined_positions = joined_positions + + async def get_last_event_pos_in_room_before_stream_ordering( + self, room_id: str, end_token: RoomStreamToken + ) -> tuple[str, PersistedEventPosition] | None: + stream = self.recheck[room_id] + if stream is None: + return None + return ("$event", PersistedEventPosition("master", stream)) + + async def bulk_get_last_event_pos_in_room_before_stream_ordering( + self, room_ids: list[str], end_token: RoomStreamToken + ) -> dict[str, int]: + return { + room_id: self.joined_positions[room_id] + for room_id in room_ids + if room_id in self.joined_positions + } + + +class _FakeUpdatesStore: + def __init__(self, rooms_with_updates: set[str]) -> None: + self.rooms_with_updates = rooms_with_updates + + async def get_rooms_that_have_updates_since_sliding_sync_table( + self, room_ids: StrCollection, from_key: RoomStreamToken + ) -> set[str]: + return {room_id for room_id in room_ids if room_id in self.rooms_with_updates} + + +class RoomListSortTestCase(unittest.TestCase): + """`rust/src/sliding_sync/room_lists/sort.rs` and `lists.rs`""" + + def test_sort_rooms(self) -> None: + rng = random.Random(11) + early = 0 + for _ in range(2000): + rooms = [_random_room(rng, room_id) for room_id in _ROOMS] + # Joined rooms are the common case. + rooms = [ + (r[0], r[1], Membership.JOIN, *r[3:]) if rng.random() < 0.6 else r + for r in rooms + ] + joined = [r[0] for r in rooms if r[2] == Membership.JOIN] + earliest = rng.randint(0, 50) + to_stream = rng.randint(40, 100) + cache = { + room_id: rng.randint(earliest + 1, 120) + for room_id in joined + if rng.random() < 0.7 + } + recheck = { + room_id: rng.choice([None, rng.randint(1, to_stream)]) + for room_id in _ROOMS + } + joined_positions = { + room_id: rng.randint(1, to_stream) + for room_id in joined + if rng.random() < 0.98 + } + limit = rng.choice([None, 1, 2, 3, 5, 8]) + + store = _FakeSortStore( + _FakeStreamCache(earliest, cache), recheck, joined_positions + ) + fake_room_lists = cast(SlidingSyncRoomLists, SimpleNamespace(store=store)) + to_token = cast( + StreamToken, + SimpleNamespace(room_key=RoomStreamToken(stream=to_stream)), + ) + try: + python: list[str] | None = [ + room.room_id + for room in _run( + SlidingSyncRoomLists.sort_rooms( + fake_room_lists, + {room[0]: _python_room(room) for room in rooms}, + to_token, + limit=limit, + ) + ) + ] + except KeyError: + python = None + try: + result: list[str] | None = rust.sort_rooms_for_test( + rooms, earliest, cache, to_stream, limit, recheck, joined_positions + ) + except KeyError: + result = None + self.assertEqual(result, python, (rooms, cache, limit)) + if python is not None and len(python) < len(rooms): + early += 1 + self.assertGreater(early, 0) + + def test_sync_list_ops(self) -> None: + rng = random.Random(12) + for _ in range(2000): + sorted_room_ids = rng.sample(_ROOMS, rng.randint(0, len(_ROOMS))) + ranges = [] + for _ in range(rng.randint(1, 3)): + start = rng.randint(0, 9) + # Pydantic allows an end before the start, but not below 0. + ranges.append((start, max(0, start + rng.randint(-1, 4)))) + config = _python_config(rng.randint(0, 10), _random_required_state(rng)) + relevant = { + room_id: _python_config(rng.randint(0, 10), _random_required_state(rng)) + for room_id in rng.sample(_ROOMS, rng.randint(0, 3)) + } + rust_relevant = [(k, _to_rust(v)) for k, v in relevant.items()] + + ops = _sync_list_ops(ranges, sorted_room_ids, config, relevant) + rust_ops, rust_map = rust.sync_list_ops_for_test( + ranges, sorted_room_ids, _to_rust(config), rust_relevant + ) + context = (ranges, sorted_room_ids, config, rust_relevant) + self.assertEqual( + rust_ops, + [(tuple(op.range), list(op.room_ids)) for op in ops], + context, + ) + self.assertEqual( + [(k, (limit, _as_sets(m))) for k, (limit, m) in rust_map], + [ + (k, (v.timeline_limit, _as_sets(v.required_state_map))) + for k, v in relevant.items() + ], + context, + ) + + def test_rooms_to_send(self) -> None: + rng = random.Random(13) + now_ms = 10 * _HOUR_MS + reset = 0 + for _ in range(1000): + many = rng.random() < 0.2 + room_ids = [f"!many{i}:test" for i in range(130)] if many else list(_ROOMS) + relevant = { + room_id: _python_config(rng.randint(0, 5), _random_required_state(rng)) + for room_id in room_ids + if rng.random() < 0.8 + } + room_configs = { + room_id: _python_config(rng.randint(0, 5), _random_required_state(rng)) + for room_id in room_ids + if rng.random() < 0.5 + } + statuses = _random_statuses(rng) + if many: + statuses = dict.fromkeys(room_ids, ("live", None)) + rooms_with_updates = { + room_id for room_id in room_ids if rng.random() < (0.9 if many else 0.3) + } + last_used_ts = rng.choice([None, now_ms - 2 * _HOUR_MS, now_ms - 60_000]) + + previous_state = PerConnectionState( + last_used_ts=last_used_ts, + rooms=RoomStatusMap( + statuses={ + k: HaveSentRoom( + HaveSentRoomFlag(flag), + RoomStreamToken(stream=token) if token else None, + ) + for k, (flag, token) in statuses.items() + } + ), + room_configs=room_configs, + ) + fake_room_lists = cast( + SlidingSyncRoomLists, + SimpleNamespace( + store=_FakeUpdatesStore(rooms_with_updates), + _clock=_FakeClock(now_ms), + ), + ) + from_token = cast( + StreamToken, SimpleNamespace(room_key=RoomStreamToken(stream=1)) + ) + try: + python: list[str] | None = list( + _run( + SlidingSyncRoomLists._filter_relevant_rooms_to_send( + fake_room_lists, previous_state, from_token, relevant + ) + ) + ) + except SlidingSyncUnknownPosition: + python = None + reset += 1 + try: + result: list[str] | None = rust.rooms_to_send_for_test( + [(k, _to_rust(v)) for k, v in relevant.items()], + {k: _to_rust(v) for k, v in room_configs.items()}, + statuses, + last_used_ts, + {room_id for room_id in relevant if room_id in rooms_with_updates}, + now_ms, + ) + except ValueError: + result = None + self.assertEqual(result, python) + self.assertGreater(reset, 0) + + +class _FakeBumpStore: + def __init__( + self, + latest_room_bump_stamp: int | None, + have_finished_bg_jobs: bool, + last_event_stream: int | None, + ) -> None: + self.latest_room_bump_stamp = latest_room_bump_stamp + self.have_finished_bg_jobs = have_finished_bg_jobs + self.last_event_stream = last_event_stream + + async def get_latest_bump_stamp_for_room(self, room_id: str) -> int | None: + return self.latest_room_bump_stamp + + async def have_finished_sliding_sync_background_jobs(self) -> bool: + return self.have_finished_bg_jobs + + async def get_last_event_pos_in_room_before_stream_ordering( + self, room_id: str, end_token: RoomStreamToken, event_types: StrCollection + ) -> tuple[str, PersistedEventPosition] | None: + assert set(event_types) == SLIDING_SYNC_DEFAULT_BUMP_EVENT_TYPES + if self.last_event_stream is None: + return None + return ("$event", PersistedEventPosition("master", self.last_event_stream)) + + +class RoomHelpersTestCase(unittest.TestCase): + """`rust/src/sliding_sync/room.rs`""" + + def test_timeline_bounds(self) -> None: + rng = random.Random(14) + for _ in range(2000): + status = _random_statuses(rng).get("!r0:test", ("never", None)) + from_stream = rng.choice([None, rng.randint(1, 50)]) + skip = rng.random() < 0.2 + prev = rng.choice( + [None, _python_config(rng.randint(0, 10), _random_required_state(rng))] + ) + config = _python_config(rng.randint(0, 10), _random_required_state(rng)) + + flag, token = status + from_bound, initial, ignore = _timeline_bounds( + HaveSentRoom( + HaveSentRoomFlag(flag), + RoomStreamToken(stream=token) if token else None, + ), + from_token=cast( + SlidingSyncStreamToken, + SimpleNamespace( + stream_token=SimpleNamespace( + room_key=RoomStreamToken(stream=from_stream) + ) + ), + ) + if from_stream is not None + else None, + skip_previous_bounds=skip, + prev_room_sync_config=prev, + room_sync_config=config, + ) + self.assertEqual( + rust.timeline_bounds_for_test( + status, + from_stream, + skip, + _to_rust(prev) if prev else None, + _to_rust(config), + ), + (from_bound.stream if from_bound else None, initial, ignore), + (status, from_stream, skip, prev, config), + ) + + def test_stripped_state(self) -> None: + rng = random.Random(15) + for _ in range(500): + membership = rng.choice( + [Membership.INVITE, Membership.KNOCK, Membership.JOIN] + ) + unsigned: dict[str, Any] = {} + for key in ("invite_room_state", "knock_room_state"): + choice = rng.random() + if choice < 0.5: + unsigned[key] = [ + {"type": "m.room.name", "state_key": "", "content": {}}, + "not an event", + ][: rng.randint(0, 2)] + elif choice < 0.7: + unsigned[key] = {"not": "a list"} + event = { + "type": EventTypes.Member, + "state_key": _USER_ID, + "sender": "@bob:test", + "content": {"membership": membership, "displayname": "Alice"}, + "unsigned": unsigned, + } + python = _stripped_state( + cast( + EventBase, + SimpleNamespace( + membership=membership, + unsigned=unsigned, + type=event["type"], + state_key=event["state_key"], + sender=event["sender"], + content=event["content"], + ), + ) + ) + self.assertEqual( + json.loads(rust.stripped_state_for_test(json.dumps(event))), python + ) + + def test_get_bump_stamp(self) -> None: + rng = random.Random(16) + types = [*sorted(SLIDING_SYNC_DEFAULT_BUMP_EVENT_TYPES), EventTypes.Member] + for _ in range(2000): + timeline = [ + (rng.choice(types), rng.choice([-3, 0, rng.randint(1, 60)])) + for _ in range(rng.randint(0, 4)) + ] + check_outside = rng.random() < 0.7 + latest = rng.choice([None, -1, 0, rng.randint(1, 60)]) + have_finished = rng.random() < 0.7 + to_stream = rng.randint(1, 60) + last_event = rng.choice([None, -2, rng.randint(1, 60)]) + + fake_handler = cast( + SlidingSyncHandler, + SimpleNamespace( + store=_FakeBumpStore(latest, have_finished, last_event) + ), + ) + python = _run( + SlidingSyncHandler._get_bump_stamp( + fake_handler, + "!r:test", + cast( + StreamToken, + SimpleNamespace(room_key=RoomStreamToken(stream=to_stream)), + ), + [ + cast( + FilteredEvent, + SimpleNamespace( + event=SimpleNamespace( + type=event_type, + internal_metadata=SimpleNamespace( + stream_ordering=stream + ), + ) + ), + ) + for event_type, stream in timeline + ], + check_outside_timeline=check_outside, + ) + ) + self.assertEqual( + rust.get_bump_stamp_for_test( + timeline, + check_outside, + latest, + have_finished, + to_stream, + last_event, + ), + python, + (timeline, check_outside, latest, have_finished, to_stream, last_event), + ) + + def test_room_config_to_persist(self) -> None: + rng = random.Random(17) + for _ in range(2000): + config = _python_config(rng.randint(0, 10), _random_required_state(rng)) + prev = rng.choice( + [None, _python_config(rng.randint(0, 10), _random_required_state(rng))] + ) + changed: dict[str, set[str]] | None = rng.choice( + [None, {}, _random_map(rng)] + ) + ignore = rng.random() < 0.3 + limited = rng.choice([None, True, False]) + + python = _room_config_to_persist( + config, + prev_room_sync_config=prev, + changed_required_state_map=changed, + ignore_timeline_bound=ignore, + limited=limited, + ) + result = rust.room_config_to_persist_for_test( + _to_rust(config), + _to_rust(prev) if prev else None, + {t: sorted(keys) for t, keys in changed.items()} + if changed is not None + else None, + ignore, + limited, + ) + context = (config, prev, changed, ignore, limited) + if python is None: + self.assertIsNone(result, context) + continue + assert result is not None + self.assertEqual( + (result[0], _as_sets(result[1])), + (python.timeline_limit, _as_sets(python.required_state_map)), + context, + ) + + +class RelevantRoomIdsTestCase(unittest.TestCase): + """`rust/src/sliding_sync/extensions/relevant.rs`""" + + def test_against_python(self) -> None: + rng = random.Random(18) + names = ["a", "b", "c", "*"] + raised = 0 + for _ in range(2000): + actual_lists = [ + ( + name, + [ + rng.sample(_ROOMS, rng.randint(0, 3)) + for _ in range(rng.choice([1, 1, 1, 1, 0, 2])) + ], + ) + for name in rng.sample(names[:3], rng.randint(0, 3)) + ] + actual_room_ids = set(rng.sample(_ROOMS, rng.randint(0, 3))) + requested_lists = rng.choice( + [None, rng.sample(names + ["missing"], rng.randint(0, 3))] + ) + requested_room_ids = rng.choice( + [None, rng.sample([*_ROOMS[:4], "*"], rng.randint(0, 3))] + ) + + try: + python: list[str] | None = sorted( + SlidingSyncExtensionHandler.find_relevant_room_ids_for_extension( + cast(SlidingSyncExtensionHandler, None), + requested_lists, + requested_room_ids, + { + name: SlidingSyncResult.SlidingWindowList( + count=0, + ops=[ + SlidingSyncResult.SlidingWindowList.Operation( + op=OperationType.SYNC, + range=(0, 0), + room_ids=room_ids, + ) + for room_ids in ops + ], + ) + for name, ops in actual_lists + }, + actual_room_ids, + ) + ) + except AssertionError: + python = None + raised += 1 + try: + result: list[str] | None = rust.find_relevant_room_ids_for_test( + requested_lists, requested_room_ids, actual_lists, actual_room_ids + ) + except AssertionError: + result = None + self.assertEqual( + result, + python, + (requested_lists, requested_room_ids, actual_lists, actual_room_ids), + ) + self.assertGreater(raised, 0) diff --git a/tests/rest/client/sliding_sync/bench_rust_sliding_sync.py b/tests/rest/client/sliding_sync/bench_rust_sliding_sync.py new file mode 100644 index 0000000000..19b3c59b6e --- /dev/null +++ b/tests/rest/client/sliding_sync/bench_rust_sliding_sync.py @@ -0,0 +1,532 @@ +# +# 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: +# . +# + +"""Sliding sync in Rust against Python, timed per request. + +Not a test: trial only collects `test_*.py`, and the cases skip unless +`SYNAPSE_BENCH_SLIDING_SYNC` is set. Build the extension in release mode +first, as a debug build makes Rust look several times slower: + + poetry run maturin develop --release + SYNAPSE_BENCH_SLIDING_SYNC=1 SYNAPSE_TEST_RUST_DRIVER=1 \\ + SYNAPSE_TEST_RUST_SLIDING_SYNC=on POSTGRES=1 \\ + run-trial.sh tests.rest.client.sliding_sync.bench_rust_sliding_sync + +Each scenario runs on two connections of one user, one served by Python and +one by Rust, in turn, so both see the same data and the same warm caches. +Every request is timed on its own. The table gives the median of `REPS`: + +- wall: the whole request, servlet included. +- reactor: CPU time on the reactor thread, less the time inside Rust calls, + plus the time in the Python methods Rust calls back into. Rust's direct + reads of Python objects under the GIL are not counted. The trial wrapper runs + each Rust call's top-level future on the reactor thread through `block_on`, + where production runs it on a tokio worker. Trial also runs Python's + database functions on the reactor thread. So this column is the time each + side holds the GIL, which is what competes with the reactor. +- process: CPU time of the whole process, Rust's threads and the database + driver's included. +- db: time in database queries, Python's and Rust's. Rust runs some queries + side by side and adds up each one's time, so this can exceed wall. +- queries: database queries, Python's and Rust's. +- txns: database transactions. Rust counts only its explicit transactions, + not single reads, so compare queries across the two sides. + +Set `SYNAPSE_BENCH_OUT` to a path to also append the tables there, and +`SYNAPSE_BENCH_SCALE` to a fraction to make the accounts smaller for a quick +check of the script itself. Set `SYNAPSE_BENCH_PROFILE` to a directory to get +a cProfile dump of one more request per scenario and side, run after the +timed ones. The dump covers that request alone. +""" + +import cProfile +import os +import re +import statistics +import time +from typing import Any, Callable +from unittest import SkipTest +from unittest.mock import patch + +import attr + +from twisted.internet.testing import MemoryReactor + +import synapse.rest.admin +from synapse.api.constants import EventTypes, ReceiptTypes +from synapse.handlers.sliding_sync import ( + rust_sliding_sync_db_time, + rust_sliding_sync_queries, + rust_sliding_sync_requests_counter, + rust_sliding_sync_transactions, +) +from synapse.metrics import SERVER_NAME_LABEL +from synapse.rest.client import account_data, login, receipts, room, sync, tags +from synapse.server import HomeServer +from synapse.storage.database import sql_query_timer, sql_txn_count +from synapse.types import JsonDict +from synapse.util.clock import Clock + +from tests.rest.client.sliding_sync.test_sliding_sync import SlidingSyncBase +from tests.server import _BlockingRustCalls + +REPS = 7 + +SCALE = float(os.environ.get("SYNAPSE_BENCH_SCALE", "1")) + +REQUIRED_STATE = [ + [EventTypes.Create, ""], + [EventTypes.Name, ""], + [EventTypes.RoomAvatar, ""], + [EventTypes.Topic, ""], + [EventTypes.RoomEncryption, ""], + [EventTypes.CanonicalAlias, ""], + [EventTypes.PowerLevels, ""], + [EventTypes.Member, "$LAZY"], + [EventTypes.Member, "$ME"], +] + +EXTENSIONS = { + "account_data": {"enabled": True}, + "receipts": {"enabled": True}, +} + + +def _list_body( + end: int, timeline_limit: int = 20, extensions: bool = False +) -> JsonDict: + body: JsonDict = { + "lists": { + "all": { + "ranges": [[0, end]], + "required_state": REQUIRED_STATE, + "timeline_limit": timeline_limit, + } + } + } + if extensions: + body["extensions"] = EXTENSIONS + return body + + +def _counter_total(counter: object, server_name: str) -> float: + """The sum over every label set of a counter for one server.""" + total = 0.0 + for metric in counter.collect(): # type: ignore[attr-defined] + for sample in metric.samples: + if ( + sample.name.endswith("_total") + and sample.labels.get(SERVER_NAME_LABEL) == server_name + ): + total += sample.value + return total + + +def _histogram_total(histogram: object, server_name: str, suffix: str) -> float: + """The sum over every label set of a histogram's `_sum` or `_count` for + one server.""" + total = 0.0 + for metric in histogram.collect(): # type: ignore[attr-defined] + for sample in metric.samples: + if ( + sample.name.endswith(suffix) + and sample.labels.get(SERVER_NAME_LABEL) == server_name + ): + total += sample.value + return total + + +class _TimedPython: + """Wraps a Python object that Rust calls back into, and adds up the CPU + time of each method call on it.""" + + def __init__(self, wrapped: Any, calls: "_TimedRustCalls") -> None: + self._wrapped = wrapped + self._calls = calls + + def __getattr__(self, name: str) -> Any: + value = getattr(self._wrapped, name) + if not callable(value): + return value + + def timed(*args: Any, **kwargs: Any) -> Any: + start = time.thread_time() + try: + return value(*args, **kwargs) + finally: + self._calls.python_cpu += time.thread_time() - start + + return timed + + +class _TimedRustCalls(_BlockingRustCalls): + """Blocking Rust calls that add up the reactor thread's CPU time inside + them, and inside that, the time in the Python objects they call back + into.""" + + def __init__(self, handler: Any) -> None: + super().__init__(handler) + self.cpu = 0.0 + self.python_cpu = 0.0 + + def __getattr__(self, name: str) -> Any: + call = super().__getattr__(name) + if getattr(self._handler, f"{name}_blocking", None) is None: + return call + + def timed(*args: Any, **kwargs: Any) -> Any: + if "events_stream_cache" in kwargs: + kwargs["events_stream_cache"] = _TimedPython( + kwargs["events_stream_cache"], self + ) + start = time.thread_time() + try: + return call(*args, **kwargs) + finally: + self.cpu += time.thread_time() - start + + return timed + + +@attr.s(auto_attribs=True, frozen=True) +class Sample: + wall: float + reactor: float + process: float + db: float + queries: float + txns: float + + +class SlidingSyncBenchmark(SlidingSyncBase): + """Times sliding sync requests in Rust and in Python.""" + + servlets = [ + synapse.rest.admin.register_servlets, + account_data.register_servlets, + login.register_servlets, + receipts.register_servlets, + room.register_servlets, + sync.register_servlets, + tags.register_servlets, + ] + + # Room creation over the fake HTTP channel is slow at this size. + timeout = 3600 + + def default_config(self) -> JsonDict: + config = super().default_config() + # Trial's config sets the event cache to one entry, which turns Python's + # off. A server keeps the events of the first screen cached. + config["event_cache_size"] = 100_000 + return config + + def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None: + if not os.environ.get("SYNAPSE_BENCH_SLIDING_SYNC"): + raise SkipTest("set SYNAPSE_BENCH_SLIDING_SYNC to run the benchmark") + 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.rust = _TimedRustCalls(rust) + # The servlet's Rust encoder releases the GIL while it encodes, so + # count it with the blocking calls. That leaves out its argument + # conversion, which holds the GIL. + encode_rooms = sync.encode_rust_sliding_sync_rooms + + def timed_encode_rooms(*args: Any, **kwargs: Any) -> bytes: + start = time.thread_time() + try: + return encode_rooms(*args, **kwargs) + finally: + self.rust.cpu += time.thread_time() - start + + patcher = patch.object( + sync, "encode_rust_sliding_sync_rooms", timed_encode_rooms + ) + patcher.start() + self.addCleanup(patcher.stop) + self.handler = hs.get_sliding_sync_handler() + self.handler._rust_sliding_sync_strict = True + self.server_name = hs.hostname + + 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") + self.conn_seq = 0 + self.lines: list[str] = [] + self.profile: cProfile.Profile | None = None + + def _rust_requests(self) -> float: + return rust_sliding_sync_requests_counter.labels( + result="rust", **{SERVER_NAME_LABEL: self.server_name} + )._value.get() + + def _new_conn(self, side: str) -> str: + self.conn_seq += 1 + return f"{side}-{self.conn_seq}" + + def _sync( + self, side: str, conn_id: str, body: JsonDict, since: str | None + ) -> tuple[JsonDict, str]: + self.handler._rust_sliding_sync = ( + self.rust if side == "rust" else None # type: ignore[assignment] + ) + before = self._rust_requests() + response, pos = self.do_sync( + {**body, "conn_id": conn_id}, since=since, tok=self.user1_tok + ) + self.assertEqual(self._rust_requests() - before, int(side == "rust")) + return response, pos + + def _db_totals(self) -> tuple[float, float, float]: + """Database time, queries and transactions so far, Python's and + Rust's.""" + name = self.server_name + return ( + _histogram_total(sql_query_timer, name, "_sum") + + _histogram_total(rust_sliding_sync_db_time, name, "_sum"), + _histogram_total(sql_query_timer, name, "_count") + + _histogram_total(rust_sliding_sync_queries, name, "_sum"), + _counter_total(sql_txn_count, name) + + _histogram_total(rust_sliding_sync_transactions, name, "_sum"), + ) + + def _timed_sync( + self, side: str, conn_id: str, body: JsonDict, since: str | None + ) -> tuple[Sample, str]: + db, queries, txns = self._db_totals() + rust_cpu = self.rust.cpu + python_cpu = self.rust.python_cpu + wall = time.perf_counter() + reactor = time.thread_time() + process = time.process_time() + + if self.profile is not None: + _, pos = self.profile.runcall(self._sync, side, conn_id, body, since) + else: + _, pos = self._sync(side, conn_id, body, since) + + wall = time.perf_counter() - wall + reactor = ( + time.thread_time() + - reactor + - (self.rust.cpu - rust_cpu) + + (self.rust.python_cpu - python_cpu) + ) + process = time.process_time() - process + db_after, queries_after, txns_after = self._db_totals() + sample = Sample( + wall=wall, + reactor=reactor, + process=process, + db=db_after - db, + queries=queries_after - queries, + txns=txns_after - txns, + ) + return sample, pos + + def _bench( + self, + name: str, + rep: Callable[[str], tuple[Sample, str]], + ) -> None: + """Run `rep` once per side to warm up, then `REPS` times per side in + turn, and record the medians.""" + for side in ("python", "rust"): + rep(side) + samples: dict[str, list[Sample]] = {"python": [], "rust": []} + for _ in range(REPS): + for side in ("python", "rust"): + samples[side].append(rep(side)[0]) + profile_dir = os.environ.get("SYNAPSE_BENCH_PROFILE") + if profile_dir: + for side in ("python", "rust"): + profile = self.profile = cProfile.Profile() + try: + rep(side) + finally: + self.profile = None + slug = re.sub(r"[^a-z0-9]+", "-", name.lower()).strip("-") + profile.dump_stats( + os.path.join( + profile_dir, f"{self._testMethodName}-{slug}-{side}.prof" + ) + ) + for side in ("python", "rust"): + s = samples[side] + self.lines.append( + f"| {name} | {side} " + f"| {statistics.median(x.wall for x in s) * 1000:.1f} " + f"| {statistics.median(x.reactor for x in s) * 1000:.1f} " + f"| {statistics.median(x.process for x in s) * 1000:.1f} " + f"| {statistics.median(x.db for x in s) * 1000:.1f} " + f"| {statistics.median(x.queries for x in s):.0f} " + f"| {statistics.median(x.txns for x in s):.0f} |" + ) + + def _report(self, title: str) -> None: + table = "\n".join( + [ + f"### {title}", + "", + "| scenario | side | wall ms | reactor ms | process ms | db ms " + "| queries | txns |", + "|---|---|---|---|---|---|---|---|", + *self.lines, + "", + ] + ) + print("\n" + table) + out = os.environ.get("SYNAPSE_BENCH_OUT") + if out: + with open(out, "a") as f: + f.write(table + "\n") + + def _make_rooms(self, count: int, messages: int) -> list[str]: + """`count` rooms that user2 created and user1 joined, each with + `messages` messages from user2, a receipt from each user and a name. + """ + count = max(int(count * SCALE), 20) + rooms = [] + for i in range(count): + room_id = self.helper.create_room_as( + self.user2_id, + tok=self.user2_tok, + extra_content={"name": f"room {i}"}, + ) + self.helper.join(room_id, self.user1_id, tok=self.user1_tok) + event_id = "" + for _ in range(messages): + event_id = self.helper.send(room_id, "hi", tok=self.user2_tok)[ + "event_id" + ] + for tok in (self.user1_tok, self.user2_tok): + channel = self.make_request( + "POST", + f"/_matrix/client/v3/rooms/{room_id}/receipt/{ReceiptTypes.READ}/{event_id}", + content={}, + access_token=tok, + ) + self.assertEqual(channel.code, 200, channel.json_body) + rooms.append(room_id) + for room_id in rooms[:10]: + channel = self.make_request( + "PUT", + f"/_matrix/client/v3/user/{self.user1_id}/rooms/{room_id}/tags/u.work", + content={}, + access_token=self.user1_tok, + ) + self.assertEqual(channel.code, 200, channel.json_body) + channel = self.make_request( + "PUT", + f"/_matrix/client/v3/user/{self.user1_id}/account_data/m.direct", + content={self.user2_id: rooms[:5]}, + access_token=self.user1_tok, + ) + self.assertEqual(channel.code, 200, channel.json_body) + return rooms + + def _initial(self, body: JsonDict) -> Callable[[str], tuple[Sample, str]]: + def rep(side: str) -> tuple[Sample, str]: + return self._timed_sync(side, self._new_conn(side), body, None) + + return rep + + def _incremental( + self, + first_body: JsonDict, + body: JsonDict, + before: Callable[[], None] | None = None, + ) -> Callable[[str], tuple[Sample, str]]: + """Requests on one connection per side, after an untimed + `first_body`. `before` runs ahead of each Python request, so both + sides see what it does.""" + conns = {} + for side in ("python", "rust"): + conn_id = self._new_conn(side) + _, pos = self._sync(side, conn_id, first_body, None) + conns[side] = (conn_id, pos) + + def rep(side: str) -> tuple[Sample, str]: + if before is not None and side == "python": + before() + conn_id, pos = conns[side] + sample, pos = self._timed_sync(side, conn_id, body, pos) + conns[side] = (conn_id, pos) + return sample, pos + + return rep + + def _expansion( + self, first_body: JsonDict, body: JsonDict + ) -> Callable[[str], tuple[Sample, str]]: + """The second request of a new connection.""" + + def rep(side: str) -> tuple[Sample, str]: + conn_id = self._new_conn(side) + _, pos = self._sync(side, conn_id, first_body, None) + return self._timed_sync(side, conn_id, body, pos) + + return rep + + def test_200_rooms(self) -> None: + """An account in 200 rooms of 20 messages each.""" + rooms = self._make_rooms(200, 20) + top = rooms[-1] + + def send_one() -> None: + self.helper.send(top, "hi", tok=self.user2_tok) + + for extensions in (False, True): + suffix = ", extensions" if extensions else "" + first = _list_body(19, extensions=extensions) + self._bench(f"initial [0,19]{suffix}", self._initial(first)) + self._bench( + f"incremental, nothing changed{suffix}", + self._incremental(first, first), + ) + self._bench( + f"incremental, one event{suffix}", + self._incremental(first, first, send_one), + ) + self._bench( + f"range [0,19] to [0,99]{suffix}", + self._expansion(first, _list_body(99, extensions=extensions)), + ) + self._report("200 rooms, 20 messages each") + + def test_2000_rooms(self) -> None: + """An account in 2000 rooms of one message each.""" + self._make_rooms(2000, 1) + + for extensions in (False, True): + suffix = ", extensions" if extensions else "" + first = _list_body(19, extensions=extensions) + self._bench(f"initial [0,19]{suffix}", self._initial(first)) + self._bench( + f"incremental, nothing changed{suffix}", + self._incremental(first, first), + ) + wide = _list_body(499, timeline_limit=1, extensions=extensions) + self._bench( + f"incremental [0,19], 500 rooms sent before{suffix}", + self._incremental(wide, first), + ) + self._bench( + f"incremental [0,499], nothing changed{suffix}", + self._incremental(wide, wide), + ) + self._report("2000 rooms, one message each") 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..f00068fdc5 --- /dev/null +++ b/tests/rest/client/sliding_sync/test_rust_lists.py @@ -0,0 +1,587 @@ +# +# 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, + RelationTypes, + RoomTypes, +) +from synapse.handlers.sliding_sync import ( + rust_sliding_sync_requests_counter, + rust_sliding_sync_rooms_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, + StateValues, +) +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_python_rooms(self) -> None: + """Python computes the rooms Rust leaves to it on a copy of their + state, and its changes reach the Rust state.""" + user3_id = self.register_user("user3", "pass") + user3_tok = self.login(user3_id, "pass") + (rust_room,) = self._joined_rooms(1) + python_room = self.helper.create_room_as(self.user2_id, tok=self.user2_tok) + self.helper.join(python_room, self.user1_id, tok=self.user1_tok) + self.helper.join(python_room, user3_id, tok=user3_tok) + self._put_account_data( + AccountDataTypes.IGNORED_USER_LIST, {"ignored_users": {user3_id: {}}} + ) + # Rust leaves a room to Python when the latest reply in a thread on + # its timeline is from an ignored user. + root_id = self._send_message(python_room) + self.helper.send_event( + python_room, + EventTypes.Message, + { + "msgtype": "m.text", + "body": "reply", + "m.relates_to": {"rel_type": RelationTypes.THREAD, "event_id": root_id}, + }, + tok=user3_tok, + ) + self._send_message(rust_room) + sync_body: JsonDict = { + "lists": { + "all": { + "ranges": [[0, 99]], + "required_state": [ + [EventTypes.Create, ""], + [EventTypes.Member, StateValues.LAZY], + ], + "timeline_limit": 3, + } + } + } + fallback = rust_sliding_sync_rooms_counter.labels( + result="fallback", **{SERVER_NAME_LABEL: self.hs.hostname} + ) + before = fallback._value.get() + response = self._sync_both(sync_body) + self.assertEqual(fallback._value.get() - before, 1) + self.assertEqual(set(response["rooms"]), {rust_room, python_room}) + + # The next requests start from the state Python's room wrote. + self._send_message(python_room) + self._send_message(rust_room) + self._sync_both(sync_body) + before = fallback._value.get() + self._sync_both( + { + "lists": { + "all": { + **sync_body["lists"]["all"], + "required_state": [[EventTypes.Member, "*"]], + "timeline_limit": 5, + } + } + } + ) + self.assertEqual(fallback._value.get() - before, 1) + + 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_request.py b/tests/rest/client/sliding_sync/test_rust_request.py new file mode 100644 index 0000000000..73af5a7f38 --- /dev/null +++ b/tests/rest/client/sliding_sync/test_rust_request.py @@ -0,0 +1,563 @@ +# +# 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 Rust parse of the sliding sync request body +(`rust/src/sliding_sync/request.rs`) against pydantic's `SlidingSyncBody`.""" + +import json +import random +from typing import Any +from unittest import SkipTest +from unittest.mock import patch + +from pydantic import ValidationError + +from twisted.internet.testing import MemoryReactor + +from synapse.rest import admin +from synapse.rest.client import login, room, sync +from synapse.server import HomeServer +from synapse.synapse_rust.sliding_sync import SlidingSyncRequest +from synapse.types import ( + Absent, + JsonDict, + RoomStreamToken, + SlidingSyncStreamToken, + StreamToken, +) +from synapse.types.handlers.sliding_sync import SlidingSyncConfig +from synapse.types.rest.client import SlidingSyncBody +from synapse.util.clock import Clock + +import tests.utils +from tests import unittest +from tests.rest.client.sliding_sync.test_sliding_sync import SlidingSyncBase + +# A JSON object as the generator builds it: its entries in order, so that a +# key can repeat. +Entries = list[tuple[str, Any]] + + +class _Object: + def __init__(self, entries: Entries) -> None: + self.entries = entries + + +def _encode(value: Any, rng: random.Random) -> str: + """Encode `value` as JSON text, keeping the repeated keys of an + `_Object` and varying the whitespace and string escapes.""" + space = rng.choice(["", " ", "\n "]) + if isinstance(value, _Object): + items = [ + f"{_encode(key, rng)}:{space}{_encode(item, rng)}" + for key, item in value.entries + ] + return "{" + f",{space}".join(items) + "}" + if isinstance(value, list): + return "[" + f",{space}".join(_encode(item, rng) for item in value) + "]" + return json.dumps(value, ensure_ascii=rng.random() < 0.5) + + +_STATE_PARTS = [ + "*", + "$LAZY", + "$ME", + "", + "m.room.member", + "m.room.name", + "m.space.child", + "@alice:test", + "org.example.é", + "☃", +] +_NAMES = ["a", "b", "list", "", "é", "x" * 64, "☃", "a b"] +_ROOM_IDS = ["!a:test", "!b:test", "!c:test", "!é:test"] + + +class _BodyGenerator: + """Random request bodies, most of which pydantic accepts.""" + + def __init__(self, rng: random.Random) -> None: + self.rng = rng + + def maybe(self, value: Any, p: float = 0.5) -> list[Any]: + """`[value]` with probability `p`, else `[]`: an optional field.""" + return [value] if self.rng.random() < p else [] + + def nullable(self, value: Any) -> Any: + return None if self.rng.random() < 0.15 else value + + def unknown_fields(self) -> Entries: + return self.maybe( + ("org.example.unknown", self.rng.choice([1, "x", None, []])), 0.1 + ) + + def coin(self) -> bool: + return self.rng.random() < 0.5 + + def strings(self, pool: list[str], allow_none: bool = False) -> list[Any]: + values: list[Any] = list(pool) + ([None] if allow_none else []) + return self.rng.sample(values, self.rng.randint(0, min(4, len(values)))) + + def required_state(self) -> list[list[str]]: + return [ + [self.rng.choice(_STATE_PARTS), self.rng.choice(_STATE_PARTS)] + for _ in range(self.rng.randint(0, 5)) + ] + + def room_params(self) -> Entries: + return [ + ("required_state", self.required_state()), + ("timeline_limit", self.rng.choice([0, 1, 10, 1000, -3])), + *self.unknown_fields(), + ] + + def filters(self) -> _Object: + rng = self.rng + fields: Entries = [ + ("is_dm", self.nullable(self.coin())), + ("spaces", self.nullable(self.strings(_ROOM_IDS))), + ("is_encrypted", self.nullable(self.coin())), + ("is_invite", self.nullable(self.coin())), + ( + "room_types", + self.nullable(self.strings(["m.space", "org.example"], True)), + ), + ( + "not_room_types", + self.nullable(self.strings(["m.space", "org.example"], True)), + ), + ("room_name_like", self.nullable(rng.choice(["foo", "", "É"]))), + ("tags", self.nullable(self.strings(["m.favourite", "u.work"]))), + ("not_tags", self.nullable(self.strings(["m.lowpriority", "u.x"]))), + ] + entries = [field for field in fields if rng.random() < 0.4] + return _Object(entries + self.unknown_fields()) + + def sliding_list(self) -> _Object: + rng = self.rng + entries = self.room_params() + if rng.random() < 0.7: + ranges = [ + sorted([rng.randint(0, 200), rng.randint(0, 200)]) + for _ in range(rng.randint(0, 3)) + ] + entries.append(("ranges", self.nullable(ranges))) + entries += self.maybe(("slow_get_all_rooms", self.nullable(self.coin())), 0.2) + entries += self.maybe(("filters", self.nullable(self.filters())), 0.4) + rng.shuffle(entries) + return _Object(entries) + + def keyed(self, keys: list[str], make: Any) -> _Object: + """An object of up to four entries, sometimes with a key repeated.""" + entries = [ + (self.rng.choice(keys), make()) for _ in range(self.rng.randint(0, 4)) + ] + return _Object(entries) + + def lists_and_rooms_extension(self) -> _Object: + entries: Entries = [ + *self.maybe(("enabled", self.nullable(self.coin()))), + *self.maybe(("lists", self.nullable(self.strings(["*", *_NAMES])))), + *self.maybe(("rooms", self.nullable(self.strings(["*", *_ROOM_IDS])))), + ] + return _Object(entries + self.unknown_fields()) + + def extensions(self) -> _Object: + rng = self.rng + candidates: Entries = [ + ( + "to_device", + _Object( + [ + *self.maybe(("enabled", self.nullable(self.coin()))), + *self.maybe(("limit", rng.choice([0, 7, 100, -1]))), + *self.maybe(("since", self.nullable(str(rng.randint(0, 99))))), + ] + ), + ), + ("e2ee", _Object(self.maybe(("enabled", self.nullable(self.coin()))))), + ("account_data", self.lists_and_rooms_extension()), + ("receipts", self.lists_and_rooms_extension()), + ("typing", self.lists_and_rooms_extension()), + ( + "io.element.msc4308.thread_subscriptions", + _Object( + [ + *self.maybe(("enabled", self.nullable(self.coin()))), + *self.maybe(("limit", rng.choice([1, 100]))), + ] + ), + ), + # Only the unstable name counts: this one is ignored. + ("thread_subscriptions", _Object([("enabled", True)])), + ( + "org.matrix.msc4354.sticky_events", + _Object( + [ + *self.maybe(("enabled", self.coin())), + *self.maybe(("limit", rng.choice([0, 5, 100]))), + *self.maybe( + ("since", rng.choice(["sticky_0", "sticky_m4~1.6"])) + ), + ] + ), + ), + ( + "org.matrix.msc4262.profiles", + _Object( + [ + *self.maybe(("enabled", self.coin())), + *self.maybe(("fields", self.strings(["displayname", "m.tz"]))), + ] + ), + ), + ("org.example.unknown_extension", _Object([("enabled", True)])), + ] + entries = [ + (name, self.nullable(value)) + if name not in _NOT_NULLABLE_EXTENSIONS + else (name, value) + for name, value in candidates + if rng.random() < 0.4 + ] + return _Object(entries) + + def body(self) -> _Object: + entries: Entries = [ + *self.maybe(("conn_id", self.nullable(self.rng.choice(["", "room-list"])))), + *self.maybe( + ("lists", self.nullable(self.keyed(_NAMES, self.sliding_list))), 0.8 + ), + *self.maybe( + ( + "room_subscriptions", + self.nullable( + self.keyed(_ROOM_IDS, lambda: _Object(self.room_params())) + ), + ), + ), + *self.maybe(("extensions", self.nullable(self.extensions()))), + *self.unknown_fields(), + ] + self.rng.shuffle(entries) + return _Object(entries) + + +# Pydantic rejects an explicit `null` for these: they are `Absent` or a value. +_NOT_NULLABLE_EXTENSIONS = { + "org.matrix.msc4354.sticky_events", + "org.matrix.msc4262.profiles", +} + + +def _expected(body: SlidingSyncBody) -> JsonDict: + """What `SlidingSyncRequest.body_for_test` should return for `body`.""" + + def room_params(params: SlidingSyncBody.CommonRoomParameters) -> JsonDict: + return { + "required_state": [list(pair) for pair in params.required_state], + "timeline_limit": params.timeline_limit, + } + + def sliding_list(sliding_list: SlidingSyncBody.SlidingSyncList) -> JsonDict: + return { + **room_params(sliding_list), + "ranges": ( + None + if sliding_list.ranges is None + else [list(r) for r in sliding_list.ranges] + ), + "slow_get_all_rooms": sliding_list.slow_get_all_rooms, + "filters": ( + None + if sliding_list.filters is None + else sliding_list.filters.model_dump() + ), + } + + def dump(model: Any) -> Any: + return None if model is None else model.model_dump() + + extensions = None + if body.extensions is not None: + ext = body.extensions + sticky = None + if ext.sticky_events is not Absent: + sticky = { + "enabled": ext.sticky_events.enabled, + "limit": ext.sticky_events.limit, + "since": ( + ext.sticky_events.since.serialise() + if ext.sticky_events.since is not Absent + else None + ), + } + profiles = None + if ext.profiles is not Absent: + profiles = { + "enabled": ext.profiles.enabled, + "fields": ( + list(ext.profiles.fields) + if ext.profiles.fields is not Absent + else None + ), + } + extensions = { + "to_device": dump(ext.to_device), + "e2ee": dump(ext.e2ee), + "account_data": dump(ext.account_data), + "receipts": dump(ext.receipts), + "typing": dump(ext.typing), + "io.element.msc4308.thread_subscriptions": dump(ext.thread_subscriptions), + "org.matrix.msc4354.sticky_events": sticky, + "org.matrix.msc4262.profiles": profiles, + } + + return { + "conn_id": body.conn_id, + "lists": ( + None + if body.lists is None + else [[name, sliding_list(value)] for name, value in body.lists.items()] + ), + "room_subscriptions": ( + None + if body.room_subscriptions is None + else [ + [room_id, room_params(value)] + for room_id, value in body.room_subscriptions.items() + ] + ), + "extensions": extensions, + } + + +def _as_json(value: Any) -> Any: + """`value` with its tuples as lists, for comparing.""" + return json.loads(json.dumps(value)) + + +def _rust_request(body: bytes) -> SlidingSyncRequest: + return SlidingSyncRequest( + body, + user_id="@alice:test", + device_id="DEVICE", + from_room_key=None, + connection_position=None, + to_room_key=RoomStreamToken(stream=1), + to_suffix="0_0", + is_server_admin=False, + return_soft_failed_events=False, + return_policy_server_spammy_events=False, + have_finished_bg_jobs=True, + rooms_to_exclude=set(), + msc4306_enabled=False, + msc4354_enabled=False, + profiles_enabled=False, + ) + + +class RustRequestParseTestCase(unittest.TestCase): + """Rust parses every body pydantic accepts the same way.""" + + def test_generated_bodies(self) -> None: + accepted = 0 + for seed in range(400): + rng = random.Random(seed) + text = _encode(_BodyGenerator(rng).body(), rng) + try: + body = SlidingSyncBody.model_validate(json.loads(text)) + except ValidationError: + continue + accepted += 1 + expected = _as_json(_expected(body)) + self.assertEqual( + _as_json(_rust_request(text.encode()).body_for_test()), + expected, + f"seed {seed}: {text}", + ) + # What `_build_rust_request` sends for a config that was not built + # from a request. + reencoded = body.model_dump_json(by_alias=True, exclude_unset=True) + self.assertEqual( + _as_json(_rust_request(reencoded.encode()).body_for_test()), + expected, + f"seed {seed}, re-encoded: {reencoded}", + ) + # The generator is no use if pydantic rejects what it makes. + self.assertGreater(accepted, 300) + + def test_repeated_list_name(self) -> None: + """A repeated key keeps its first position and its last value, as in + the dict `json.loads` builds.""" + text = ( + '{"lists": {' + '"b": {"required_state": [], "timeline_limit": 1},' + '"a": {"required_state": [], "timeline_limit": 2},' + '"b": {"required_state": [], "timeline_limit": 3}}}' + ) + body = SlidingSyncBody.model_validate(json.loads(text)) + rust = _rust_request(text.encode()).body_for_test() + self.assertEqual(_as_json(rust), _as_json(_expected(body))) + self.assertEqual([name for name, _ in rust["lists"]], ["b", "a"]) + self.assertEqual(rust["lists"][0][1]["timeline_limit"], 3) + + def test_unparseable_body(self) -> None: + """A body Rust does not cover raises `ValueError`, and Python computes + the request.""" + text = ( + '{"lists": {"a": {"required_state": [], "timeline_limit": 1,' + ' "ranges": [[0, 100000000000000000000000]]}}}' + ) + SlidingSyncBody.model_validate(json.loads(text)) + with self.assertRaises(ValueError): + _rust_request(text.encode()) + + def test_tokens_go_together(self) -> None: + with self.assertRaises(ValueError): + SlidingSyncRequest( + b"{}", + user_id="@alice:test", + device_id=None, + from_room_key=RoomStreamToken(stream=1), + connection_position=None, + to_room_key=RoomStreamToken(stream=1), + to_suffix="0", + is_server_admin=False, + return_soft_failed_events=False, + return_policy_server_spammy_events=False, + have_finished_bg_jobs=True, + rooms_to_exclude=set(), + msc4306_enabled=False, + msc4354_enabled=False, + profiles_enabled=False, + ) + + +class RustRequestBuildTestCase(SlidingSyncBase): + """The servlet keeps the body's bytes, and the handler builds the Rust + request from them.""" + + servlets = [ + admin.register_servlets, + login.register_servlets, + room.register_servlets, + sync.register_servlets, + ] + + def default_config(self) -> JsonDict: + config = super().default_config() + config["experimental_features"]["rust_sliding_sync"] = "on" + config["exclude_rooms_from_sync"] = ["!excluded:test"] + return config + + def make_homeserver(self, reactor: MemoryReactor, clock: Clock) -> HomeServer: + if tests.utils.RUST_SLIDING_SYNC_FOR_TESTS not in (None, "on"): + raise SkipTest("SYNAPSE_TEST_RUST_SLIDING_SYNC overrides the config") + hs = self.setup_test_homeserver() + if hs.get_sliding_sync_handler()._rust_sliding_sync is None: + raise SkipTest("needs the native Rust Postgres driver") + return hs + + def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None: + super().prepare(reactor, clock, hs) + self.handler = hs.get_sliding_sync_handler() + self.user = self.register_user("alice", "pass", admin=True) + self.tok = self.login("alice", "pass") + + def _capture_configs(self, text: str, since: str | None = None) -> list[Any]: + """Send `text` as a sliding sync body, and return the arguments of each + `current_sync_for_user` call it led to.""" + calls: list[Any] = [] + original = self.handler.current_sync_for_user + + async def capture( + sync_config: SlidingSyncConfig, + to_token: StreamToken, + from_token: SlidingSyncStreamToken | None = None, + ) -> Any: + calls.append((sync_config, to_token, from_token)) + return await original(sync_config, to_token=to_token, from_token=from_token) + + path = self.sync_endpoint + (f"?pos={since}" if since else "") + with patch.object(self.handler, "current_sync_for_user", side_effect=capture): + channel = self.make_request( + "POST", path, content=text.encode(), access_token=self.tok + ) + self.assertEqual(channel.code, 200, channel.json_body) + return calls + + def test_body_bytes_and_request(self) -> None: + text = ( + '{"lists": {"a": {"ranges": [[0, 9]],' + ' "required_state": [["m.room.name", ""]], "timeline_limit": 3}},' + ' "conn_id": "c"}' + ) + [(sync_config, to_token, from_token)] = self._capture_configs(text) + self.assertEqual(sync_config.body_bytes, text.encode()) + self.assertIsNone(from_token) + + request = self.get_success( + self.handler._build_rust_request(sync_config, from_token, to_token) + ) + self.assertEqual(request.user_id, self.user) + self.assertEqual(request.device_id, sync_config.requester.device_id) + self.assertIsNone(request.from_room_key) + self.assertIsNone(request.connection_position) + self.assertEqual(request.to_room_key, to_token.room_key) + self.assertEqual( + f"{self.get_success(to_token.room_key.to_string(self.store))}_" + f"{request.to_suffix}", + self.get_success(to_token.to_string(self.store)), + ) + self.assertTrue(request.is_server_admin) + self.assertTrue(request.have_finished_bg_jobs) + self.assertEqual(request.rooms_to_exclude, {"!excluded:test"}) + self.assertEqual( + _as_json(request.body_for_test()), + _as_json(_expected(SlidingSyncBody.model_validate(json.loads(text)))), + ) + + def test_incremental_request(self) -> None: + text = '{"lists": {"a": {"required_state": [], "timeline_limit": 1}}}' + [(sync_config, to_token, _)] = self._capture_configs(text) + channel = self.make_request( + "POST", self.sync_endpoint, content=text.encode(), access_token=self.tok + ) + pos = channel.json_body["pos"] + [(sync_config, to_token, from_token)] = self._capture_configs(text, pos) + assert from_token is not None + request = self.get_success( + self.handler._build_rust_request(sync_config, from_token, to_token) + ) + self.assertEqual(request.from_room_key, from_token.stream_token.room_key) + self.assertEqual(request.connection_position, from_token.connection_position) + + def test_config_without_body_bytes(self) -> None: + """A config not built from a request is re-encoded for Rust.""" + text = '{"lists": {"a": {"required_state": [["*", "*"]], "timeline_limit": 2}}}' + [(sync_config, to_token, _)] = self._capture_configs(text) + config = sync_config.model_copy(update={"body_bytes": None}) + request = self.get_success( + self.handler._build_rust_request(config, None, to_token) + ) + self.assertEqual( + _as_json(request.body_for_test()), + _as_json(_expected(SlidingSyncBody.model_validate(json.loads(text)))), + ) + + @property + def store(self) -> Any: + return self.hs.get_datastores().main diff --git a/tests/rest/client/sliding_sync/test_rust_rooms.py b/tests/rest/client/sliding_sync/test_rust_rooms.py new file mode 100644 index 0000000000..6b5b256fa1 --- /dev/null +++ b/tests/rest/client/sliding_sync/test_rust_rooms.py @@ -0,0 +1,727 @@ +# +# 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 rooms of a sliding sync response computed and encoded in Rust +(`rust/src/sliding_sync/room_entry.rs`, `encode.rs`) against Python's. + +Each request goes to two connections of the same user, one served by Python +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 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 +from synapse.api.constants import ( + AccountDataTypes, + EventTypes, + JoinRules, + RelationTypes, +) +from synapse.events import EventBase, StrippedStateEvent +from synapse.handlers.sliding_sync import rust_sliding_sync_rooms_counter +from synapse.metrics import SERVER_NAME_LABEL +from synapse.rest.client import account_data, 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, + SlidingSyncConfig, +) +from synapse.util.clock import Clock + +from tests.rest.client.sliding_sync.test_sliding_sync import SlidingSyncBase +from tests.server import _BlockingRustCalls +from tests.test_utils.event_injection import create_event + + +def _without_age(value: Any) -> Any: + """`value` without the `age` of each event's unsigned data, which moves on + with the clock between the two requests.""" + if isinstance(value, dict): + return { + key: _without_age(item) + for key, item in value.items() + if not (key == "age" and isinstance(item, int)) + } + if isinstance(value, list): + return [_without_age(item) for item in value] + return value + + +def _normalized(rooms: JsonDict) -> JsonDict: + """`rooms` without the ages, and with each room's `required_state` in + order: it is a set, which neither side orders.""" + rooms = _without_age(rooms) + for entry in rooms.values(): + if "required_state" in entry: + entry["required_state"].sort(key=lambda e: (e["type"], e["state_key"])) + return rooms + + +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. `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: ( + timeline_limit, + { + state_type: set(state_keys) + for state_type, state_keys in required_state_map.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 returned.items() + }, + set(invalidated), + ) + 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",), + [ + (True,), + (False,), + ], + class_name_func=lambda cls, + num, + params_dict: f"{cls.__name__}_{'new' if params_dict['use_new_tables'] else 'fallback'}", +) +class SlidingSyncRustRoomsTestCase(SlidingSyncBase): + maxDiff = None + + servlets = [ + synapse.rest.admin.register_servlets, + account_data.register_servlets, + knock.register_servlets, + login.register_servlets, + room.register_servlets, + sync.register_servlets, + ] + + 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") + # 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} + + 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] = 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") + self.user3_id = self.register_user("user3", "pass") + self.user3_tok = self.login(self.user3_id, "pass") + + def _rust_rooms_computed(self) -> float: + return rust_sliding_sync_rooms_counter.labels( + result="rust", **{SERVER_NAME_LABEL: self.hs.hostname} + )._value.get() + + def _sync_both(self, sync_body: JsonDict) -> JsonDict: + """Sync `sync_body` on the Python and the Rust connection, assert the + two `rooms` and the two sides' changes to the connection state are + equal, and return Rust's response.""" + 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_rooms_computed() + response, self.pos[side] = self.do_sync( + {**sync_body, "conn_id": side}, + since=self.pos[side], + tok=self.user1_tok, + ) + computed = self._rust_rooms_computed() - before + 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"]) + rust_rooms = _normalized(responses["rust"]["rooms"]) + self.assertEqual(python_rooms.keys(), rust_rooms.keys()) + for room_id, python_room in python_rooms.items(): + self.assertEqual(python_room, rust_rooms[room_id], room_id) + self.assertEqual( + _unordered_lists(responses["python"]), + _unordered_lists(responses["rust"]), + ) + self.assertEqual( + _connection_changes(self.recorded["python"]), + _connection_changes(self.recorded["rust"]), + ) + return responses["rust"] + + def _named_room(self) -> str: + """A room with a name, an avatar, a topic and a busy timeline: an edit, + a thread, a redaction.""" + room_id = self.helper.create_room_as(self.user2_id, tok=self.user2_tok) + self.helper.send_state( + room_id, EventTypes.Name, {"name": "named"}, tok=self.user2_tok + ) + self.helper.send_state( + room_id, + EventTypes.RoomAvatar, + {"url": "mxc://test/avatar"}, + tok=self.user2_tok, + ) + self.helper.send_state( + room_id, EventTypes.Topic, {"topic": "a topic"}, tok=self.user2_tok + ) + self.helper.join(room_id, self.user1_id, tok=self.user1_tok) + root = self.helper.send(room_id, "root", tok=self.user2_tok)["event_id"] + self.helper.send_event( + room_id, + "m.room.message", + { + "msgtype": "m.text", + "body": "* edited", + "m.new_content": {"msgtype": "m.text", "body": "edited"}, + "m.relates_to": {"rel_type": RelationTypes.REPLACE, "event_id": root}, + }, + tok=self.user2_tok, + ) + doomed = self.helper.send(room_id, "doomed", tok=self.user1_tok)["event_id"] + channel = self.make_request( + "PUT", + f"/rooms/{room_id}/redact/{doomed}/txn1", + content={"reason": "test"}, + access_token=self.user1_tok, + ) + self.assertEqual(channel.code, 200, channel.json_body) + self.helper.send_event( + room_id, + "m.room.message", + { + "msgtype": "m.text", + "body": "in thread", + "m.relates_to": {"rel_type": RelationTypes.THREAD, "event_id": root}, + }, + tok=self.user1_tok, + ) + self.helper.send(room_id, "last", tok=self.user2_tok) + return room_id + + def _rooms(self) -> dict[str, str]: + """One room of each kind the room computation tells apart.""" + rooms = {"named": self._named_room()} + + # No name: heroes, with their display names and avatars. + unnamed = self.helper.create_room_as(self.user2_id, tok=self.user2_tok) + self.helper.join(unnamed, self.user3_id, tok=self.user3_tok) + self.helper.send_state( + unnamed, + EventTypes.Member, + {"membership": "join", "displayname": "Three", "avatar_url": "mxc://a/b"}, + state_key=self.user3_id, + tok=self.user3_tok, + ) + self.helper.join(unnamed, self.user1_id, tok=self.user1_tok) + self.helper.send(unnamed, "hello", tok=self.user3_tok) + rooms["unnamed"] = unnamed + + # An empty name counts as none. + empty_name = self.helper.create_room_as(self.user2_id, tok=self.user2_tok) + self.helper.send_state( + empty_name, EventTypes.Name, {"name": ""}, tok=self.user2_tok + ) + self.helper.join(empty_name, self.user1_id, tok=self.user1_tok) + rooms["empty_name"] = empty_name + + invited = self.helper.create_room_as(self.user2_id, tok=self.user2_tok) + self.helper.invite(invited, self.user2_id, self.user1_id, tok=self.user2_tok) + rooms["invited"] = invited + + knocked = self.helper.create_room_as( + self.user2_id, tok=self.user2_tok, room_version="7" + ) + self.helper.send_state( + knocked, + EventTypes.JoinRules, + {"join_rule": JoinRules.KNOCK}, + tok=self.user2_tok, + ) + self.helper.knock(knocked, self.user1_id, tok=self.user1_tok) + rooms["knocked"] = knocked + + left = self.helper.create_room_as(self.user2_id, tok=self.user2_tok) + self.helper.join(left, self.user1_id, tok=self.user1_tok) + self.helper.send(left, "before leaving", tok=self.user2_tok) + self.helper.leave(left, self.user1_id, tok=self.user1_tok) + self.helper.send(left, "after leaving", tok=self.user2_tok) + rooms["left"] = left + + banned = self.helper.create_room_as(self.user2_id, tok=self.user2_tok) + self.helper.join(banned, self.user1_id, tok=self.user1_tok) + self.helper.ban(banned, self.user2_id, self.user1_id, tok=self.user2_tok) + rooms["banned"] = banned + + rooms["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, + ) + + remote_invite, _ = self._create_remote_invite_room_for_user( + self.user1_id, + [ + StrippedStateEvent( + type=EventTypes.Name, + state_key="", + sender="@inviter:remote_server", + content={"name": "remote"}, + ) + ], + ) + rooms["remote_invite"] = remote_invite + return rooms + + def _list(self, timeline_limit: int, required_state: list[list[str]]) -> JsonDict: + return { + "lists": { + "all": { + "ranges": [[0, 99]], + "required_state": required_state, + "timeline_limit": timeline_limit, + } + } + } + + def test_initial(self) -> None: + """Initial requests with several configs.""" + rooms = self._rooms() + bodies = [ + self._list( + 2, + [ + [EventTypes.Name, ""], + [EventTypes.Member, "$LAZY"], + [EventTypes.Member, "$ME"], + ], + ), + self._list(20, [["*", "*"]]), + self._list(0, []), + # Back to the thread root, so that it has bundled aggregations. + self._list(6, [[EventTypes.Member, "*"], [EventTypes.Topic, ""]]), + ] + for body in bodies: + self.pos = {"python": None, "rust": None} + response = self._sync_both(body) + entries = response["rooms"] + for kind in ("named", "unnamed", "invited", "knocked", "dm"): + self.assertIn(rooms[kind], entries) + self.assertIn("heroes", entries[rooms["unnamed"]]) + self.assertIn("invite_state", entries[rooms["invited"]]) + self.assertIn("invite_state", entries[rooms["remote_invite"]]) + + # The last request covered the bundles and the redaction. + timeline = entries[rooms["named"]]["timeline"] + self.assertTrue(entries[rooms["named"]]["limited"]) + self.assertTrue(any("m.relations" in e["unsigned"] for e in timeline)) + self.assertTrue(any("redacted_because" in e["unsigned"] for e in timeline)) + + def test_incremental(self) -> None: + """Incremental requests: new events, state changes, memberships, and a + changed config.""" + rooms = self._rooms() + body = self._list( + 3, + [ + [EventTypes.Name, ""], + [EventTypes.Member, "$LAZY"], + ], + ) + self._sync_both(body) + + # Nothing happened. + response = self._sync_both(body) + self.assertEqual(response["rooms"], {}) + + self.helper.send(rooms["named"], "new message", tok=self.user2_tok) + self.helper.send_state( + rooms["unnamed"], EventTypes.Name, {"name": "now named"}, tok=self.user2_tok + ) + self.helper.leave(rooms["dm"], self.user2_id, tok=self.user2_tok) + self.helper.join(rooms["named"], self.user3_id, tok=self.user3_tok) + self.helper.send(rooms["named"], "from three", tok=self.user3_tok) + self.helper.join(rooms["invited"], self.user1_id, tok=self.user1_tok) + self.helper.leave(rooms["empty_name"], self.user1_id, tok=self.user1_tok) + response = self._sync_both(body) + self.assertIn(rooms["named"], response["rooms"]) + + # The client asks for more: more state, and a longer timeline. + body = self._list( + 10, + [ + [EventTypes.Name, ""], + [EventTypes.Topic, ""], + [EventTypes.Member, "$LAZY"], + [EventTypes.Member, self.user2_id], + ], + ) + self._sync_both(body) + + # And back to lazy members only, after more traffic. + self.helper.send(rooms["named"], "more", tok=self.user1_tok) + self.helper.send(rooms["unnamed"], "more", tok=self.user3_tok) + self._sync_both(self._list(1, [[EventTypes.Member, "$LAZY"]])) + + def test_room_subscription(self) -> None: + """A room subscription on top of a list, in a room the list leaves out.""" + rooms = self._rooms() + body: JsonDict = { + **self._list(1, [[EventTypes.Name, ""]]), + "room_subscriptions": { + rooms["named"]: { + "required_state": [["*", "*"]], + "timeline_limit": 4, + } + }, + } + body["lists"]["all"]["ranges"] = [[0, 0]] + self._sync_both(body) + self.helper.send(rooms["named"], "new", tok=self.user2_tok) + self._sync_both(body) + + def test_incremental_state_and_membership(self) -> None: + """Incremental requests after an avatar change, a ban, a rejected + remote invite and a state reset out of a room.""" + rooms = self._rooms() + + # A room to be state reset out of: a join rules event that points + # back to before user1 joined undoes the join. + reset = self.helper.create_room_as( + self.user2_id, is_public=True, tok=self.user2_tok + ) + before_join = self.helper.send(reset, "before the join", tok=self.user2_tok) + self.helper.join(reset, self.user1_id, tok=self.user1_tok) + + body = self._list( + 3, + [ + [EventTypes.Name, ""], + [EventTypes.RoomAvatar, ""], + [EventTypes.Member, "$ME"], + [EventTypes.Member, "$LAZY"], + ], + ) + self._sync_both(body) + + self.helper.send_state( + rooms["named"], + EventTypes.RoomAvatar, + {"url": "mxc://test/new_avatar"}, + tok=self.user2_tok, + ) + self.helper.ban( + rooms["unnamed"], self.user2_id, self.user1_id, tok=self.user2_tok + ) + self.helper.leave(rooms["remote_invite"], self.user1_id, tok=self.user1_tok) + + persistence = self.hs.get_storage_controllers().persistence + assert persistence is not None + join_rules, context = self.get_success( + create_event( + self.hs, + prev_event_ids=[before_join["event_id"]], + type=EventTypes.JoinRules, + state_key="", + content={"join_rule": JoinRules.INVITE}, + sender=self.user2_id, + room_id=reset, + room_version=self.get_success( + self.hs.get_datastores().main.get_room_version_id(reset) + ), + ) + ) + self.get_success(persistence.persist_event(join_rules, context)) + users_in_room = self.get_success( + self.hs.get_datastores().main.get_users_in_room(reset) + ) + self.assertEqual(users_in_room, [self.user2_id]) + self.helper.send_state( + reset, EventTypes.Name, {"name": "after the reset"}, tok=self.user2_tok + ) + + response = self._sync_both(body) + for kind in ("named", "unnamed"): + self.assertIn(rooms[kind], response["rooms"]) + self.assertIn(reset, response["rooms"]) + self.assertIn("avatar", response["rooms"][rooms["named"]]) + + def test_multi_writer_prev_batch(self) -> None: + """A `prev_batch` that names a writer: a timeline with no new events + ends at the from token, whose room key names another writer.""" + rooms = self._rooms() + body = self._list(2, [[EventTypes.Name, ""]]) + self._sync_both(body) + self._name_other_writer_in_pos() + + # More state for the named room, which has no new events. + body = self._list(2, [[EventTypes.Name, ""], [EventTypes.Topic, ""]]) + response = self._sync_both(body) + named = response["rooms"][rooms["named"]] + self.assertNotIn("timeline", named) + self.assertTrue(named["prev_batch"].startswith("m"), named["prev_batch"]) + + def test_writer_added_after_instance_map_cached(self) -> None: + """Stored tokens that name a writer added after Rust cached the + `instance_map` table.""" + rooms = self._rooms() + body = self._list(2, [[EventTypes.Name, ""]]) + self._sync_both(body) + self._name_other_writer_in_pos() + + # Rooms outside the narrower range go unsent, and are stored with the + # from token, which names the writer. + for kind in ("named", "unnamed", "empty_name"): + self.helper.send(rooms[kind], "again", tok=self.user2_tok) + narrow = {"lists": {"all": {**body["lists"]["all"], "ranges": [[0, 0]]}}} + self._sync_both(narrow) + self._sync_both(body) + + def _name_other_writer_in_pos(self) -> None: + """Rewrite both sides' `pos` so that its room key names another + writer.""" + # Move the stream on, so that the token's position for the other + # writer is not in the future. + self.helper.create_room_as(self.user2_id, tok=self.user2_tok) + writer_id = self.get_success( + self.hs.get_datastores().main.get_id_for_instance("other_writer") + ) + for side, pos in self.pos.items(): + assert pos is not None + connection_position, stream_token = pos.split("/", 1) + room_key, rest = stream_token.split("_", 1) + self.assertTrue(room_key.startswith("s"), room_key) + stream = int(room_key[1:]) + self.pos[side] = ( + f"{connection_position}/m{stream}~{writer_id}.{stream + 1}_{rest}" + ) + + def test_relations_change(self) -> None: + """Bundled aggregations after the relations Rust cached change: a + second edit, a second threaded reply and a reference.""" + room_id = self._named_room() + # Only a limited timeline gets bundles. Six events reach back to the + # thread root. + response = self._sync_both(self._list(6, [[EventTypes.Name, ""]])) + root = next( + event + for event in response["rooms"][room_id]["timeline"] + if event["content"].get("body") == "root" + ) + root_id = root["event_id"] + self.assertEqual( + root["unsigned"]["m.relations"][RelationTypes.THREAD]["count"], 1 + ) + + for content in ( + { + "msgtype": "m.text", + "body": "* edited again", + "m.new_content": {"msgtype": "m.text", "body": "edited again"}, + "m.relates_to": { + "rel_type": RelationTypes.REPLACE, + "event_id": root_id, + }, + }, + { + "msgtype": "m.text", + "body": "in thread again", + "m.relates_to": {"rel_type": RelationTypes.THREAD, "event_id": root_id}, + }, + { + "msgtype": "m.text", + "body": "a reference", + "m.relates_to": { + "rel_type": RelationTypes.REFERENCE, + "event_id": root_id, + }, + }, + ): + self.helper.send_event( + room_id, "m.room.message", content, tok=self.user2_tok + ) + + self.pos = {"python": None, "rust": None} + response = self._sync_both(self._list(9, [[EventTypes.Name, ""]])) + root = next( + event + for event in response["rooms"][room_id]["timeline"] + if event["event_id"] == root_id + ) + relations = root["unsigned"]["m.relations"] + self.assertEqual(relations[RelationTypes.THREAD]["count"], 2) + self.assertEqual( + relations[RelationTypes.REPLACE]["content"]["body"], "* edited again" + ) + self.assertEqual(len(relations[RelationTypes.REFERENCE]["chunk"]), 1) + + def test_ignore_after_sync(self) -> None: + """Ignoring a user after Rust cached the user's ignore list.""" + room_id = self._named_room() + body = self._list(20, [[EventTypes.Name, ""]]) + response = self._sync_both(body) + senders = {e["sender"] for e in response["rooms"][room_id]["timeline"]} + self.assertIn(self.user2_id, senders) + + channel = self.make_request( + "PUT", + f"/user/{self.user1_id}/account_data/{AccountDataTypes.IGNORED_USER_LIST}", + content={"ignored_users": {self.user2_id: {}}}, + access_token=self.user1_tok, + ) + self.assertEqual(channel.code, 200, channel.json_body) + + self.pos = {"python": None, "rust": None} + response = self._sync_both(body) + # Their state events stay. + senders = { + e["sender"] + for e in response["rooms"][room_id]["timeline"] + if "state_key" not in e + } + self.assertNotIn(self.user2_id, senders) + + def test_unsigned_callbacks(self) -> None: + """Module callbacks add to the unsigned data of Rust rooms' events.""" + + async def add_extra_fields(event: EventBase) -> JsonDict: + return {"io.test.id": event.event_id} + + self.hs.get_event_client_serializer().register_add_extra_fields_to_unsigned_client_event_callback( + add_extra_fields + ) + self._named_room() + response = self._sync_both(self._list(3, [["*", "*"]])) + (room_entry,) = response["rooms"].values() + for event in room_entry["timeline"] + room_entry["required_state"]: + self.assertEqual(event["unsigned"]["io.test.id"], event["event_id"]) diff --git a/tests/server.py b/tests/server.py index 25de86370b..1084385fac 100644 --- a/tests/server.py +++ b/tests/server.py @@ -112,6 +112,9 @@ from synapse.storage.engines import BaseDatabaseEngine, create_engine from synapse.storage.prepare_database import prepare_database from synapse.synapse_rust.persist import Persister as RustPersister from synapse.synapse_rust.runtime import RustRuntime, set_virtual_time_msec +from synapse.synapse_rust.sliding_sync import ( + SlidingSyncHandler as RustSlidingSyncHandler, +) from synapse.synapse_rust.state import ( StateResolutionHandler as RustStateResolutionHandler, ) @@ -131,6 +134,8 @@ from tests.utils import ( POSTGRES_USER, RUST_PERSIST_FOR_TESTS, RUST_PERSIST_STRICT_FOR_TESTS, + RUST_SLIDING_SYNC_FOR_TESTS, + RUST_SLIDING_SYNC_STRICT_FOR_TESTS, RUST_STATE_RES_FOR_TESTS, RUST_SYNC_FOR_TESTS, RUST_SYNC_STRICT_FOR_TESTS, @@ -1345,6 +1350,13 @@ def setup_test_homeserver( RustEventPersistenceLevel, RUST_PERSIST_FOR_TESTS ) + if RUST_SLIDING_SYNC_FOR_TESTS: + if RUST_SLIDING_SYNC_FOR_TESTS not in ("off", "on"): + raise ValueError("SYNAPSE_TEST_RUST_SLIDING_SYNC must be off or on") + config.experimental.rust_sliding_sync = cast( + Literal["off", "on"], RUST_SLIDING_SYNC_FOR_TESTS + ) + if USE_POSTGRES_FOR_TESTS: test_db = "synapse_test_%s" % uuid.uuid4().hex @@ -1552,6 +1564,17 @@ def setup_test_homeserver( ) sync_handler._rust_sync_strict = RUST_SYNC_STRICT_FOR_TESTS + if RUST_SLIDING_SYNC_FOR_TESTS: + # Likewise for the Rust sliding sync handler. + sliding_sync_handler = hs.get_sliding_sync_handler() + if sliding_sync_handler._rust_sliding_sync is not None: + sliding_sync_handler._rust_sliding_sync = cast( + Any, _BlockingRustCalls(sliding_sync_handler._rust_sliding_sync) + ) + sliding_sync_handler._rust_sliding_sync_strict = ( + RUST_SLIDING_SYNC_STRICT_FOR_TESTS + ) + # Likewise for event loads, when they run in Rust: code that loads events # in the background would otherwise not have finished when a test checks. # `GetEventCancellationTestCase` exercises the deferreds. @@ -1623,7 +1646,7 @@ class _BlockingRustStateResolution: class _BlockingRustCalls: - """Wraps the Rust `SyncHandler`, `RustStore` or `Persister` so that each of its entry + """Wraps a Rust handler, the `RustStore` or the `Persister` so that each of its entry points blocks, for the reason `_BlockingRustStateResolution` gives. Every entry point `name` that returns a Deferred has a `name_blocking` @@ -1631,7 +1654,10 @@ class _BlockingRustCalls: that twin's result in an already-fired Deferred. """ - def __init__(self, handler: "RustSyncHandler | RustStore | RustPersister") -> None: + def __init__( + self, + handler: "RustSyncHandler | RustSlidingSyncHandler | RustStore | RustPersister", + ) -> None: self._handler = handler def __getattr__(self, name: str) -> Any: diff --git a/tests/storage/test_sliding_sync_store.py b/tests/storage/test_sliding_sync_store.py new file mode 100644 index 0000000000..bed2b7c877 --- /dev/null +++ b/tests/storage/test_sliding_sync_store.py @@ -0,0 +1,1143 @@ +# +# 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 Rust sliding sync store (`rust/src/storage/sliding_sync_store.rs`) +against the Python store methods it ports, on the same database. These need +the native Rust Postgres driver: + + SYNAPSE_TEST_RUST_DRIVER=1 POSTGRES=1 run-trial.sh tests.storage.test_sliding_sync_store + +The two connection-state transactions are compared by running the same +operations through each on two connections of one device, and diffing the +`sliding_sync_connection_*` rows of the two with their IDs replaced by what +they stand for. +""" + +import random +from typing import Any +from unittest import SkipTest +from unittest.mock import patch + +from immutabledict import immutabledict +from parameterized import parameterized + +from twisted.internet.testing import MemoryReactor + +from synapse.api.constants import EventContentFields, EventTypes, RoomTypes +from synapse.api.errors import SlidingSyncUnknownPosition +from synapse.rest import admin +from synapse.rest.client import login, room +from synapse.server import HomeServer +from synapse.storage.databases.main.state import ROOM_UNKNOWN_SENTINEL +from synapse.storage.roommember import RoomsForUserSlidingSync +from synapse.storage.rust_pool import RustConnectionPool +from synapse.synapse_rust.sliding_sync import SlidingSyncStoreForTest +from synapse.types import ( + MultiWriterStreamToken, + PersistedEventPosition, + RoomStreamToken, + StreamKeyType, + StreamToken, +) +from synapse.types.handlers.sliding_sync import ( + HaveSentRoom, + MutablePerConnectionState, + PerConnectionState, + RoomLazyMembershipChanges, + RoomSyncConfig, +) +from synapse.util.clock import Clock + +from tests import unittest +from tests.sync.fixtures import SEEDS, WRITERS + +DEVICE = "DEVICE" + + +def _snapshot(room: RoomsForUserSlidingSync) -> tuple[Any, ...]: + """A Python snapshot in the Rust test form.""" + return ( + room.room_id, + room.sender, + room.membership, + room.event_id, + room.room_version_id, + room.event_pos.instance_name, + room.event_pos.stream, + room.has_known_state, + room.room_type, + bool(room.is_encrypted), + ) + + +def _position(position: PersistedEventPosition | None) -> tuple[str, int] | None: + if position is None: + return None + # A null instance name reads as "master" on the Rust side. + return (position.instance_name or "master", position.stream) + + +class SlidingSyncStoreTestCase(unittest.HomeserverTestCase): + """Each test runs a Rust query and the Python store method against the + same fixture and compares.""" + + servlets = [ + admin.register_servlets, + room.register_servlets, + login.register_servlets, + ] + + def make_homeserver(self, reactor: MemoryReactor, clock: Clock) -> HomeServer: + hs = self.setup_test_homeserver() + if not isinstance( + hs.get_datastores().main.db_pool._db_pool, RustConnectionPool + ): + raise SkipTest("needs the native Rust Postgres driver") + return hs + + def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None: + self.store = hs.get_datastores().main + db_pool = self.store.db_pool._db_pool + assert isinstance(db_pool, RustConnectionPool) + assert db_pool._pool is not None + rust_handler = hs.get_rust_handlers().sliding_sync + assert rust_handler is not None + self.rust = SlidingSyncStoreForTest(db_pool._pool, rust_handler) + + self.alice = self.register_user("alice", "pass") + self.alice_tok = self.login("alice", "pass") + self.bob = self.register_user("bob", "pass") + self.bob_tok = self.login("bob", "pass") + self.room_id = self.helper.create_room_as(self.alice, tok=self.alice_tok) + self.helper.join(self.room_id, self.bob, tok=self.bob_tok) + + # Fixtures. + + def _forget(self, room_id: str, tok: str) -> None: + channel = self.make_request( + "POST", f"/_matrix/client/r0/rooms/{room_id}/forget", {}, access_token=tok + ) + self.assertEqual(channel.code, 200, channel.json_body) + + def _bob_memberships(self) -> list[str]: + """Rooms with every kind of membership snapshot for Bob: joined, + invited, left himself, kicked, banned, forgotten, a space, an + encrypted room. Returns them.""" + rooms = [self.room_id] + invited = self.helper.create_room_as(self.alice, tok=self.alice_tok) + self.helper.invite(invited, self.alice, self.bob, tok=self.alice_tok) + rooms.append(invited) + left = self.helper.create_room_as(self.alice, tok=self.alice_tok) + self.helper.join(left, self.bob, tok=self.bob_tok) + self.helper.leave(left, self.bob, tok=self.bob_tok) + rooms.append(left) + kicked = self.helper.create_room_as(self.alice, tok=self.alice_tok) + self.helper.join(kicked, self.bob, tok=self.bob_tok) + self.helper.change_membership( + kicked, self.alice, self.bob, "leave", tok=self.alice_tok + ) + rooms.append(kicked) + banned = self.helper.create_room_as(self.alice, tok=self.alice_tok) + self.helper.join(banned, self.bob, tok=self.bob_tok) + self.helper.ban(banned, self.alice, self.bob, tok=self.alice_tok) + rooms.append(banned) + forgotten = self.helper.create_room_as(self.alice, tok=self.alice_tok) + self.helper.join(forgotten, self.bob, tok=self.bob_tok) + self.helper.leave(forgotten, self.bob, tok=self.bob_tok) + self._forget(forgotten, self.bob_tok) + rooms.append(forgotten) + space = self.helper.create_room_as( + self.alice, + tok=self.alice_tok, + extra_content={"creation_content": {"type": RoomTypes.SPACE}}, + ) + self.helper.join(space, self.bob, tok=self.bob_tok) + rooms.append(space) + encrypted = self.helper.create_room_as(self.alice, tok=self.alice_tok) + self.helper.send_state( + encrypted, + EventTypes.RoomEncryption, + {EventContentFields.ENCRYPTION_ALGORITHM: "m.megolm.v1.aes-sha2"}, + tok=self.alice_tok, + ) + self.helper.join(encrypted, self.bob, tok=self.bob_tok) + rooms.append(encrypted) + return rooms + + def _scatter_writers(self, rng: random.Random, user_id: str) -> None: + """Pretend the user's memberships were persisted by random writers, + in every table the queries take an instance name from.""" + rows = self.get_success( + self.store.db_pool.execute( + "snapshots", + "SELECT room_id FROM sliding_sync_membership_snapshots WHERE user_id = ?", + user_id, + ) + ) + for (room_id,) in rows: + self.get_success( + self.store.db_pool.simple_update_one( + "sliding_sync_membership_snapshots", + {"room_id": room_id, "user_id": user_id}, + {"event_instance_name": rng.choice(WRITERS)}, + ) + ) + rows = self.get_success( + self.store.db_pool.execute( + "deltas", + "SELECT stream_id, room_id, type, state_key" + " FROM current_state_delta_stream WHERE state_key = ?", + user_id, + ) + ) + for stream_id, room_id, event_type, state_key in rows: + self.get_success( + self.store.db_pool.simple_update( + "current_state_delta_stream", + { + "stream_id": stream_id, + "room_id": room_id, + "type": event_type, + "state_key": state_key, + }, + {"instance_name": rng.choice(WRITERS + [None])}, + desc="scatter", + ) + ) + rows = self.get_success( + self.store.db_pool.execute( + "events", + "SELECT event_id FROM room_memberships WHERE user_id = ?", + user_id, + ) + ) + for (event_id,) in rows: + self.get_success( + self.store.db_pool.simple_update_one( + "events", + {"event_id": event_id}, + {"instance_name": rng.choice(WRITERS)}, + ) + ) + # The changes behind the caches' backs. + self.store.get_sliding_sync_rooms_for_user_from_membership_snapshots.invalidate_all() + self.store._get_event_cache.clear() + + def _stream_range(self) -> tuple[int, int]: + rows = self.get_success( + self.store.db_pool.execute( + "range", "SELECT MIN(stream_ordering), MAX(stream_ordering) FROM events" + ) + ) + return rows[0][0], rows[0][1] + + def _random_token(self, rng: random.Random, lo: int, hi: int) -> RoomStreamToken: + stream = rng.randint(lo - 2, hi + 2) + instance_map = { + name: stream + rng.randint(1, 6) + for name in rng.sample(WRITERS, rng.randint(0, len(WRITERS))) + } + return RoomStreamToken(stream=stream, instance_map=immutabledict(instance_map)) + + # Membership snapshots. + + def test_rooms_from_snapshots(self) -> None: + self._bob_memberships() + for user in (self.alice, self.bob, "@nobody:test"): + expected = self.get_success( + self.store.get_sliding_sync_rooms_for_user_from_membership_snapshots( + user + ) + ) + self.assertEqual( + self.rust.rooms_from_snapshots(user), + {room_id: _snapshot(room) for room_id, room in expected.items()}, + user, + ) + + # A membership change invalidates the Python cache, and the Rust one + # with it. + other = self.helper.create_room_as(self.alice, tok=self.alice_tok) + self.helper.join(other, self.bob, tok=self.bob_tok) + self.helper.leave(self.room_id, self.bob, tok=self.bob_tok) + expected = self.get_success( + self.store.get_sliding_sync_rooms_for_user_from_membership_snapshots( + self.bob + ) + ) + self.assertIn(other, expected) + self.assertNotIn(self.room_id, expected) + self.assertEqual( + self.rust.rooms_from_snapshots(self.bob), + {room_id: _snapshot(room) for room_id, room in expected.items()}, + ) + + def test_rooms_from_snapshots_skips_unknown_versions(self) -> None: + other = self.helper.create_room_as(self.alice, tok=self.alice_tok) + self.helper.join(other, self.bob, tok=self.bob_tok) + self.get_success( + self.store.db_pool.simple_update_one( + "rooms", {"room_id": other}, {"room_version": "unknown-version"} + ) + ) + self.store.get_sliding_sync_rooms_for_user_from_membership_snapshots.invalidate_all() + expected = self.get_success( + self.store.get_sliding_sync_rooms_for_user_from_membership_snapshots( + self.bob + ) + ) + self.assertNotIn(other, expected) + self.assertEqual( + self.rust.rooms_from_snapshots(self.bob), + {room_id: _snapshot(room) for room_id, room in expected.items()}, + ) + # The single and batch reads keep it. + kept = self.get_success( + self.store.get_sliding_sync_room_for_user(self.bob, other) + ) + assert kept is not None + self.assertEqual( + self.rust.snapshot_room(self.bob, other), + _snapshot(kept), + ) + + @parameterized.expand(SEEDS) + def test_self_leave_rooms_after(self, seed: int) -> None: + rng = random.Random(seed) + self._bob_memberships() + for _ in range(3): + room_id = self.helper.create_room_as(self.alice, tok=self.alice_tok) + self.helper.join(room_id, self.bob, tok=self.bob_tok) + self.helper.leave(room_id, self.bob, tok=self.bob_tok) + self._scatter_writers(rng, self.bob) + lo, hi = self._stream_range() + for _ in range(10): + to_token = self._random_token(rng, lo, hi) + expected = self.get_success( + self.store.get_sliding_sync_self_leave_rooms_after_to_token( + self.bob, + StreamToken.START.copy_and_replace(StreamKeyType.ROOM, to_token), + ) + ) + self.assertEqual( + self.rust.self_leave_rooms_after(self.bob, to_token), + {room_id: _snapshot(room) for room_id, room in expected.items()}, + to_token, + ) + + def test_snapshot_rooms(self) -> None: + rooms = self._bob_memberships() + ["!unknown:test"] + for room_id in rooms: + expected = self.get_success( + self.store.get_sliding_sync_room_for_user(self.bob, room_id) + ) + self.assertEqual( + self.rust.snapshot_room(self.bob, room_id), + None if expected is None else _snapshot(expected), + room_id, + ) + expected_batch = self.get_success( + self.store.get_sliding_sync_room_for_user_batch(self.bob, rooms) + ) + self.assertEqual( + self.rust.snapshot_rooms(self.bob, rooms), + {room_id: _snapshot(room) for room_id, room in expected_batch.items()}, + ) + self.assertEqual(self.rust.snapshot_rooms(self.bob, []), {}) + + # Membership changes. + + def _random_memberships(self, rng: random.Random) -> list[str]: + """Random joins, leaves, display name changes, kicks and invites of + Bob in a few rooms.""" + rooms = [self.room_id] + [ + self.helper.create_room_as(self.alice, tok=self.alice_tok) for _ in range(3) + ] + for i in range(14): + room_id = rng.choice(rooms) + membership = self.get_success( + self.store.get_local_current_membership_for_user_in_room( + self.bob, room_id + ) + )[0] + roll = rng.random() + if membership == "join" and roll < 0.3: + self.helper.leave(room_id, self.bob, tok=self.bob_tok) + elif membership == "join" and roll < 0.45: + self.helper.change_membership( + room_id, self.alice, self.bob, "leave", tok=self.alice_tok + ) + elif membership == "join": + self.helper.send_state( + room_id, + EventTypes.Member, + {"membership": "join", "displayname": str(i)}, + state_key=self.bob, + tok=self.bob_tok, + ) + elif membership != "invite" and roll < 0.3: + self.helper.invite(room_id, self.alice, self.bob, tok=self.alice_tok) + else: + self.helper.join(room_id, self.bob, tok=self.bob_tok) + return rooms + + def _null_event_deltas(self, rng: random.Random, rooms: list[str]) -> None: + """Add rows without an event to Bob's membership deltas, as the server + leaving a room or a state reset write them, at random positions. + Their previous event is one of Bob's memberships of the room, or + none.""" + lo, hi = self._stream_range() + for _ in range(6): + room_id = rng.choice(rooms) + rows = self.get_success( + self.store.db_pool.execute( + "memberships", + "SELECT event_id FROM room_memberships" + " WHERE user_id = ? AND room_id = ? ORDER BY event_id", + self.bob, + room_id, + ) + ) + prev_event_id = rng.choice([row[0] for row in rows] + [None]) + self.get_success( + self.store.db_pool.simple_insert( + "current_state_delta_stream", + { + "stream_id": rng.randint(lo, hi), + "room_id": room_id, + "type": EventTypes.Member, + "state_key": self.bob, + "event_id": None, + "prev_event_id": prev_event_id, + "instance_name": rng.choice(WRITERS + [None]), + }, + desc="null_event_delta", + ) + ) + + @parameterized.expand(SEEDS) + def test_snapshot_membership_changes(self, seed: int) -> None: + rng = random.Random(seed) + rooms = self._random_memberships(rng) + self._scatter_writers(rng, self.bob) + self._null_event_deltas(rng, rooms) + lo, hi = self._stream_range() + for _ in range(10): + from_token = self._random_token(rng, lo, hi) + to_token = self._random_token(rng, lo, hi) + excluded = set(rng.sample(rooms, rng.randint(0, 2))) + may_have_changed = self.store._membership_stream_cache.has_entity_changed( + self.bob, from_token.stream + ) + expected = self.get_success( + self.store.get_sliding_sync_membership_changes( + self.bob, from_token, to_token, excluded_room_ids=excluded + ) + ) + self.assertEqual( + self.rust.snapshot_membership_changes( + self.bob, from_token, to_token, list(excluded), may_have_changed + ), + [ + ( + change.room_id, + change.sender, + change.membership, + change.event_id, + _position(change.event_pos), + change.room_version_id, + ) + for change in expected.values() + ], + f"from {from_token} to {to_token} excluding {excluded}", + ) + + @parameterized.expand(SEEDS) + def test_membership_delta_changes(self, seed: int) -> None: + rng = random.Random(seed) + rooms = self._random_memberships(rng) + self._scatter_writers(rng, self.bob) + self._null_event_deltas(rng, rooms) + lo, hi = self._stream_range() + for _ in range(10): + from_token = self._random_token(rng, lo, hi) + to_token = self._random_token(rng, lo, hi) + excluded = rng.sample(rooms, rng.randint(0, 2)) + may_have_changed = self.store._membership_stream_cache.has_entity_changed( + self.bob, from_token.stream + ) + expected = self.get_success( + self.store.get_current_state_delta_membership_changes_for_user( + self.bob, from_token, to_token, excluded_room_ids=excluded + ) + ) + self.assertEqual( + self.rust.membership_delta_changes( + self.bob, from_token, to_token, excluded, may_have_changed + ), + [ + ( + change.room_id, + change.event_id, + _position(change.event_pos), + change.membership, + change.sender, + change.prev_event_id, + _position(change.prev_event_pos), + change.prev_membership, + change.prev_sender, + ) + for change in expected + ], + f"from {from_token} to {to_token} excluding {excluded}", + ) + + def test_membership_changes_without_a_change(self) -> None: + """Without a change since the token, as the membership stream cache + says, there is nothing to read.""" + token = self.store.get_room_max_token() + self.helper.leave(self.room_id, self.bob, tok=self.bob_tok) + later = self.store.get_room_max_token() + self.assertEqual( + self.rust.snapshot_membership_changes(self.bob, token, later, [], False), + [], + ) + self.assertEqual( + len( + self.rust.snapshot_membership_changes(self.bob, token, later, [], True) + ), + 1, + ) + self.assertEqual( + self.rust.membership_delta_changes(self.bob, token, token, [], True), [] + ) + + # Stream positions. + + def _rooms_with_events(self, rng: random.Random) -> list[str]: + rooms = [self.room_id] + [ + self.helper.create_room_as(self.alice, tok=self.alice_tok) for _ in range(4) + ] + for _ in range(12): + self.helper.send(rng.choice(rooms), "hi", tok=self.alice_tok) + # A room the server has left. + left = self.helper.create_room_as(self.bob, tok=self.bob_tok) + self.helper.leave(left, self.bob, tok=self.bob_tok) + return rooms + [left, "!unknown:test"] + + @parameterized.expand(SEEDS) + def test_rooms_with_updates_since(self, seed: int) -> None: + rng = random.Random(seed) + rooms = self._rooms_with_events(rng) + lo, hi = self._stream_range() + for _ in range(8): + from_token = RoomStreamToken(stream=rng.randint(lo - 2, hi + 2)) + subset = rng.sample(rooms, rng.randint(0, len(rooms))) + + # The table, with the stream cache out of reach. Python checks + # the background updates without awaiting the check, so the + # coroutine is always true. Only the finished path is ported. + cache = self.store._events_stream_cache + with ( + patch.object( + cache, "get_earliest_known_position", return_value=hi + 10 + ), + patch.object( + self.store, + "have_finished_sliding_sync_background_jobs", + new=lambda: True, + ), + ): + expected = self.get_success( + self.store.get_rooms_that_have_updates_since_sliding_sync_table( + subset, from_token + ) + ) + self.assertEqual( + self.rust.rooms_with_updates_since(subset, from_token, None), + set(expected), + from_token, + ) + + # The stream cache's answer, as Python passes it. + changed = cache.get_all_entities_changed(from_token.stream) + if changed.hit: + expected = self.get_success( + self.store.get_rooms_that_have_updates_since_sliding_sync_table( + subset, from_token + ) + ) + self.assertEqual( + self.rust.rooms_with_updates_since( + subset, from_token, set(changed.entities) + ), + set(expected), + from_token, + ) + + @parameterized.expand(SEEDS) + def test_max_and_last_event_pos(self, seed: int) -> None: + rng = random.Random(seed) + rooms = self._rooms_with_events(rng) + lo, hi = self._stream_range() + cache = self.store._events_stream_cache + for _ in range(8): + # A "current" token behind some rooms' latest events, so both + # sides recheck them. + now_token = self._random_token(rng, (lo + hi) // 2, hi) + end_token = self._random_token(rng, lo, hi) + subset = rng.sample(rooms, rng.randint(0, len(rooms))) + # What the events stream cache knows of some of the rooms. A + # position may be anything, as both sides take it as it is. + last_change_pos = { + room_id: rng.randint(lo, hi) + for room_id in rng.sample(subset, rng.randint(0, len(subset))) + } + # Start each round cold, which the mirror carries over. + self.store._get_max_event_pos.invalidate_all() + with ( + patch.object( + cache, "get_max_pos_of_last_change", side_effect=last_change_pos.get + ), + patch.object(self.store, "get_room_max_token", return_value=now_token), + ): + expected_max = self.get_success( + self.store._bulk_get_max_event_pos(subset) + ) + self.assertEqual( + self.rust.max_event_pos(subset, now_token, last_change_pos), + dict(expected_max), + now_token, + ) + self.store._get_max_event_pos.invalidate_all() + expected_last = self.get_success( + self.store.bulk_get_last_event_pos_in_room_before_stream_ordering( + subset, end_token + ) + ) + self.assertEqual( + self.rust.last_event_pos_before( + subset, end_token, now_token, last_change_pos + ), + expected_last, + f"{end_token} now {now_token}", + ) + + def test_max_event_pos_invalidated_by_new_events(self) -> None: + now = self.store.get_room_max_token() + before = self.rust.max_event_pos([self.room_id], now, {})[self.room_id] + self.helper.send(self.room_id, "hi", tok=self.alice_tok) + now = self.store.get_room_max_token() + after = self.rust.max_event_pos([self.room_id], now, {})[self.room_id] + assert before is not None and after is not None + self.assertGreater(after, before) + self.assertEqual(after, now.stream) + + # Room metadata. + + def test_room_metadata(self) -> None: + rooms = self._bob_memberships() + ["!unknown:test"] + for room_id in rooms: + self.store.get_room_type.invalidate((room_id,)) + room_type = self.get_success(self.store.get_room_type(room_id)) + self.assertEqual( + self.rust.room_type(room_id), + (False, None) + if room_type is ROOM_UNKNOWN_SENTINEL + else (True, room_type), + room_id, + ) + try: + version: str | None = self.get_success( + self.store.get_room_version_id(room_id) + ) + except AssertionError: + version = None + self.assertEqual(self.rust.room_version_id(room_id), version, room_id) + self.assertEqual( + self.rust.cached_member_counts(room_id), + dict(self.get_success(self.store.get_member_counts(room_id))), + room_id, + ) + + # A join invalidates the counts. + self.helper.join(self.room_id, self.alice, tok=self.alice_tok) + carol = self.register_user("carol", "pass") + self.helper.join(self.room_id, carol, tok=self.login("carol", "pass")) + self.assertEqual( + self.rust.cached_member_counts(self.room_id), + dict(self.get_success(self.store.get_member_counts(self.room_id))), + ) + + def test_partial_rooms(self) -> None: + self.assertEqual( + self.rust.partial_rooms(), + set(self.get_success(self.store.get_partial_rooms())), + ) + other = self.helper.create_room_as(self.alice, tok=self.alice_tok) + join = self.helper.join(other, self.bob, tok=self.bob_tok)["event_id"] + self.get_success( + self.store.db_pool.simple_insert( + "partial_state_rooms", + { + "room_id": other, + "device_lists_stream_id": 0, + "join_event_id": join, + "joined_via": "other", + }, + ) + ) + # Behind the cache's back. Python invalidates the whole cache when + # it changes the table, and the mirror with it. + self.store.get_partial_rooms.invalidate_all() + self.assertEqual(self.rust.partial_rooms(), {other}) + self.assertEqual( + self.rust.partial_rooms(), + set(self.get_success(self.store.get_partial_rooms())), + ) + + def test_tags_and_account_data(self) -> None: + other = self.helper.create_room_as(self.bob, tok=self.bob_tok) + self.assertEqual( + self.rust.tags_for_user(self.bob), + self.get_success(self.store.get_tags_for_user(self.bob)), + ) + self.get_success( + self.store.add_tag_to_room(self.bob, self.room_id, "m.favourite", {}) + ) + self.get_success( + self.store.add_tag_to_room( + self.bob, other, "u.work", {"order": 0.5, "é": "ü"} + ) + ) + self.get_success(self.store.add_tag_to_room(self.bob, other, "u.home", {})) + expected = self.get_success(self.store.get_tags_for_user(self.bob)) + self.assertEqual(len(expected), 2) + self.assertEqual(self.rust.tags_for_user(self.bob), expected) + self.get_success(self.store.remove_tag_from_room(self.bob, other, "u.home")) + self.assertEqual( + self.rust.tags_for_user(self.bob), + self.get_success(self.store.get_tags_for_user(self.bob)), + ) + + @parameterized.expand(SEEDS[:5]) + def test_is_encrypted_at(self, seed: int) -> None: + rng = random.Random(seed) + room_lists = self.hs.get_sliding_sync_handler().room_lists + room_id = self.helper.create_room_as(self.alice, tok=self.alice_tok) + for _ in range(4): + self.helper.send(room_id, "hi", tok=self.alice_tok) + if rng.random() < 0.5: + self.helper.send_state( + room_id, + EventTypes.RoomEncryption, + rng.choice( + [ + {EventContentFields.ENCRYPTION_ALGORITHM: "m.megolm.v1"}, + {EventContentFields.ENCRYPTION_ALGORITHM: None}, + {}, + ] + ), + tok=self.alice_tok, + ) + lo, hi = self._stream_range() + now_token = self.store.get_room_max_token() + for stream in range(lo - 1, hi + 1): + to_token = RoomStreamToken(stream=stream) + self.assertEqual( + self.rust.is_encrypted_at(room_id, to_token, now_token), + self.get_success( + room_lists.get_is_encrypted_for_room_at_token(room_id, to_token) + ), + stream, + ) + + # Connection state. + + def _random_state( + self, rng: random.Random, rooms: list[str], users: list[str] + ) -> MutablePerConnectionState: + """A request's random updates to a connection's state.""" + state = PerConnectionState().get_mutable() + for room_id in rng.sample(rooms, rng.randint(0, len(rooms))): + if rng.random() < 0.5: + state.rooms._statuses[room_id] = HaveSentRoom.live() + else: + state.rooms._statuses[room_id] = HaveSentRoom.previously( + RoomStreamToken( + stream=rng.randint(1, 50), + instance_map=immutabledict( + {"master": rng.randint(51, 60)} + if rng.random() < 0.3 + else {} + ), + ) + ) + for room_id in rng.sample(rooms, rng.randint(0, 2)): + state.receipts._statuses[room_id] = ( + HaveSentRoom.live() + if rng.random() < 0.5 + else HaveSentRoom.previously( + MultiWriterStreamToken(stream=rng.randint(1, 50)) + ) + ) + for room_id in rng.sample(rooms, rng.randint(0, 2)): + state.account_data._statuses[room_id] = ( + HaveSentRoom.live() + if rng.random() < 0.5 + else HaveSentRoom.previously(rng.randint(1, 50)) + ) + for room_id in rng.sample(rooms, rng.randint(0, len(rooms))): + required_state: dict[str, set[str]] = {} + for event_type in rng.sample( + [EventTypes.Member, EventTypes.Name, "*", "org.é"], rng.randint(0, 3) + ): + required_state[event_type] = set( + rng.sample(["", "$ME", "$LAZY", "*", self.bob], rng.randint(1, 2)) + ) + state.room_configs[room_id] = RoomSyncConfig( + timeline_limit=rng.randint(0, 3), + required_state_map=required_state, + ) + for room_id in rng.sample(rooms, rng.randint(0, 2)): + state.room_lazy_membership[room_id] = RoomLazyMembershipChanges( + returned_user_id_to_last_seen_ts_map={ + user_id: rng.choice( + [ + None, + self.clock.time_msec(), + self.clock.time_msec() - 4_000_000, + ] + ) + for user_id in rng.sample(users, rng.randint(0, len(users))) + }, + invalidated_user_ids=set(rng.sample(users, rng.randint(0, 1))), + ) + return state + + def _rust_updates(self, state: MutablePerConnectionState) -> dict[str, Any]: + """The state's updates as the Rust test entry point takes them.""" + + def statuses(updates: Any, to_string: Any) -> dict[str, tuple[str, Any]]: + return { + room_id: ( + status.status.value, + None + if status.last_token is None + else self.get_success(to_string(status.last_token)), + ) + for room_id, status in updates.items() + } + + async def to_string(token: Any) -> str: + if isinstance(token, int): + return str(token) + return await token.to_string(self.store) + + return { + "rooms": statuses(state.rooms.get_updates(), to_string), + "receipts": statuses(state.receipts.get_updates(), to_string), + "account_data": statuses(state.account_data.get_updates(), to_string), + "room_configs": { + room_id: ( + config.timeline_limit, + { + event_type: sorted(state_keys) + for event_type, state_keys in config.required_state_map.items() + }, + ) + for room_id, config in state.room_configs.maps[0].items() + }, + "room_lazy_membership": { + room_id: ( + dict(changes.returned_user_id_to_last_seen_ts_map), + set(changes.invalidated_user_ids), + ) + for room_id, changes in state.room_lazy_membership.items() + }, + "now_ms": self.clock.time_msec(), + } + + def _python_loaded(self, state: PerConnectionState) -> tuple[Any, ...]: + """A Python loaded state in the Rust test form.""" + + def statuses(status_map: Any) -> dict[str, tuple[str, Any]]: + return { + room_id: ( + status.status.value, + None + if status.last_token is None + else self.get_success(_to_string(self.store, status.last_token)), + ) + for room_id, status in status_map._statuses.items() + } + + return ( + state.last_used_ts, + statuses(state.rooms), + statuses(state.receipts), + statuses(state.account_data), + { + room_id: ( + config.timeline_limit, + { + event_type: sorted(state_keys) + for event_type, state_keys in config.required_state_map.items() + }, + ) + for room_id, config in state.room_configs.items() + }, + ) + + def _connection_rows(self, conn_id: str) -> dict[str, Any]: + """Every row of the connection, with IDs replaced: positions by their + order, required state IDs by their JSON.""" + db = self.store.db_pool + + def select(sql: str, *args: Any) -> list[tuple[Any, ...]]: + return self.get_success(db.execute("rows", sql, *args)) + + connections = select( + "SELECT connection_key, created_ts, last_used_ts FROM sliding_sync_connections" + " WHERE user_id = ? AND effective_device_id = ? AND conn_id = ?", + self.bob, + DEVICE, + conn_id, + ) + if not connections: + return {} + [(key, created_ts, last_used_ts)] = connections + positions = [ + row[0] + for row in select( + "SELECT connection_position FROM sliding_sync_connection_positions" + " WHERE connection_key = ? ORDER BY connection_position", + key, + ) + ] + index = {position: i for i, position in enumerate(positions)} + required_state = dict( + select( + "SELECT required_state_id, required_state" + " FROM sliding_sync_connection_required_state WHERE connection_key = ?", + key, + ) + ) + configs = { + (index[position], room_id): (limit, required_state[state_id]) + for position, room_id, limit, state_id in select( + "SELECT connection_position, room_id, timeline_limit, required_state_id" + " FROM sliding_sync_connection_room_configs" + " WHERE connection_position = ANY(?)", + positions, + ) + } + streams = { + (index[position], stream, room_id): (status, token) + for position, stream, room_id, status, token in select( + "SELECT connection_position, stream, room_id, room_status, last_token" + " FROM sliding_sync_connection_streams WHERE connection_position = ANY(?)", + positions, + ) + } + lazy = { + (room_id, user_id): (None if position is None else index[position], ts) + for position, room_id, user_id, ts in select( + "SELECT connection_position, room_id, user_id, last_seen_ts" + " FROM sliding_sync_connection_lazy_members WHERE connection_key = ?", + key, + ) + } + return { + "created_ts": created_ts, + "last_used_ts": last_used_ts, + "positions": len(positions), + "required_state": sorted(required_state.values()), + "configs": configs, + "streams": streams, + "lazy": lazy, + } + + @parameterized.expand(SEEDS) + def test_connection_state(self, seed: int) -> None: + """The same operations through Python's transactions on one + connection and Rust's on another leave the same rows and load the + same states.""" + rng = random.Random(seed) + rooms = ["!a:test", "!b:test", "!c:test", "!d:test"] + users = [self.alice, self.bob, "@carol:test"] + # Each side's positions, in the order the operations made them. + positions: dict[str, list[int]] = {"py": [], "rs": []} + + for step in range(12): + context = f"step {step}" + self.reactor.advance(rng.choice([0, 1, 200, 400, 4000])) + op = rng.random() + # A position to build on or load: usually the latest, sometimes + # an older one (a fork), sometimes none or a stranger's. + choice = rng.random() + index: int | None + if not positions["py"] or choice < 0.1: + index = None + elif choice < 0.75: + index = len(positions["py"]) - 1 + else: + index = rng.randrange(len(positions["py"])) + + if op < 0.6: + state = self._random_state(rng, rooms, users) + updates = self._rust_updates(state) + outcomes = {} + py_prev = None if index is None else positions["py"][index] + rs_prev = None if index is None else positions["rs"][index] + try: + py_pos = self.get_success_or_raise( + self.store.persist_per_connection_state( + self.bob, DEVICE, "py", py_prev, state + ) + ) + outcomes["py"] = "ok" + except SlidingSyncUnknownPosition: + outcomes["py"] = "unknown" + try: + rs_pos = self.rust.persist_connection_state( + self.bob, DEVICE, "rs", rs_prev, **updates + ) + outcomes["rs"] = "ok" + except SlidingSyncUnknownPosition: + outcomes["rs"] = "unknown" + self.assertEqual(outcomes["py"], outcomes["rs"], context) + if outcomes["py"] == "ok": + positions["py"].append(py_pos) + positions["rs"].append(rs_pos) + elif index is not None: + now = self.clock.time_msec() + try: + py_loaded = self._python_loaded( + self.get_success_or_raise( + self.store.get_and_clear_connection_positions( + self.bob, DEVICE, "py", positions["py"][index] + ) + ) + ) + except SlidingSyncUnknownPosition: + py_loaded = None + try: + rs_loaded = self.rust.load_connection_state( + self.bob, DEVICE, "rs", positions["rs"][index], now + ) + except SlidingSyncUnknownPosition: + rs_loaded = None + self.assertEqual(rs_loaded, py_loaded, context) + + py_rows = self._connection_rows("py") + rs_rows = self._connection_rows("rs") + self.assertEqual(rs_rows, py_rows, context) + + # Someone else's position is unknown to both. + if positions["py"]: + with self.assertRaises(SlidingSyncUnknownPosition): + self.rust.load_connection_state( + self.alice, DEVICE, "rs", positions["rs"][-1], 0 + ) + with self.assertRaises(SlidingSyncUnknownPosition): + self.rust.persist_connection_state( + self.bob, DEVICE, "py", positions["rs"][-1], **updates + ) + + def test_connection_lazy_members(self) -> None: + state = PerConnectionState().get_mutable() + state.room_lazy_membership[self.room_id] = RoomLazyMembershipChanges( + returned_user_id_to_last_seen_ts_map={self.alice: None, self.bob: None} + ) + first = self.get_success( + self.store.persist_per_connection_state( + self.bob, DEVICE, "conn", None, state + ) + ) + # Two forks off the first position, each sending another member. + forks = [] + for user_id in ("@carol:test", "@dave:test"): + state = PerConnectionState().get_mutable() + state.room_lazy_membership[self.room_id] = RoomLazyMembershipChanges( + returned_user_id_to_last_seen_ts_map={user_id: None} + ) + forks.append( + self.get_success( + self.store.persist_per_connection_state( + self.bob, DEVICE, "conn", first, state + ) + ) + ) + + users = [self.alice, self.bob, "@carol:test", "@dave:test", "@eve:test"] + + def compare() -> None: + for position in [first, *forks, 12345]: + for room_id in (self.room_id, "!other:test"): + self.assertEqual( + self.rust.connection_lazy_members(position, room_id, users), + dict( + self.get_success( + self.store.get_sliding_sync_connection_lazy_members( + position, room_id, set(users) + ) + ) + ), + (position, room_id), + ) + self.assertEqual( + self.rust.connection_lazy_members(first, self.room_id, []), {} + ) + + compare() + # Loading a fork commits to it, which drops the other. + self.get_success( + self.store.get_and_clear_connection_positions( + self.bob, DEVICE, "conn", forks[0] + ) + ) + compare() + + def test_load_memo(self) -> None: + """A position loads once: the second load is answered from the memo, + without the transaction's side effects.""" + state = PerConnectionState().get_mutable() + state.rooms._statuses[self.room_id] = HaveSentRoom.live() + first = self.rust.persist_connection_state( + self.bob, DEVICE, "conn", None, **self._rust_updates(state) + ) + second = self.rust.persist_connection_state( + self.bob, DEVICE, "conn", first, **self._rust_updates(state) + ) + loaded = self.rust.load_connection_state(self.bob, DEVICE, "conn", first, 0) + self.assertEqual(loaded[1], {self.room_id: ("live", None)}) + # The load deleted the other position. + self.assertEqual(self._connection_rows("conn")["positions"], 1) + with self.assertRaises(SlidingSyncUnknownPosition): + self.rust.load_connection_state(self.bob, DEVICE, "conn", second, 0) + # Again from the memo, though the row's last use is now different. + self.assertEqual( + self.rust.load_connection_state(self.bob, DEVICE, "conn", first, 10**12), + loaded, + ) + self.assertEqual( + self._connection_rows("conn")["last_used_ts"], self.clock.time_msec() + ) + + +async def _to_string(store: Any, token: Any) -> str: + if isinstance(token, int): + return str(token) + return await token.to_string(store) diff --git a/tests/storage/test_sync_store.py b/tests/storage/test_sync_store.py index 79604e77ef..6e456d7bd0 100644 --- a/tests/storage/test_sync_store.py +++ b/tests/storage/test_sync_store.py @@ -654,6 +654,38 @@ class SyncStoreTestCase(unittest.HomeserverTestCase): str(end_token), ) + @parameterized.expand(SEEDS) + def test_last_event_of_types_before(self, seed: int) -> None: + rng = random.Random(seed) + self._write_events(self.room_id, 15) + lo, hi = self._stream_range(self.room_id) + type_choices = [ + EventTypes.Message, + EventTypes.Topic, + EventTypes.Member, + EventTypes.Create, + "org.example.unknown", + ] + for _ in range(12): + end_token = self._random_token(rng, lo, hi, False) + event_types = rng.sample(type_choices, rng.randint(0, 3)) + expected = self.get_success( + self.store.get_last_event_pos_in_room_before_stream_ordering( + self.room_id, end_token, event_types=event_types + ) + ) + last = self.rust.last_event_of_types_before( + self.room_id, end_token, event_types + ) + context = f"{event_types} at {end_token}" + if expected is None: + self.assertIsNone(last, context) + else: + event_id, position = expected + self.assertEqual( + last, (event_id, position.instance_name, position.stream), context + ) + @parameterized.expand(SEEDS) def test_membership_changes(self, seed: int) -> None: rng = random.Random(seed) @@ -849,6 +881,101 @@ class SyncStoreTestCase(unittest.HomeserverTestCase): f"from {from_token} to {to_token}", ) + @parameterized.expand(SEEDS) + def test_state_deltas_by_stream_id(self, seed: int) -> None: + rng = random.Random(seed) + event_ids = self._write_events(self.room_id, 15) + other = self.helper.create_room_as(self.alice, tok=self.alice_tok) + self._write_events(other, 4) + # The delta rows of the rewritten events move to the same writer. + for event_id in event_ids: + rows = self.get_success( + self.store.db_pool.simple_select_list( + "events", {"event_id": event_id}, ["instance_name"] + ) + ) + self.get_success( + self.store.db_pool.simple_update( + "current_state_delta_stream", + {"event_id": event_id}, + {"instance_name": rows[0][0]}, + desc="rewrite", + ) + ) + lo, hi = self._stream_range(self.room_id) + for _ in range(12): + from_token = ( + self._random_token(rng, lo, hi, False) if rng.random() < 0.8 else None + ) + to_token = ( + self._random_token(rng, lo, hi, False) if rng.random() < 0.8 else None + ) + expected = self.get_success( + self.store.get_current_state_deltas_for_room( + self.room_id, from_token=from_token, to_token=to_token + ) + ) + # In `stream_id` order. Deltas of one batch share a `stream_id`, so + # compare them as a sorted list. + rust = self.rust.state_deltas_by_stream_id( + self.room_id, from_token, to_token + ) + self.assertEqual( + [delta[0] for delta in rust], sorted(delta[0] for delta in rust) + ) + self.assertEqual( + sorted(rust), + sorted( + ( + d.stream_id, + d.event_type, + d.state_key, + d.event_id, + d.prev_event_id, + ) + for d in expected + ), + f"from {from_token} to {to_token}", + ) + + def test_member_counts(self) -> None: + carol = self.register_user("carol", "pass") + carol_tok = self.login("carol", "pass") + dave = self.register_user("dave", "pass") + self.helper.invite(self.room_id, self.alice, carol, tok=self.alice_tok) + self.helper.invite(self.room_id, self.alice, dave, tok=self.alice_tok) + self.helper.join(self.room_id, carol, tok=carol_tok) + self.helper.leave(self.room_id, carol, tok=carol_tok) + self.helper.ban(self.room_id, self.alice, self.bob, tok=self.alice_tok) + other = self.helper.create_room_as(self.bob, tok=self.bob_tok) + + for room_id in (self.room_id, other, "!unknown:test"): + self.assertEqual( + self.rust.member_counts(room_id), + dict(self.get_success(self.store.get_member_counts(room_id))), + room_id, + ) + + def test_latest_bump_stamp(self) -> None: + self.helper.send(self.room_id, "hi", tok=self.alice_tok) + # A room whose row has no bump stamp. + no_stamp = self.helper.create_room_as(self.alice, tok=self.alice_tok) + self.get_success( + self.store.db_pool.simple_update_one( + "sliding_sync_joined_rooms", + {"room_id": no_stamp}, + {"bump_stamp": None}, + ) + ) + for room_id in (self.room_id, no_stamp, "!unknown:test"): + self.assertEqual( + self.rust.latest_bump_stamp(room_id), + self.get_success(self.store.get_latest_bump_stamp_for_room(room_id)), + room_id, + ) + self.assertIsNotNone(self.rust.latest_bump_stamp(self.room_id)) + self.assertIsNone(self.rust.latest_bump_stamp(no_stamp)) + # Notifications. def _unread_counts( diff --git a/tests/utils.py b/tests/utils.py index 6efc8983ff..e452685af8 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -70,6 +70,14 @@ RUST_PERSIST_FOR_TESTS = os.environ.get("SYNAPSE_TEST_RUST_PERSIST") RUST_PERSIST_STRICT_FOR_TESTS = ( os.environ.get("SYNAPSE_TEST_RUST_PERSIST_STRICT", "1") != "0" ) +# As above, for `experimental_features.rust_sliding_sync` ("off" or "on"). +RUST_SLIDING_SYNC_FOR_TESTS = os.environ.get("SYNAPSE_TEST_RUST_SLIDING_SYNC") +# A Rust error fails the test rather than falling back to Python, so that a +# test fails on the room or request Rust got wrong. Set to 0 to fall back and +# count, as production does. +RUST_SLIDING_SYNC_STRICT_FOR_TESTS = ( + os.environ.get("SYNAPSE_TEST_RUST_SLIDING_SYNC_STRICT", "1") != "0" +) # At `rust_federation_sending: transactions`, check every transaction body Rust # builds against the one Python would build, and fail the test if they differ. RUST_FED_STRICT_FOR_TESTS = os.environ.get("SYNAPSE_TEST_RUST_FED_STRICT", "1") != "0"