diff --git a/rust/src/sync/ll_members.rs b/rust/src/sync/ll_members.rs new file mode 100644 index 0000000000..839e71992b --- /dev/null +++ b/rust/src/sync/ll_members.rs @@ -0,0 +1,356 @@ +/* + * 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 membership events a lazy-loading client has been sent, so that later +//! syncs can leave them out: `LazyLoadedMembersCache` in +//! `synapse.handlers.sync`, one per `(user, device)`. +//! +//! A member may only be left out if the client processed the response that +//! carried it. Each entry therefore records the `next_batch` position of +//! the response that carried it. A client that has processed a response +//! syncs from its `next_batch` next, so a request from `since` means the +//! client processed every response up to `since` and none after it. +//! [`LazyLoadedMembersCache::start_request`] forgets the entries recorded +//! after `since`: they were in responses the client lost, by timing out a +//! slow request and retrying it, or by reloading from an older persisted +//! token. Entries at or before `since` stay. An initial sync forgets +//! everything. +//! +//! [`LazyLoadedMembersCaches`] holds the per-device caches and drops the +//! ones not used for [`LAZY_LOADED_MEMBERS_CACHE_MAX_AGE`], as the Python +//! `ExpiringCache` does. + +use std::{ + collections::HashMap, + num::NonZeroUsize, + sync::Arc, + time::{Duration, SystemTime}, +}; + +use pyo3::prelude::*; + +use crate::types::stream_token::RoomStreamToken; + +/// Remember the last 100 members we sent to a client. +pub const LAZY_LOADED_MEMBERS_CACHE_MAX_SIZE: usize = 100; + +/// Keep a device's cache for no more than 30 minutes without use. +pub const LAZY_LOADED_MEMBERS_CACHE_MAX_AGE: Duration = Duration::from_secs(30 * 60); + +/// The members sent to one device. See the module docs. +#[pyclass(name = "LazyLoadedMembersCache")] +pub struct LazyLoadedMembersCache { + /// User ID to the event ID of the membership we last sent and the room + /// position of the `next_batch` of the response that carried it. Least + /// recently used first. + sent: lru::LruCache, (Arc, RoomStreamToken)>, + + /// An upper bound on the positions of the entries, so that + /// `start_request` need not look at them when nothing was recorded + /// after `since`. That is the steady state. + newest: Option, +} + +impl Default for LazyLoadedMembersCache { + fn default() -> Self { + Self::new() + } +} + +impl LazyLoadedMembersCache { + pub fn new() -> Self { + LazyLoadedMembersCache { + sent: lru::LruCache::new( + NonZeroUsize::new(LAZY_LOADED_MEMBERS_CACHE_MAX_SIZE).expect("non-zero"), + ), + newest: None, + } + } + + /// Forget the members the client did not receive. Called at the start + /// of each sync request for this device, with the request's `since` + /// room position, or `None` for an initial sync. + pub fn start_request(&mut self, since: Option<&RoomStreamToken>) { + let Some(since) = since else { + self.sent.clear(); + self.newest = None; + return; + }; + + // The common case: nothing has been recorded after `since`. + if self + .newest + .as_ref() + .is_none_or(|newest| newest.is_before_or_eq(since)) + { + return; + } + + let lost: Vec> = self + .sent + .iter() + .filter(|(_, (_, sent_at))| !sent_at.is_before_or_eq(since)) + .map(|(user_id, _)| Arc::clone(user_id)) + .collect(); + for user_id in lost { + self.sent.pop(&user_id); + } + + // Everything left is at or before `since`. + self.newest = Some(since.clone()); + } + + /// Whether `event_id` is the membership of `user_id` we last sent. A hit + /// counts as a use. + pub fn was_sent(&mut self, user_id: &str, event_id: &str) -> bool { + match self.sent.get(user_id) { + Some((sent_event_id, _)) => &**sent_event_id == event_id, + None => false, + } + } + + /// Record that `event_id` is being sent in the response ending at + /// `next_batch`. + pub fn mark_sent(&mut self, user_id: &str, event_id: &str, next_batch: &RoomStreamToken) { + self.sent.push( + Arc::from(user_id), + (Arc::from(event_id), next_batch.clone()), + ); + + // Only move the bound forwards. Positions taken by this process only + // ever advance, but a response that started earlier may record its + // entries after a later one has. + if self + .newest + .as_ref() + .is_none_or(|newest| newest.is_before_or_eq(next_batch)) + { + self.newest = Some(next_batch.clone()); + } + } + + pub fn len(&self) -> usize { + self.sent.len() + } + + pub fn is_empty(&self) -> bool { + self.sent.is_empty() + } +} + +#[pymethods] +impl LazyLoadedMembersCache { + #[new] + fn py_new() -> Self { + Self::new() + } + + #[pyo3(name = "start_request", signature = (since))] + fn py_start_request(&mut self, since: Option) { + self.start_request(since.as_ref()) + } + + #[pyo3(name = "was_sent")] + fn py_was_sent(&mut self, user_id: &str, event_id: &str) -> bool { + self.was_sent(user_id, event_id) + } + + #[pyo3(name = "mark_sent")] + fn py_mark_sent(&mut self, user_id: &str, event_id: &str, next_batch: RoomStreamToken) { + self.mark_sent(user_id, event_id, &next_batch) + } + + fn __len__(&self) -> usize { + self.len() + } +} + +/// A device's cache, and when it was last used. +struct Entry { + cache: LazyLoadedMembersCache, + last_used: SystemTime, +} + +/// The caches of every `(user, device)`, dropped after +/// [`LAZY_LOADED_MEMBERS_CACHE_MAX_AGE`] without use. The Python +/// `ExpiringCache` prunes on a timer; this prunes on use, at most once per +/// half the max age, which bounds the work the same way. +pub struct LazyLoadedMembersCaches { + entries: HashMap<(Arc, Option>), Entry>, + last_pruned: Option, +} + +impl Default for LazyLoadedMembersCaches { + fn default() -> Self { + Self::new() + } +} + +impl LazyLoadedMembersCaches { + pub fn new() -> Self { + LazyLoadedMembersCaches { + entries: HashMap::new(), + last_pruned: None, + } + } + + /// The cache for `(user_id, device_id)`, created if it is not here, as + /// used at `now`. + pub fn get_or_create( + &mut self, + user_id: &str, + device_id: Option<&str>, + now: SystemTime, + ) -> &mut LazyLoadedMembersCache { + self.maybe_prune(now); + let entry = self + .entries + .entry((Arc::from(user_id), device_id.map(Arc::from))) + .or_insert_with(|| Entry { + cache: LazyLoadedMembersCache::new(), + last_used: now, + }); + entry.last_used = now; + &mut entry.cache + } + + /// Drop the caches not used since `now - max_age`, if it has been half a + /// max age since the last prune. + fn maybe_prune(&mut self, now: SystemTime) { + let due = match self.last_pruned { + None => true, + Some(last) => now + .duration_since(last) + .is_ok_and(|since| since >= LAZY_LOADED_MEMBERS_CACHE_MAX_AGE / 2), + }; + if !due { + return; + } + self.last_pruned = Some(now); + self.entries.retain(|_, entry| { + now.duration_since(entry.last_used) + .is_ok_and(|unused| unused <= LAZY_LOADED_MEMBERS_CACHE_MAX_AGE) + || now < entry.last_used + }); + } + + pub fn len(&self) -> usize { + self.entries.len() + } + + pub fn is_empty(&self) -> bool { + self.entries.is_empty() + } +} + +#[cfg(test)] +mod tests { + use std::time::UNIX_EPOCH; + + use super::*; + + fn at(stream: i64) -> RoomStreamToken { + RoomStreamToken::live(stream) + } + + #[test] + fn forgets_only_entries_after_since() { + let mut cache = LazyLoadedMembersCache::new(); + cache.mark_sent("@a:test", "$a", &at(5)); + cache.mark_sent("@b:test", "$b", &at(7)); + + // The client synced past both responses. + cache.start_request(Some(&at(7))); + assert!(cache.was_sent("@a:test", "$a")); + assert!(cache.was_sent("@b:test", "$b")); + + // The client lost the response that ended at 7. + cache.start_request(Some(&at(6))); + assert!(cache.was_sent("@a:test", "$a")); + assert!(!cache.was_sent("@b:test", "$b")); + + // A later request from an even older token forgets the rest. + cache.start_request(Some(&at(4))); + assert!(!cache.was_sent("@a:test", "$a")); + } + + #[test] + fn initial_sync_forgets_everything() { + let mut cache = LazyLoadedMembersCache::new(); + cache.mark_sent("@a:test", "$a", &at(5)); + cache.start_request(None); + assert!(!cache.was_sent("@a:test", "$a")); + assert!(cache.is_empty()); + } + + #[test] + fn different_event_is_not_sent() { + let mut cache = LazyLoadedMembersCache::new(); + cache.mark_sent("@a:test", "$a1", &at(5)); + assert!(!cache.was_sent("@a:test", "$a2")); + } + + #[test] + fn out_of_order_marks_are_still_forgotten() { + let mut cache = LazyLoadedMembersCache::new(); + cache.mark_sent("@b:test", "$b", &at(7)); + cache.mark_sent("@a:test", "$a", &at(5)); + cache.start_request(Some(&at(6))); + assert!(cache.was_sent("@a:test", "$a")); + assert!(!cache.was_sent("@b:test", "$b")); + } + + #[test] + fn evicts_least_recently_used() { + let mut cache = LazyLoadedMembersCache::new(); + for i in 0..LAZY_LOADED_MEMBERS_CACHE_MAX_SIZE { + cache.mark_sent(&format!("@u{i}:test"), &format!("$u{i}"), &at(1)); + } + // Touch the oldest entry so that it is no longer the eviction + // candidate. + assert!(cache.was_sent("@u0:test", "$u0")); + cache.mark_sent("@new:test", "$new", &at(1)); + assert!(cache.was_sent("@u0:test", "$u0")); + assert!(!cache.was_sent("@u1:test", "$u1")); + assert!(cache.was_sent("@new:test", "$new")); + assert_eq!(cache.len(), LAZY_LOADED_MEMBERS_CACHE_MAX_SIZE); + } + + #[test] + fn per_device_caches_expire_without_use() { + let mut caches = LazyLoadedMembersCaches::new(); + let t0 = UNIX_EPOCH + Duration::from_secs(1_000_000); + caches + .get_or_create("@a:test", Some("D1"), t0) + .mark_sent("@b:test", "$b", &at(1)); + caches.get_or_create("@a:test", None, t0); + assert_eq!(caches.len(), 2); + + // Used within the max age: kept, and the use restamps it. + let t1 = t0 + LAZY_LOADED_MEMBERS_CACHE_MAX_AGE - Duration::from_secs(1); + assert!(caches + .get_or_create("@a:test", Some("D1"), t1) + .was_sent("@b:test", "$b")); + assert_eq!(caches.len(), 2); + + // At the next prune, the one unused for over a max age goes and + // the one used at `t1` stays. + let t2 = t1 + LAZY_LOADED_MEMBERS_CACHE_MAX_AGE / 2; + caches.get_or_create("@c:test", None, t2); + assert_eq!(caches.len(), 2); + assert!(caches + .get_or_create("@a:test", Some("D1"), t2) + .was_sent("@b:test", "$b")); + } +} diff --git a/rust/src/sync/mod.rs b/rust/src/sync/mod.rs index 6512099809..8939fe7c06 100644 --- a/rust/src/sync/mod.rs +++ b/rust/src/sync/mod.rs @@ -22,21 +22,27 @@ //! //! The pure parts of the per-room loop are each a port of the Python they //! are named after, checked against it by the tests in `tests/sync/`: -//! [`filter`] ports `synapse.api.filtering`, and [`state_delta`] the state -//! delta arithmetic of the sync handler. +//! [`filter`] ports `synapse.api.filtering`, [`state_delta`] the state +//! delta arithmetic of the sync handler, and [`ll_members`] the per-device cache +//! of lazy-loaded members sent. //! [`crate::visibility`] and [`crate::types::stream_token`] are shared with //! later ports. -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use pyo3::{create_exception, exceptions::PyException, prelude::*, types::PyModule}; +use crate::caches::lock_ignoring_poison; use crate::runtime::RustRuntime; use crate::storage::store::Store; +use crate::types::stream_token::RoomStreamToken; pub mod filter; +pub mod ll_members; pub mod state_delta; +use ll_members::LazyLoadedMembersCaches; + create_exception!( synapse.synapse_rust.sync, RustSyncFallback, @@ -59,6 +65,10 @@ pub struct SyncHandler { /// The Rust store, for its event and state group caches. #[allow(dead_code)] store: Arc, + /// The lazy-loaded members sent to each device. Held here rather than + /// on the store: they are per device, and no data change invalidates + /// them. + ll_members: Mutex, } impl SyncHandler { @@ -67,15 +77,36 @@ impl SyncHandler { db_pool, runtime, store, + ll_members: Mutex::new(LazyLoadedMembersCaches::new()), } } } +#[pymethods] +impl SyncHandler { + /// Forget the lazy-loaded members a device did not receive, at the + /// start of its request from `since`, the request's room position + /// (`None` for an initial sync). See [`ll_members`]. + #[pyo3(signature = (user_id, device_id, since))] + fn start_request( + &self, + user_id: &str, + device_id: Option<&str>, + since: Option, + ) { + let now = self.runtime.clock().now(); + lock_ignoring_poison(&self.ll_members) + .get_or_create(user_id, device_id, now) + .start_request(since.as_ref()); + } +} + /// Called when registering modules with python. pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { let child_module = PyModule::new(py, "sync")?; child_module.add_class::()?; child_module.add_class::()?; + child_module.add_class::()?; state_delta::register_test_functions(py, &child_module)?; child_module.add("RustSyncFallback", py.get_type::())?; diff --git a/synapse/synapse_rust/sync.pyi b/synapse/synapse_rust/sync.pyi index f1479f981b..283b82ce7e 100644 --- a/synapse/synapse_rust/sync.pyi +++ b/synapse/synapse_rust/sync.pyi @@ -13,7 +13,7 @@ from collections.abc import Mapping from synapse.synapse_rust.events import Event -from synapse.types import JsonDict, JsonMapping +from synapse.types import JsonDict, JsonMapping, RoomStreamToken class RustSyncFallback(Exception): """The Rust sync handler cannot compute this room's entry. Compute it @@ -27,6 +27,12 @@ class SyncHandler: uses that driver. """ + def start_request( + self, user_id: str, device_id: str | None, since: RoomStreamToken | None + ) -> None: + """Forget the lazy-loaded members the device did not receive, at the + start of its request from `since` (None for an initial sync).""" + class FilterCollection: """The Rust port of `synapse.api.filtering.FilterCollection`, built from the same filter JSON. The keyword flags are what the Python one reads @@ -74,6 +80,19 @@ class FilterCollection: def check_presence(self, user_id: str) -> bool: ... def check_global_account_data(self, event: JsonDict) -> bool: ... +class LazyLoadedMembersCache: + """The Rust port of `synapse.handlers.sync.LazyLoadedMembersCache`: the + members a lazy-loading device has been sent. Positions are the room part + of the sync tokens.""" + + def __init__(self) -> None: ... + def start_request(self, since: RoomStreamToken | None) -> None: ... + def was_sent(self, user_id: str, event_id: str) -> bool: ... + def mark_sent( + self, user_id: str, event_id: str, next_batch: RoomStreamToken + ) -> None: ... + def __len__(self) -> int: ... + # The Rust ports of the state delta arithmetic in `synapse.handlers.sync`. # Deltas are `(event_type, state_key, event_id, prev_event_id)` in stream # order. diff --git a/tests/sync/test_rust_ll_members.py b/tests/sync/test_rust_ll_members.py new file mode 100644 index 0000000000..e02d4261a5 --- /dev/null +++ b/tests/sync/test_rust_ll_members.py @@ -0,0 +1,116 @@ +# +# 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: +# . +# + +"""`LazyLoadedMembersCache` against `rust/src/sync/ll_members.rs`: the +hand-written cases of the Python one, and random operation sequences.""" + +import random + +from parameterized import parameterized + +from synapse.handlers.sync import ( + LAZY_LOADED_MEMBERS_CACHE_MAX_SIZE, + LazyLoadedMembersCache, +) +from synapse.synapse_rust.sync import ( + LazyLoadedMembersCache as RustLazyLoadedMembersCache, +) +from synapse.types import RoomStreamToken + +from tests import unittest +from tests.sync.fixtures import ( + SEEDS, + random_live_token, +) + + +class LazyLoadedMembersCacheTestCase(unittest.TestCase): + """`LazyLoadedMembersCache` against `rust/src/sync/ll_members.rs`.""" + + @staticmethod + def _token(stream: int) -> RoomStreamToken: + return RoomStreamToken(stream=stream) + + def test_forgets_only_entries_after_since(self) -> None: + cache = RustLazyLoadedMembersCache() + cache.mark_sent("@a:test", "$a", self._token(5)) + cache.mark_sent("@b:test", "$b", self._token(7)) + + cache.start_request(self._token(7)) + self.assertTrue(cache.was_sent("@a:test", "$a")) + self.assertTrue(cache.was_sent("@b:test", "$b")) + + cache.start_request(self._token(6)) + self.assertTrue(cache.was_sent("@a:test", "$a")) + self.assertFalse(cache.was_sent("@b:test", "$b")) + + cache.start_request(self._token(4)) + self.assertFalse(cache.was_sent("@a:test", "$a")) + + cache.mark_sent("@a:test", "$a", self._token(5)) + cache.start_request(None) + self.assertFalse(cache.was_sent("@a:test", "$a")) + self.assertEqual(len(cache), 0) + + def test_evicts_least_recently_used(self) -> None: + cache = RustLazyLoadedMembersCache() + for i in range(LAZY_LOADED_MEMBERS_CACHE_MAX_SIZE): + cache.mark_sent(f"@u{i}:test", f"$u{i}", self._token(1)) + self.assertTrue(cache.was_sent("@u0:test", "$u0")) + cache.mark_sent("@new:test", "$new", self._token(1)) + self.assertTrue(cache.was_sent("@u0:test", "$u0")) + self.assertFalse(cache.was_sent("@u1:test", "$u1")) + self.assertTrue(cache.was_sent("@new:test", "$new")) + + @parameterized.expand(SEEDS) + def test_random_operations_agree(self, seed: int) -> None: + """The same sequence of operations, with multi-writer positions, gives + the same answers from both caches.""" + from synapse.types import StreamKeyType, StreamToken + + rng = random.Random(seed) + python = LazyLoadedMembersCache() + rust = RustLazyLoadedMembersCache() + users = [f"@u{i}:test" for i in range(LAZY_LOADED_MEMBERS_CACHE_MAX_SIZE + 20)] + + def token(rng: random.Random) -> RoomStreamToken: + return random_live_token(rng) + + for _ in range(500): + op = rng.choice(["mark", "mark", "mark", "was", "was", "start"]) + if op == "mark": + user = rng.choice(users) + event_id = rng.choice(["$1", "$2", "$3"]) + position = token(rng) + python.mark_sent( + user, + event_id, + StreamToken.START.copy_and_replace(StreamKeyType.ROOM, position), + ) + rust.mark_sent(user, event_id, position) + elif op == "was": + user = rng.choice(users) + event_id = rng.choice(["$1", "$2", "$3"]) + self.assertEqual( + rust.was_sent(user, event_id), python.was_sent(user, event_id) + ) + else: + since = None if rng.random() < 0.1 else token(rng) + python.start_request( + None + if since is None + else StreamToken.START.copy_and_replace(StreamKeyType.ROOM, since) + ) + rust.start_request(since) + self.assertEqual(len(rust), len(python._sent))