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/mod.rs b/rust/src/sliding_sync/mod.rs index 6c5058f7d2..a831ae3146 100644 --- a/rust/src/sliding_sync/mod.rs +++ b/rust/src/sliding_sync/mod.rs @@ -37,6 +37,7 @@ use crate::runtime::RustRuntime; use crate::storage::sync_store::{SyncDbStats, SyncStore}; pub mod config; +pub mod connection; pub mod request; pub mod required_state; @@ -164,6 +165,7 @@ pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> child_module.add_class::()?; child_module.add_class::()?; config::register_test_functions(&child_module)?; + connection::register_test_functions(&child_module)?; required_state::register_test_functions(&child_module)?; m.add_submodule(&child_module)?; diff --git a/synapse/synapse_rust/sliding_sync.pyi b/synapse/synapse_rust/sliding_sync.pyi index 9908d1db06..78dcc44c4e 100644 --- a/synapse/synapse_rust/sliding_sync.pyi +++ b/synapse/synapse_rust/sliding_sync.pyi @@ -139,3 +139,30 @@ def required_state_request_for_test( 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))`.""" diff --git a/tests/handlers/test_sliding_sync_rust_pure.py b/tests/handlers/test_sliding_sync_rust_pure.py index 0ca18a6632..33c7fd42eb 100644 --- a/tests/handlers/test_sliding_sync_rust_pure.py +++ b/tests/handlers/test_sliding_sync_rust_pure.py @@ -18,7 +18,7 @@ inputs.""" import random from types import SimpleNamespace -from typing import AbstractSet, Mapping, cast +from typing import AbstractSet, Any, Mapping, cast from synapse.api.constants import EventTypes, Membership from synapse.events import EventBase @@ -29,12 +29,19 @@ from synapse.handlers.sliding_sync import ( _state_filter_with_meta, ) from synapse.synapse_rust import sliding_sync as rust +from synapse.types import RoomStreamToken from synapse.types.handlers.sliding_sync import ( + HaveSentRoom, + HaveSentRoomFlag, + PerConnectionState, + RoomLazyMembershipChanges, + RoomStatusMap, RoomSyncConfig, SlidingSyncConfig, StateValues, ) from synapse.types.state import StateFilter +from synapse.util.clock import Clock from tests import unittest @@ -388,3 +395,201 @@ class RequiredStateRequestTestCase(unittest.TestCase): ), 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), + ), + )