diff --git a/rust/src/handlers/mod.rs b/rust/src/handlers/mod.rs index eba284f001..659bd1732c 100644 --- a/rust/src/handlers/mod.rs +++ b/rust/src/handlers/mod.rs @@ -81,6 +81,11 @@ impl RustHandlers { 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( + connection_pool.clone_ref(py), + runtime.clone(), + Arc::clone(&store), + ); let persister = if config.experimental.rust_event_persistence != "off" { Persister::for_homeserver( py, @@ -88,6 +93,7 @@ impl RustHandlers { connection_pool.clone_ref(py), runtime.clone(), Arc::clone(&store), + state_resolution.conflict_cache(), )? .map(|persister| Py::new(py, persister)) .transpose()? @@ -95,14 +101,7 @@ impl RustHandlers { None }; ( - Some(Py::new( - py, - StateResolutionHandler::new( - connection_pool.clone_ref(py), - runtime.clone(), - Arc::clone(&store), - ), - )?), + Some(Py::new(py, state_resolution)?), Some(Py::new( py, SyncHandler::new(connection_pool, runtime.clone(), store, &config), diff --git a/rust/src/persist/context.rs b/rust/src/persist/context.rs index 16806c5c93..18d2755165 100644 --- a/rust/src/persist/context.rs +++ b/rust/src/persist/context.rs @@ -42,7 +42,8 @@ pub struct PersistContext { pub state_delta_due_to_event: Option, /// Deltas between pairs of state groups `(prev, current)` known when the /// context was built, which can save a state resolution. - pub state_group_deltas: HashMap<(i64, i64), StateMap>, + /// Shared, since every event of a chunk carries a copy. + pub state_group_deltas: Arc>, /// Whether the state at the event is partial. pub partial_state: bool, /// Whether an application service sent the event. Only used to label @@ -79,7 +80,7 @@ impl<'a, 'py> FromPyObject<'a, 'py> for PersistContext { .getattr(intern!(py, "state_group_before_event"))? .extract()?, state_delta_due_to_event, - state_group_deltas, + state_group_deltas: Arc::new(state_group_deltas), partial_state: context.getattr(intern!(py, "partial_state"))?.extract()?, app_service: !context.getattr(intern!(py, "app_service"))?.is_none(), }) diff --git a/rust/src/persist/extremities.rs b/rust/src/persist/extremities.rs index 68214dfd8a..41adf27b6e 100644 --- a/rust/src/persist/extremities.rs +++ b/rust/src/persist/extremities.rs @@ -31,49 +31,94 @@ use std::{ use anyhow::{anyhow, bail}; use async_trait::async_trait; use log::{error, info}; +use serde_json::Value; use super::{context::PersistContext, metrics::PersistCounts, reader::PersistReader}; -use crate::{events::Event, storage::events::RedactBehaviour}; +use crate::{events::Event, state::StateMap, storage::events::RedactBehaviour}; /// `EventTypes.Dummy`. const DUMMY: &str = "org.matrix.dummy_event"; +/// `EventTypes.Create`. +const CREATE: &str = "m.room.create"; -/// What the extremity calculations read from an event being persisted. +/// What the pre-transaction computation reads from an event being +/// persisted and its context. #[derive(Clone, Debug, Default)] pub struct NewEvent { pub event_id: Arc, + pub event_type: Arc, + /// The state key, if it is a state event. In an MSC4242 room, that is + /// what puts it in the state DAG. + pub state_key: Option>, pub sender: Arc, pub depth: i64, pub prev_event_ids: Vec>, /// The MSC4242 `prev_state_events`, if the event's format has them. pub prev_state_events: Option>>, - /// Whether the event has a state key, which in an MSC4242 room is - /// what puts it in the state DAG. - pub is_state: bool, + /// `content.membership`, if it is a string. + pub membership: Option>, + /// For a create event, `content.room_version` as + /// `_get_new_state_after_events` reads it: `"1"` if it is missing, and + /// `None` if it is falsy in Python's terms, which sends the reader to + /// the database. A value that is not a string fails to parse as a + /// room version, as Python's `KNOWN_ROOM_VERSIONS` lookup fails for it. + pub create_room_version: Option>, pub outlier: bool, pub soft_failed: bool, pub rejected: bool, + /// The context's state group, the state after the event. `None` for an + /// outlier. + pub state_group: Option, + /// The context's `state_group_deltas`. + pub state_group_deltas: Arc>, } impl NewEvent { pub fn new(event: &Event, context: &PersistContext) -> anyhow::Result { let metadata = event.internal_metadata(); + let content = event.content(); + let membership = match content.get_field("membership") { + Some(Value::String(membership)) => Some(Arc::from(membership.as_str())), + _ => None, + }; + let create_room_version = if event.type_() == CREATE && event.state_key() == Some("") { + match content.get_field("room_version") { + None => Some(Arc::from("1")), + Some(Value::String(version)) if version.is_empty() => None, + Some(Value::String(version)) => Some(Arc::from(version.as_str())), + Some(value) if is_falsy(value) => None, + // Prefixed so that a number such as `10` is not read as + // the room version `"10"`. + Some(value) => Some(Arc::from(format!("not a string: {value}"))), + } + } else { + None + }; Ok(NewEvent { event_id: Arc::from(event.event_id()), + event_type: Arc::from(event.type_()), + state_key: event.state_key().map(Arc::from), sender: Arc::from(event.sender()), depth: event.depth(), prev_event_ids: event.prev_event_ids().into_iter().map(Arc::from).collect(), prev_state_events: event .prev_state_event_ids() .map(|ids| ids.iter().map(|id| Arc::from(id.as_str())).collect()), - is_state: event.state_key().is_some(), + membership, + create_room_version, outlier: metadata.is_outlier()?, soft_failed: metadata.soft_failed()?, rejected: context.is_rejected(), + state_group: context.state_group, + state_group_deltas: Arc::clone(&context.state_group_deltas), }) } - /// The `prev_state_events`, which an event in the state DAG must have. + pub fn is_state(&self) -> bool { + self.state_key.is_some() + } + + /// The prev state events, which an event in the state DAG must have. fn prev_state(&self) -> anyhow::Result<&[Arc]> { self.prev_state_events .as_deref() @@ -81,6 +126,18 @@ impl NewEvent { } } +/// Whether Python would read a JSON value as false. +fn is_falsy(value: &Value) -> bool { + match value { + Value::Null => true, + Value::Bool(b) => !b, + Value::Number(n) => n.as_f64() == Some(0.0), + Value::String(s) => s.is_empty(), + Value::Array(a) => a.is_empty(), + Value::Object(o) => o.is_empty(), + } +} + /// What `_prune_extremities` reads from an extremity it might drop, and /// from the prev events of a local dummy event. #[derive(Clone, Debug, Default)] @@ -245,7 +302,7 @@ pub async fn calculate_new_state_dag_extremities( events: &[NewEvent], counts: &mut PersistCounts, ) -> anyhow::Result>> { - let state_events: Vec<&NewEvent> = events.iter().filter(|event| event.is_state).collect(); + let state_events: Vec<&NewEvent> = events.iter().filter(|event| event.is_state()).collect(); if state_events.is_empty() { return Ok(existing.clone()); } @@ -434,7 +491,7 @@ fn domain_from_id(id: &str) -> anyhow::Result<&str> { } /// `HomeServer.is_mine_id`: false for a malformed ID. -fn is_mine_id(id: &str, server_name: &str) -> bool { +pub(super) fn is_mine_id(id: &str, server_name: &str) -> bool { id.split_once(':') .is_some_and(|(_, domain)| domain == server_name) } @@ -556,7 +613,7 @@ mod tests { fn state_event(id: &str, prev_state: &[&str]) -> NewEvent { NewEvent { - is_state: true, + state_key: Some(Arc::from("")), prev_state_events: Some(prev_state.iter().map(|p| Arc::from(*p)).collect()), ..new_event(id, &[]) } diff --git a/rust/src/persist/mod.rs b/rust/src/persist/mod.rs index f95dce91cc..2b9d88a562 100644 --- a/rust/src/persist/mod.rs +++ b/rust/src/persist/mod.rs @@ -25,6 +25,7 @@ //! - [`metrics`]: counts for the Python metrics, flushed when scraped. //! - [`python`]: the [`python::Persister`] and test helpers. //! - [`reader`]: what persisting reads before its write transaction. +//! - [`state_delta`]: the room's new current state after a chunk. //! //! Every GIL hold in this tree goes through [`crate::gil::with_gil`] with a //! site named `persist_*`, so that tests can pin the count per site. @@ -37,6 +38,7 @@ pub mod extremities; pub mod metrics; pub mod python; pub mod reader; +pub mod state_delta; /// Called when registering modules with python. pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { diff --git a/rust/src/persist/python.rs b/rust/src/persist/python.rs index 9b382b0a70..a46de3f35c 100644 --- a/rust/src/persist/python.rs +++ b/rust/src/persist/python.rs @@ -29,6 +29,7 @@ use crate::{ events::Event, homeserver::HomeServer, runtime::RustRuntime, + state::ConflictCache, storage::{ native_pool::NativePool, store::Store, @@ -64,6 +65,8 @@ pub(super) struct PersisterInner { pub server_name: String, pub instance_name: String, pub metrics: PersistMetrics, + /// The Rust state resolution handler's conflict cache. + pub conflict_cache: Arc, } impl Persister { @@ -76,6 +79,7 @@ impl Persister { connection_pool: Py, runtime: RustRuntime, store: Arc, + conflict_cache: Arc, ) -> PyResult> { let instance_name = homeserver.instance_name(py)?; if !homeserver.event_writers(py)?.contains(&instance_name) { @@ -112,6 +116,7 @@ impl Persister { server_name: homeserver.hostname(py)?, instance_name, metrics: PersistMetrics::default(), + conflict_cache, }), })) } @@ -178,7 +183,7 @@ pub fn persist_context_for_test<'py>( .transpose()?, )?; let deltas = PyDict::new(py); - for (key, delta) in &context.state_group_deltas { + for (key, delta) in context.state_group_deltas.iter() { deltas.set_item(key, state_map_to_py(py, delta)?)?; } dict.set_item("state_group_deltas", deltas)?; diff --git a/rust/src/persist/reader.rs b/rust/src/persist/reader.rs index 1e575d0ee9..ed66b28251 100644 --- a/rust/src/persist/reader.rs +++ b/rust/src/persist/reader.rs @@ -50,13 +50,17 @@ use super::{ extremities::{self, NewEvent}, metrics::PersistCounts, python::persist_error_to_py, + state_delta::{self, DeltaState, NativeStateResolver, StateAndExtremities}, }; use crate::database::postgres::pool::{ConnectionPool, PyConnectionPool}; use crate::events::Event; +use crate::state::{ConflictCache, EventMap, StateKey, StateMap, StateResEvent}; +use crate::storage::caches::StoreCaches; use crate::storage::events::Querier; use crate::storage::sync_store::{ PersistedEventPosition, SyncReader, SyncStore, SELECT_MANY_BATCH_SIZE, }; +use crate::types::state_filter::state_map_to_py; /// How many events `_get_events_which_are_prevs` and /// `_get_prevs_before_rejected` handle per transaction. @@ -692,6 +696,8 @@ pub struct PyPersistReaderForTest { store: Arc, pool: ConnectionPool, handle: Handle, + caches: Arc, + conflict_cache: Arc, } impl PyPersistReaderForTest { @@ -707,6 +713,62 @@ impl PyPersistReaderForTest { }) .map_err(persist_error_to_py) } + + /// Run one of the two orchestrating steps. + fn run_orchestration( + &self, + py: Python<'_>, + room_id: String, + server_name: String, + events_and_contexts: Vec<(Py, PersistContext)>, + now_ms: i64, + state_dag: bool, + ) -> PyResult<(StateAndExtremities, PersistCounts)> { + let (events, event_map) = chunk(py, &events_and_contexts)?; + let resolver = self.resolver(); + self.run(py, |reader| { + async move { + let mut counts = PersistCounts::default(); + let result = if state_dag { + state_delta::process_state_dag_forward_extremities_and_state_delta( + &reader, + &resolver, + &room_id, + &server_name, + &events, + &event_map, + now_ms, + &mut counts, + ) + .await? + } else { + state_delta::calculate_new_forward_extremities_and_state_delta( + &reader, + &resolver, + &room_id, + &server_name, + &events, + &event_map, + now_ms, + &mut counts, + ) + .await? + }; + Ok((result, counts)) + } + .boxed() + }) + } + + /// A state resolver over the same pool. + fn resolver(&self) -> NativeStateResolver { + NativeStateResolver { + pool: self.pool.clone(), + handle: self.handle.clone(), + caches: Arc::clone(&self.caches), + conflict_cache: Arc::clone(&self.conflict_cache), + } + } } #[pymethods] @@ -718,6 +780,8 @@ impl PyPersistReaderForTest { store: Arc::clone(&persister.inner.sync), pool: pool.pool.clone(), handle: pool.handle.clone(), + caches: Arc::clone(&persister.inner.store.caches), + conflict_cache: Arc::clone(&persister.inner.conflict_cache), } } @@ -1076,6 +1140,218 @@ impl PyPersistReaderForTest { })?; Ok((strings(result), counts.to_py(py)?)) } + + // The state delta over the reader and the native state resolver, each + // returning its result and what it counted. + + /// `(current_state, delta_ids, new_latest_event_ids, counts)`. + #[allow(clippy::too_many_arguments, clippy::type_complexity)] + fn get_new_state_after_events<'py>( + &self, + py: Python<'py>, + room_id: String, + server_name: String, + events_and_contexts: Vec<(Py, PersistContext)>, + old_latest_event_ids: Vec, + new_latest_event_ids: Vec, + should_prune: bool, + now_ms: i64, + ) -> PyResult<( + Option>, + Option>, + HashSet, + Bound<'py, PyDict>, + )> { + let (events, event_map) = chunk(py, &events_and_contexts)?; + let old: HashSet> = arcs(old_latest_event_ids).into_iter().collect(); + let new: HashSet> = arcs(new_latest_event_ids).into_iter().collect(); + let resolver = self.resolver(); + let (result, counts) = self.run(py, |reader| { + async move { + let mut counts = PersistCounts::default(); + let result = state_delta::get_new_state_after_events( + &reader, + &resolver, + &room_id, + &server_name, + &events, + &event_map, + &old, + new, + should_prune, + now_ms, + &mut counts, + ) + .await?; + Ok((result, counts)) + } + .boxed() + })?; + Ok(( + result + .current_state + .map(|state| state_map_to_py(py, &state)) + .transpose()?, + result + .delta_ids + .map(|delta| state_map_to_py(py, &delta)) + .transpose()?, + strings(result.new_latest_event_ids), + counts.to_py(py)?, + )) + } + + /// `(to_delete, to_insert)`. + #[allow(clippy::type_complexity)] + fn calculate_state_delta<'py>( + &self, + py: Python<'py>, + room_id: String, + current_state: HashMap<(String, String), String>, + ) -> PyResult<(Vec<(String, String)>, Bound<'py, PyDict>)> { + let current_state = state_map(current_state); + let delta = self.run(py, |reader| { + async move { + state_delta::calculate_state_delta(&reader, &room_id, ¤t_state).await + } + .boxed() + })?; + let (to_delete, to_insert, _) = delta_to_py(py, &delta)?; + Ok((to_delete, to_insert)) + } + + fn is_server_still_joined( + &self, + py: Python<'_>, + room_id: String, + server_name: String, + events_and_contexts: Vec<(Py, PersistContext)>, + to_delete: Vec<(String, String)>, + to_insert: HashMap<(String, String), String>, + ) -> PyResult { + let (events, _) = chunk(py, &events_and_contexts)?; + let delta = DeltaState { + to_delete: to_delete + .iter() + .map(|(event_type, state_key)| StateKey::new(event_type, state_key)) + .collect(), + to_insert: state_map(to_insert), + no_longer_in_room: false, + }; + self.run(py, |reader| { + async move { + state_delta::is_server_still_joined( + &reader, + &room_id, + &server_name, + &events, + &delta, + ) + .await + } + .boxed() + }) + } + + /// `(new_forward_extremities, delta, counts)`, the delta as + /// `(to_delete, to_insert, no_longer_in_room)`. + #[allow(clippy::type_complexity)] + fn calculate_new_forward_extremities_and_state_delta<'py>( + &self, + py: Python<'py>, + room_id: String, + server_name: String, + events_and_contexts: Vec<(Py, PersistContext)>, + now_ms: i64, + ) -> PyResult<( + Option>, + Option>, + Bound<'py, PyDict>, + )> { + let (result, counts) = + self.run_orchestration(py, room_id, server_name, events_and_contexts, now_ms, false)?; + Ok(( + result.new_forward_extremities.map(strings), + result + .state_delta + .map(|delta| delta_to_py(py, &delta)) + .transpose()?, + counts.to_py(py)?, + )) + } + + /// `(new_forward_extremities, delta, new_state_dag_extremities, counts)`, + /// the delta as `(to_delete, to_insert, no_longer_in_room)`. + #[allow(clippy::type_complexity)] + fn process_state_dag_forward_extremities_and_state_delta<'py>( + &self, + py: Python<'py>, + room_id: String, + server_name: String, + events_and_contexts: Vec<(Py, PersistContext)>, + now_ms: i64, + ) -> PyResult<( + Option>, + Option>, + Option>, + Bound<'py, PyDict>, + )> { + let (result, counts) = + self.run_orchestration(py, room_id, server_name, events_and_contexts, now_ms, true)?; + Ok(( + result.new_forward_extremities.map(strings), + result + .state_delta + .map(|delta| delta_to_py(py, &delta)) + .transpose()?, + result.new_state_dag_extremities.map(strings), + counts.to_py(py)?, + )) + } +} + +/// A delta as Python's test reads it: `(to_delete, to_insert, +/// no_longer_in_room)`. +type PyDelta<'py> = (Vec<(String, String)>, Bound<'py, PyDict>, bool); + +fn delta_to_py<'py>(py: Python<'py>, delta: &DeltaState) -> PyResult> { + Ok(( + delta + .to_delete + .iter() + .map(|key| (key.0.to_string(), key.1.to_string())) + .collect(), + state_map_to_py(py, &delta.to_insert)?, + delta.no_longer_in_room, + )) +} + +fn state_map(map: HashMap<(String, String), String>) -> StateMap { + map.into_iter() + .map(|((event_type, state_key), event_id)| { + (StateKey::new(&event_type, &state_key), Arc::from(event_id)) + }) + .collect() +} + +/// The chunk's events as the computation reads them, and as state +/// resolution's map of events in flight. +fn chunk( + py: Python<'_>, + events_and_contexts: &[(Py, PersistContext)], +) -> PyResult<(Vec, EventMap)> { + let events = new_events(py, events_and_contexts)?; + let event_map = events_and_contexts + .iter() + .map(|(event, _)| { + let event = event.bind(py).get(); + ( + Arc::from(event.event_id()), + StateResEvent::new(event.shallow_clone()), + ) + }) + .collect(); + Ok((events, event_map)) } fn new_events( diff --git a/rust/src/persist/state_delta.rs b/rust/src/persist/state_delta.rs new file mode 100644 index 0000000000..451a23f113 --- /dev/null +++ b/rust/src/persist/state_delta.rs @@ -0,0 +1,1603 @@ +/* + * 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's new current state after a chunk of events: ports of +//! `_get_new_state_after_events`, `_calculate_state_delta`, +//! `_is_server_still_joined` and the two steps that call them with the +//! extremity calculations, `_calculate_new_forward_extremities_and_state_delta` +//! and `_process_state_dag_forward_extremities_and_state_delta`, from +//! `synapse/storage/controllers/persist_events.py`. +//! +//! The reads go through [`StateDeltaReads`], and state resolution through a +//! [`StateResolver`], so that the unit tests can fake both. +//! +//! State resolution always runs in Rust, with the conflict cache of the +//! Rust state resolution handler. It differs from the Python +//! `StateResolutionHandler.resolve_state_groups` in these ways: +//! +//! - There is no `_state_cache` of earlier resolutions by group set and no +//! linearizer. Every call resolves, through the conflict cache, and checks +//! and bumps the input groups as a cache miss does in Python. +//! - Rooms whose state resolution is v1, and rooms without a chain cover +//! index, fail with [`StateResFallback`]. The caller should run the +//! Python step instead. +//! - When two input groups have the same state, the one the pruning keeps +//! is the lowest group ID. Python takes the first in its set's order. +//! - [`NativeStateResolver`] reads and bumps the state groups on the main +//! database's pool. Python uses the state database. Like the Rust state +//! resolution store, this needs the two to be the same database. +//! +//! A step can fail with [`StateResFallback`] after it has added to the +//! caller's [`PersistCounts`] and bumped the input groups. A caller that +//! falls back to Python should drop those counts, or the chunk is counted +//! twice. The second bump is harmless. + +use std::{ + collections::{HashMap, HashSet}, + fmt, + sync::Arc, +}; + +use anyhow::{anyhow, bail}; +use async_trait::async_trait; +use futures::FutureExt; +use log::info; +use tokio::runtime::Handle; + +use super::{ + extremities::{ + calculate_new_extremities, calculate_new_state_dag_extremities, is_mine_id, + prune_extremities, ExtremityReads, NewEvent, + }, + metrics::PersistCounts, + reader::PersistReader, +}; +use crate::{ + clock::Clock, + database::postgres::{ + pool::ConnectionPool, + write::{run_write_interaction, WriteError}, + }, + room_versions::{RoomVersion, StateResolutionVersions}, + state::{ + matching_group, resolve_on_pool, store::NoChainCoverIndex, v2, ConflictCache, EventMap, + StateKey, StateMap, StateResStats, + }, + storage::caches::StoreCaches, + types::state_filter::StateFilter, +}; + +/// `EventTypes.Member`. +const MEMBER: &str = "m.room.member"; +/// `EventTypes.Create`. +const CREATE: &str = "m.room.create"; + +/// `DeltaState`: how the room's current state changes. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct DeltaState { + /// The keys to remove, in key order. + pub to_delete: Vec, + /// The entries to add or replace. + pub to_insert: StateMap, + /// The server has left the room, so its current state goes. + pub no_longer_in_room: bool, +} + +/// What `_get_new_state_after_events` returns. If both state maps are +/// `None`, the state has not changed. +#[derive(Debug, Default)] +pub struct NewState { + /// The full new current state, if it was worked out. + pub current_state: Option, + /// The change to the current state, if the contexts knew it. + pub delta_ids: Option, + /// The new forward extremities, possibly pruned. + pub new_latest_event_ids: HashSet>, + /// What the state resolution recorded, if one ran. + pub state_res_stats: Option, +} + +/// What the two orchestrating steps return: the inputs of the write +/// transaction. +#[derive(Debug, Default)] +pub struct StateAndExtremities { + /// The room's new forward extremities, or `None` if they do not change. + pub new_forward_extremities: Option>>, + /// The change to the room's current state, if any. + pub state_delta: Option, + /// The new MSC4242 state DAG forward extremities. Only for a state DAG + /// room. + pub new_state_dag_extremities: Option>>, + /// What the state resolution recorded, if one ran. + pub state_res_stats: Option, +} + +/// The Rust state resolution cannot resolve this room: the caller should +/// run the Python step. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct StateResFallback(pub String); + +impl fmt::Display for StateResFallback { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "the Rust state resolution cannot resolve this: {}", + self.0 + ) + } +} + +impl std::error::Error for StateResFallback {} + +/// The reads the state delta makes, on top of the extremity reads. +#[async_trait] +pub trait StateDeltaReads: ExtremityReads { + /// See [`PersistReader::latest_event_ids_in_room`]. + async fn latest_event_ids_in_room(&self, room_id: &str) -> anyhow::Result>>; + + /// See [`PersistReader::state_dag_extremities`]. + async fn state_dag_extremities(&self, room_id: &str) -> anyhow::Result>>; + + /// `_get_state_group_for_events`: the state group of each event that + /// has one. + async fn state_group_for_events( + &self, + event_ids: &[Arc], + ) -> anyhow::Result, i64>>; + + /// `_get_state_for_groups`: the full state of each group. + async fn state_for_groups(&self, groups: &[i64]) + -> anyhow::Result>>; + + /// See [`PersistReader::room_version_id`]. + async fn room_version_id(&self, room_id: &str) -> anyhow::Result>>; + + /// `get_partial_current_state_ids`: the room's current state, partial + /// or not. + async fn current_state_ids(&self, room_id: &str) -> anyhow::Result; + + /// `get_membership_from_event_ids`: the membership of each of the + /// events that is a membership event. + async fn membership_from_event_ids( + &self, + event_ids: &[Arc], + ) -> anyhow::Result, Arc>>; + + /// See [`PersistReader::is_local_host_in_room_ignoring_users`]. + async fn is_local_host_in_room_ignoring_users( + &self, + room_id: &str, + ignore_users: &[Arc], + ) -> anyhow::Result; +} + +#[async_trait] +impl StateDeltaReads for PersistReader<'_> { + async fn latest_event_ids_in_room(&self, room_id: &str) -> anyhow::Result>> { + PersistReader::latest_event_ids_in_room(self, room_id).await + } + + async fn state_dag_extremities(&self, room_id: &str) -> anyhow::Result>> { + PersistReader::state_dag_extremities(self, room_id).await + } + + async fn state_group_for_events( + &self, + event_ids: &[Arc], + ) -> anyhow::Result, i64>> { + (**self).state_group_for_events(event_ids).await + } + + async fn state_for_groups( + &self, + groups: &[i64], + ) -> anyhow::Result>> { + (**self).state_for_groups(groups).await + } + + async fn room_version_id(&self, room_id: &str) -> anyhow::Result>> { + PersistReader::room_version_id(self, room_id).await + } + + async fn current_state_ids(&self, room_id: &str) -> anyhow::Result { + (**self) + .current_state_ids(room_id, &StateFilter::all()) + .await + } + + async fn membership_from_event_ids( + &self, + event_ids: &[Arc], + ) -> anyhow::Result, Arc>> { + let memberships = PersistReader::membership_from_event_ids(self, event_ids).await?; + Ok(memberships + .into_iter() + .map(|(event_id, membership)| (event_id, membership.membership)) + .collect()) + } + + async fn is_local_host_in_room_ignoring_users( + &self, + room_id: &str, + ignore_users: &[Arc], + ) -> anyhow::Result { + PersistReader::is_local_host_in_room_ignoring_users(self, room_id, ignore_users).await + } +} + +/// What the state delta needs from state resolution. +#[async_trait] +pub trait StateResolver: Sync { + /// `check_state_groups_and_bump_deletion`: those of the groups that do + /// not exist. The deletion of the others, if pending, is put off. + async fn check_state_groups_and_bump_deletion( + &self, + groups: &[i64], + ) -> anyhow::Result>; + + /// Resolve the state sets. `event_map` holds the events being + /// persisted. Fails with [`StateResFallback`] if the Rust state + /// resolution cannot resolve the room. + async fn resolve( + &self, + room_id: &str, + room_version: &'static RoomVersion, + state_sets: &[Arc], + event_map: &EventMap, + ) -> anyhow::Result<(StateMap, StateResStats)>; +} + +/// The [`StateResolver`] over the native pool. +pub struct NativeStateResolver { + pub pool: ConnectionPool, + pub handle: Handle, + pub caches: Arc, + pub conflict_cache: Arc, +} + +#[async_trait] +impl StateResolver for NativeStateResolver { + async fn check_state_groups_and_bump_deletion( + &self, + groups: &[i64], + ) -> anyhow::Result> { + let existing = run_write_interaction( + &self.pool, + "check_state_groups_and_bump_deletion", + |txn, _| { + async move { + // Read on each attempt, as Python reads the clock inside + // the transaction. + let now_ms = Clock::new().now_millis() as i64; + let rows = txn + .query("SELECT id FROM state_groups WHERE id = ANY($1)", &[&groups]) + .await?; + let existing = rows + .iter() + .map(|row| row.try_get::<_, i64>(0)) + .collect::, _>>() + .map_err(|e| WriteError::Other(e.into()))?; + if !existing.is_empty() { + txn.execute( + "UPDATE state_groups_pending_deletion + SET sequence_number = DEFAULT, insertion_ts = $1 + WHERE state_group = ANY($2)", + &[&now_ms, &existing], + ) + .await?; + } + Ok(existing) + } + .boxed() + }, + ) + .await?; + let existing: HashSet = existing.into_iter().collect(); + Ok(groups + .iter() + .copied() + .filter(|group| !existing.contains(group)) + .collect()) + } + + async fn resolve( + &self, + room_id: &str, + room_version: &'static RoomVersion, + state_sets: &[Arc], + event_map: &EventMap, + ) -> anyhow::Result<(StateMap, StateResStats)> { + if room_version.state_res == StateResolutionVersions::V1 { + return Err(v1_fallback(room_version)); + } + let (unconflicted, conflicted) = v2::separate(state_sets); + if conflicted.is_empty() { + return Ok((unconflicted, StateResStats::default())); + } + let output = resolve_on_pool( + &self.pool, + self.handle.clone(), + Arc::clone(&self.caches), + &self.conflict_cache, + room_id, + room_version, + state_sets, + (unconflicted, conflicted), + event_map.clone(), + StateResStats::default(), + ) + .await + .map_err(|e| match e.downcast::() { + Ok(no_index) => StateResFallback(no_index.to_string()).into(), + Err(e) => e, + })?; + Ok((output.state, output.stats)) + } +} + +/// The [`StateResFallback`] for a room on state resolution v1. +fn v1_fallback(room_version: &RoomVersion) -> anyhow::Error { + StateResFallback(format!( + "room version {} uses state resolution v1", + room_version.identifier + )) + .into() +} + +/// Look up the state group of each event, failing for one without, as +/// Python's dict lookup does. +fn groups_of<'a>( + event_ids: impl IntoIterator>, + event_id_to_state_group: &HashMap, i64>, +) -> anyhow::Result> { + event_ids + .into_iter() + .map(|event_id| { + event_id_to_state_group + .get(event_id) + .copied() + .ok_or_else(|| anyhow!("no state group for {event_id}")) + }) + .collect() +} + +/// `_get_new_state_after_events`: the room's state after `events`, given +/// that its forward extremities go from `old_latest_event_ids` to +/// `new_latest_event_ids`. +/// +/// If the extremities' state groups do not change, both state maps are +/// `None`. If a context knows the delta between the old and new group, that +/// comes back as `delta_ids`. Otherwise the full state comes back, resolved +/// if there are several new groups. After a resolution whose result is one +/// of the new groups, and if `should_prune`, the extremities in other groups +/// may be pruned: see [`prune_extremities`]. +#[allow(clippy::too_many_arguments)] +pub async fn get_new_state_after_events( + reads: &dyn StateDeltaReads, + resolver: &dyn StateResolver, + room_id: &str, + server_name: &str, + events: &[NewEvent], + event_map: &EventMap, + old_latest_event_ids: &HashSet>, + new_latest_event_ids: HashSet>, + should_prune: bool, + now_ms: i64, + counts: &mut PersistCounts, +) -> anyhow::Result { + // The deltas the contexts know, from (prev group, new group). A later + // event's replaces an earlier one's. + let mut state_group_deltas: HashMap<(i64, i64), &StateMap> = HashMap::new(); + for event in events { + // Python's `context.state_group` raises for a rejected event. + if event.rejected { + bail!("Attempt to access state_group of rejected event"); + } + if event.state_group.is_none() { + // This should only happen for outlier events. + if !event.outlier { + bail!( + "Context for new event {} has no state group", + event.event_id + ); + } + continue; + } + state_group_deltas.extend(event.state_group_deltas.iter().map(|(k, v)| (*k, v))); + } + + // The state group of each extremity: from the chunk if the event is + // being persisted, and otherwise from the database. None of the old + // extremities are in the chunk. + let mut missing: HashSet> = old_latest_event_ids.clone(); + let mut event_id_to_state_group: HashMap, i64> = HashMap::new(); + for event_id in &new_latest_event_ids { + let in_chunk = events + .iter() + .find_map(|event| event.state_group.filter(|_| event.event_id == *event_id)); + match in_chunk { + Some(group) => { + event_id_to_state_group.insert(Arc::clone(event_id), group); + } + None => { + missing.insert(Arc::clone(event_id)); + } + } + } + if !missing.is_empty() { + let missing: Vec> = missing.into_iter().collect(); + event_id_to_state_group.extend(reads.state_group_for_events(&missing).await?); + } + + let old_state_groups = groups_of(old_latest_event_ids, &event_id_to_state_group)?; + let new_state_groups = groups_of(&new_latest_event_ids, &event_id_to_state_group)?; + + if old_state_groups == new_state_groups { + return Ok(NewState { + new_latest_event_ids, + ..NewState::default() + }); + } + + if new_state_groups.len() == 1 && old_state_groups.len() == 1 { + // Going from one group to another: a context may know the delta. + let new_group = *new_state_groups.iter().next().expect("one group"); + let old_group = *old_state_groups.iter().next().expect("one group"); + if let Some(delta) = state_group_deltas.get(&(old_group, new_group)) { + return Ok(NewState { + delta_ids: Some((*delta).clone()), + new_latest_event_ids, + ..NewState::default() + }); + } + } + + let mut new_groups: Vec = new_state_groups.iter().copied().collect(); + new_groups.sort_unstable(); + let mut state_groups_map = reads.state_for_groups(&new_groups).await?; + + if new_groups.len() == 1 { + let state = state_groups_map.remove(&new_groups[0]).unwrap_or_default(); + return Ok(NewState { + current_state: Some((*state).clone()), + new_latest_event_ids, + ..NewState::default() + }); + } + + // Several groups: resolve them. The room version is in the create + // event, which may be in the chunk. + let room_version = match events + .iter() + .find(|event| &*event.event_type == CREATE && event.state_key.as_deref() == Some("")) + { + Some(create) => create.create_room_version.clone(), + None => None, + }; + let room_version = match room_version { + Some(version) => version, + None => reads + .room_version_id(room_id) + .await? + .ok_or_else(|| anyhow!("Room {room_id} not found"))?, + }; + let room_version: &'static RoomVersion = room_version + .parse() + .map_err(|e| anyhow!("{e}: {room_version:?}"))?; + if room_version.state_res == StateResolutionVersions::V1 { + return Err(v1_fallback(room_version)); + } + + let groups: Vec<(i64, Arc)> = new_groups + .iter() + .map(|group| { + let state = state_groups_map.get(group).cloned().unwrap_or_default(); + (*group, state) + }) + .collect(); + + info!("Resolving state for {room_id} with groups {new_groups:?}"); + let missing_groups = resolver + .check_state_groups_and_bump_deletion(&new_groups) + .await?; + if !missing_groups.is_empty() { + bail!("State groups have been deleted: {missing_groups:?}"); + } + let state_sets: Vec> = + groups.iter().map(|(_, state)| Arc::clone(state)).collect(); + let (state, stats) = resolver + .resolve(room_id, room_version, &state_sets, event_map) + .await?; + counts.state_resolutions += 1; + + // If the resolved state is one of the new groups, the extremities in + // the others may be stuck ones to prune. + let mut new_latest_event_ids = new_latest_event_ids; + if should_prune { + if let Some(resolved_group) = matching_group(&state, &groups) { + new_latest_event_ids = prune_extremities( + reads, + room_id, + server_name, + new_latest_event_ids, + resolved_group, + &event_id_to_state_group, + events, + now_ms, + counts, + ) + .await?; + } + } + + Ok(NewState { + current_state: Some(state), + delta_ids: None, + new_latest_event_ids, + state_res_stats: Some(stats), + }) +} + +/// `_calculate_state_delta`: the change from the room's current state to +/// `current_state`. +pub async fn calculate_state_delta( + reads: &dyn StateDeltaReads, + room_id: &str, + current_state: &StateMap, +) -> anyhow::Result { + let existing_state = reads.current_state_ids(room_id).await?; + + let mut to_delete: Vec = existing_state + .keys() + .filter(|key| !current_state.contains_key(*key)) + .cloned() + .collect(); + to_delete.sort_unstable(); + + let to_insert: StateMap = current_state + .iter() + .filter(|(key, event_id)| existing_state.get(*key) != Some(event_id)) + .map(|(key, event_id)| (key.clone(), Arc::clone(event_id))) + .collect(); + + Ok(DeltaState { + to_delete, + to_insert, + no_longer_in_room: false, + }) +} + +/// `_is_server_still_joined`: whether a local user is still joined after +/// `delta`. +pub async fn is_server_still_joined( + reads: &dyn StateDeltaReads, + room_id: &str, + server_name: &str, + events: &[NewEvent], + delta: &DeltaState, +) -> anyhow::Result { + let is_local_member = |key: &StateKey| &*key.0 == MEMBER && is_mine_id(&key.1, server_name); + + if !delta + .to_delete + .iter() + .chain(delta.to_insert.keys()) + .any(is_local_member) + { + // No local user's membership changed. + return Ok(true); + } + + // A local join in the new state, among the events being persisted or + // in the database. + let mut events_to_check = Vec::new(); + for (key, event_id) in &delta.to_insert { + if !is_local_member(key) { + continue; + } + for event in events { + if *event_id == event.event_id && event.membership.as_deref() == Some("join") { + return Ok(true); + } + } + // Python reads the event from the database whether or not it was in + // the chunk. + events_to_check.push(Arc::clone(event_id)); + } + if !events_to_check.is_empty() { + let memberships = reads.membership_from_event_ids(&events_to_check).await?; + if memberships + .values() + .any(|membership| &**membership == "join") + { + return Ok(true); + } + } + + // Otherwise, any other local user still joined. The users whose state + // changed are covered above. + let users_to_ignore: Vec> = delta + .to_insert + .keys() + .chain(delta.to_delete.iter()) + .filter(|key| is_local_member(key)) + .map(|key| Arc::clone(&key.1)) + .collect(); + reads + .is_local_host_in_room_ignoring_users(room_id, &users_to_ignore) + .await +} + +/// The delta from what [`get_new_state_after_events`] found, with +/// `no_longer_in_room` set if the last local user left. +async fn delta_from_new_state( + reads: &dyn StateDeltaReads, + room_id: &str, + server_name: &str, + events: &[NewEvent], + new_state: NewState, +) -> anyhow::Result> { + let mut delta = if let Some(delta_ids) = new_state.delta_ids { + // A known delta only adds or replaces state. + DeltaState { + to_insert: delta_ids, + ..DeltaState::default() + } + } else if let Some(current_state) = &new_state.current_state { + calculate_state_delta(reads, room_id, current_state).await? + } else { + return Ok(None); + }; + + if !is_server_still_joined(reads, room_id, server_name, events, &delta).await? { + info!("Server no longer in room {room_id}"); + delta.no_longer_in_room = true; + } + Ok(Some(delta)) +} + +/// `_calculate_new_forward_extremities_and_state_delta`: the new forward +/// extremities and the change to the current state after `events`. +/// +/// Fails if no forward extremity is left, where Python's `assert` raises +/// an `AssertionError`. +#[allow(clippy::too_many_arguments)] +pub async fn calculate_new_forward_extremities_and_state_delta( + reads: &dyn StateDeltaReads, + resolver: &dyn StateResolver, + room_id: &str, + server_name: &str, + events: &[NewEvent], + event_map: &EventMap, + now_ms: i64, + counts: &mut PersistCounts, +) -> anyhow::Result { + let latest_event_ids = reads.latest_event_ids_in_room(room_id).await?; + let new_latest_event_ids = + calculate_new_extremities(reads, events, &latest_event_ids, counts).await?; + + if new_latest_event_ids == latest_event_ids { + // No change in extremities, so no change in state. + return Ok(StateAndExtremities::default()); + } + + if new_latest_event_ids.is_empty() { + bail!("No forward extremities left!"); + } + + // A chain of single-prev non-state events cannot change the state. + if latest_event_ids.len() == 1 + && new_latest_event_ids.len() == 1 + && events + .iter() + .all(|event| event.prev_event_ids.len() == 1 && !event.is_state()) + { + return Ok(StateAndExtremities { + new_forward_extremities: Some(new_latest_event_ids), + ..StateAndExtremities::default() + }); + } + + counts.state_delta += 1; + if new_latest_event_ids.len() == 1 { + counts.state_delta_single_event += 1; + + // A rough guess at whether the delta could have been known when one + // of the events was created: a local event takes the forward + // extremities as its prev events. + if events.iter().any(|event| { + let prevs: HashSet<&Arc> = event.prev_event_ids.iter().collect(); + prevs.len() == latest_event_ids.len() + && latest_event_ids.iter().all(|id| prevs.contains(id)) + }) { + counts.state_delta_reuse_delta += 1; + } + } + + let mut new_state = get_new_state_after_events( + reads, + resolver, + room_id, + server_name, + events, + event_map, + &latest_event_ids, + new_latest_event_ids, + true, + now_ms, + counts, + ) + .await?; + if new_state.new_latest_event_ids.is_empty() { + bail!("No forward extremities left!"); + } + let new_forward_extremities = std::mem::take(&mut new_state.new_latest_event_ids); + let state_res_stats = new_state.state_res_stats.take(); + + let state_delta = delta_from_new_state(reads, room_id, server_name, events, new_state).await?; + + Ok(StateAndExtremities { + new_forward_extremities: Some(new_forward_extremities), + state_delta, + new_state_dag_extremities: None, + state_res_stats, + }) +} + +/// `_process_state_dag_forward_extremities_and_state_delta`: for an MSC4242 +/// state DAG room, the new forward extremities of both DAGs and the change +/// to the current state, which follows the state DAG's extremities. Those +/// are never pruned, as that would lose eventual delivery. +/// +/// Fails if no room DAG forward extremity is left, where Python's `assert` +/// raises an `AssertionError`. +#[allow(clippy::too_many_arguments)] +pub async fn process_state_dag_forward_extremities_and_state_delta( + reads: &dyn StateDeltaReads, + resolver: &dyn StateResolver, + room_id: &str, + server_name: &str, + events: &[NewEvent], + event_map: &EventMap, + now_ms: i64, + counts: &mut PersistCounts, +) -> anyhow::Result { + let existing_state_dag = reads.state_dag_extremities(room_id).await?; + let new_state_dag = + calculate_new_state_dag_extremities(reads, room_id, &existing_state_dag, events, counts) + .await?; + + let existing_room_dag = reads.latest_event_ids_in_room(room_id).await?; + let new_room_dag = calculate_new_extremities(reads, events, &existing_room_dag, counts).await?; + if new_room_dag.is_empty() { + bail!("No room dag forward extremities left in room {room_id}!"); + } + + if new_state_dag == existing_state_dag { + return Ok(StateAndExtremities { + new_forward_extremities: Some(new_room_dag), + state_delta: None, + new_state_dag_extremities: Some(new_state_dag), + state_res_stats: None, + }); + } + + let mut new_state = get_new_state_after_events( + reads, + resolver, + room_id, + server_name, + events, + event_map, + &existing_state_dag, + new_state_dag.clone(), + false, + now_ms, + counts, + ) + .await?; + let state_res_stats = new_state.state_res_stats.take(); + + let state_delta = delta_from_new_state(reads, room_id, server_name, events, new_state).await?; + + Ok(StateAndExtremities { + new_forward_extremities: Some(new_room_dag), + state_delta, + new_state_dag_extremities: Some(new_state_dag), + state_res_stats, + }) +} + +#[cfg(test)] +mod tests { + use std::sync::Mutex; + + use super::*; + use crate::persist::extremities::DagEvent; + + /// The database as the state delta reads it. + #[derive(Default)] + struct FakeReads { + latest: HashSet>, + state_dag: HashSet>, + event_groups: HashMap, i64>, + groups: HashMap>, + room_version: Option>, + current_state: StateMap, + memberships: HashMap, Arc>, + /// Local users joined to the room. + local_joined: Vec>, + dag_events: HashMap, DagEvent>, + /// What `is_local_host_in_room_ignoring_users` was asked to ignore. + ignored: Mutex>>>, + } + + #[async_trait] + impl ExtremityReads for FakeReads { + async fn events_which_are_prevs( + &self, + _event_ids: &[Arc], + ) -> anyhow::Result>> { + Ok(Vec::new()) + } + + async fn prevs_before_rejected( + &self, + _event_ids: &[Arc], + _include_soft_failed: bool, + ) -> anyhow::Result>> { + Ok(HashSet::new()) + } + + async fn have_seen_events( + &self, + event_ids: &[Arc], + ) -> anyhow::Result>> { + Ok(event_ids.iter().cloned().collect()) + } + + async fn dag_events(&self, event_ids: &[Arc]) -> anyhow::Result> { + Ok(event_ids + .iter() + .filter_map(|id| self.dag_events.get(id).cloned()) + .collect()) + } + } + + #[async_trait] + impl StateDeltaReads for FakeReads { + async fn latest_event_ids_in_room(&self, _: &str) -> anyhow::Result>> { + Ok(self.latest.clone()) + } + + async fn state_dag_extremities(&self, _: &str) -> anyhow::Result>> { + Ok(self.state_dag.clone()) + } + + async fn state_group_for_events( + &self, + event_ids: &[Arc], + ) -> anyhow::Result, i64>> { + Ok(event_ids + .iter() + .filter_map(|id| self.event_groups.get(id).map(|g| (Arc::clone(id), *g))) + .collect()) + } + + async fn state_for_groups( + &self, + groups: &[i64], + ) -> anyhow::Result>> { + Ok(groups + .iter() + .filter_map(|g| self.groups.get(g).map(|s| (*g, Arc::clone(s)))) + .collect()) + } + + async fn room_version_id(&self, _: &str) -> anyhow::Result>> { + Ok(self.room_version.clone()) + } + + async fn current_state_ids(&self, _: &str) -> anyhow::Result { + Ok(self.current_state.clone()) + } + + async fn membership_from_event_ids( + &self, + event_ids: &[Arc], + ) -> anyhow::Result, Arc>> { + Ok(event_ids + .iter() + .filter_map(|id| { + self.memberships + .get(id) + .map(|m| (Arc::clone(id), Arc::clone(m))) + }) + .collect()) + } + + async fn is_local_host_in_room_ignoring_users( + &self, + _: &str, + ignore_users: &[Arc], + ) -> anyhow::Result { + *self.ignored.lock().unwrap() = Some(ignore_users.to_vec()); + Ok(self + .local_joined + .iter() + .any(|user| !ignore_users.contains(user))) + } + } + + /// Resolves to a fixed state, and records what it was asked. + #[derive(Default)] + struct FakeResolver { + result: StateMap, + deleted_groups: Vec, + checked: Mutex>>, + resolved: Mutex>, + } + + #[async_trait] + impl StateResolver for FakeResolver { + async fn check_state_groups_and_bump_deletion( + &self, + groups: &[i64], + ) -> anyhow::Result> { + self.checked.lock().unwrap().push(groups.to_vec()); + Ok(groups + .iter() + .copied() + .filter(|g| self.deleted_groups.contains(g)) + .collect()) + } + + async fn resolve( + &self, + _room_id: &str, + room_version: &'static RoomVersion, + state_sets: &[Arc], + _event_map: &EventMap, + ) -> anyhow::Result<(StateMap, StateResStats)> { + self.resolved + .lock() + .unwrap() + .push((room_version.identifier, state_sets.len())); + Ok((self.result.clone(), StateResStats::default())) + } + } + + const ROOM: &str = "!r:here"; + const SERVER: &str = "here"; + + fn ids(ids: &[&str]) -> HashSet> { + ids.iter().map(|id| Arc::from(*id)).collect() + } + + fn state(entries: &[(&str, &str, &str)]) -> StateMap { + entries + .iter() + .map(|(t, k, e)| (StateKey::new(t, k), Arc::from(*e))) + .collect() + } + + fn event(id: &str, prevs: &[&str], group: i64) -> NewEvent { + NewEvent { + event_id: Arc::from(id), + event_type: Arc::from("m.room.message"), + sender: Arc::from("@a:remote"), + depth: 10, + prev_event_ids: prevs.iter().map(|p| Arc::from(*p)).collect(), + state_group: Some(group), + ..Default::default() + } + } + + fn member(id: &str, user: &str, membership: &str, group: i64) -> NewEvent { + NewEvent { + event_type: Arc::from(MEMBER), + state_key: Some(Arc::from(user)), + membership: Some(Arc::from(membership)), + ..event(id, &["$p"], group) + } + } + + fn block_on(f: F) -> F::Output { + futures::executor::block_on(f) + } + + fn new_state( + reads: &FakeReads, + resolver: &FakeResolver, + events: &[NewEvent], + old: &[&str], + new: &[&str], + ) -> anyhow::Result<(NewState, PersistCounts)> { + let mut counts = PersistCounts::default(); + let result = block_on(get_new_state_after_events( + reads, + resolver, + ROOM, + SERVER, + events, + &EventMap::new(), + &ids(old), + ids(new), + true, + 0, + &mut counts, + ))?; + Ok((result, counts)) + } + + #[test] + fn same_groups_change_nothing() { + let reads = FakeReads { + event_groups: [(Arc::from("$old"), 5)].into(), + ..Default::default() + }; + let events = [event("$new", &["$old"], 5)]; + let (result, _) = new_state( + &reads, + &FakeResolver::default(), + &events, + &["$old"], + &["$new"], + ) + .unwrap(); + assert!(result.current_state.is_none()); + assert!(result.delta_ids.is_none()); + assert_eq!(result.new_latest_event_ids, ids(&["$new"])); + } + + #[test] + fn a_known_delta_is_reused() { + let reads = FakeReads { + event_groups: [(Arc::from("$old"), 5)].into(), + ..Default::default() + }; + let delta = state(&[("m.room.topic", "", "$new")]); + let mut first = event("$new", &["$old"], 6); + // A delta between other groups is no use. + first.state_group_deltas = Arc::new([((4, 6), state(&[("x", "", "$y")]))].into()); + let mut second = event("$other", &["$old"], 6); + second.state_group_deltas = Arc::new([((5, 6), delta.clone())].into()); + let (result, _) = new_state( + &reads, + &FakeResolver::default(), + &[first, second], + &["$old"], + &["$new"], + ) + .unwrap(); + assert_eq!(result.delta_ids, Some(delta)); + assert!(result.current_state.is_none()); + } + + #[test] + fn one_new_group_is_the_state() { + let group_state = state(&[("m.room.create", "", "$c"), ("m.room.topic", "", "$new")]); + let reads = FakeReads { + event_groups: [(Arc::from("$old"), 5)].into(), + groups: [(6, Arc::new(group_state.clone()))].into(), + ..Default::default() + }; + let resolver = FakeResolver::default(); + let events = [event("$new", &["$old"], 6)]; + let (result, counts) = new_state(&reads, &resolver, &events, &["$old"], &["$new"]).unwrap(); + assert_eq!(result.current_state, Some(group_state)); + assert!(resolver.resolved.lock().unwrap().is_empty()); + assert_eq!(counts.state_resolutions, 0); + } + + #[test] + fn extremities_in_the_chunk_without_a_group_come_from_the_database() { + // `$out` is an outlier in the chunk, so its group is read. + let reads = FakeReads { + event_groups: [(Arc::from("$old"), 5), (Arc::from("$out"), 5)].into(), + ..Default::default() + }; + let mut outlier = event("$out", &["$old"], 0); + outlier.state_group = None; + outlier.outlier = true; + let (result, _) = new_state( + &reads, + &FakeResolver::default(), + &[outlier], + &["$old"], + &["$out"], + ) + .unwrap(); + assert!(result.current_state.is_none() && result.delta_ids.is_none()); + } + + #[test] + fn several_groups_are_resolved() { + let resolved = state(&[("m.room.create", "", "$c"), ("m.room.topic", "", "$a")]); + let reads = FakeReads { + event_groups: [(Arc::from("$old"), 5), (Arc::from("$b"), 7)].into(), + groups: [ + ( + 6, + Arc::new(state(&[ + ("m.room.create", "", "$c"), + ("m.room.topic", "", "$x"), + ])), + ), + (7, Arc::new(state(&[("m.room.create", "", "$c")]))), + ] + .into(), + room_version: Some(Arc::from("10")), + ..Default::default() + }; + let resolver = FakeResolver { + result: resolved.clone(), + ..Default::default() + }; + let events = [event("$a", &["$old"], 6)]; + let (result, counts) = + new_state(&reads, &resolver, &events, &["$old"], &["$a", "$b"]).unwrap(); + assert_eq!(result.current_state, Some(resolved)); + assert_eq!(*resolver.checked.lock().unwrap(), vec![vec![6, 7]]); + assert_eq!(*resolver.resolved.lock().unwrap(), vec![("10", 2)]); + assert_eq!(counts.state_resolutions, 1); + // The result is neither group, so nothing is pruned. + assert_eq!(counts.potential_times_prune_extremities, 0); + assert_eq!(result.new_latest_event_ids, ids(&["$a", "$b"])); + } + + /// A resolution that comes out as `$a`'s group, with `$b` an old remote + /// extremity in another group. + fn prunable() -> (FakeReads, FakeResolver, Vec) { + let group_a = state(&[("m.room.create", "", "$c"), ("m.room.topic", "", "$a")]); + let reads = FakeReads { + event_groups: [(Arc::from("$old"), 5), (Arc::from("$b"), 7)].into(), + groups: [ + (6, Arc::new(group_a.clone())), + (7, Arc::new(state(&[("m.room.create", "", "$c")]))), + ] + .into(), + room_version: Some(Arc::from("10")), + dag_events: [( + Arc::from("$b"), + DagEvent { + // The chunk's server, and over 20 events behind it. + sender: Arc::from("@b:remote"), + event_type: Arc::from("m.room.message"), + origin_server_ts: 0, + depth: 1, + ..Default::default() + }, + )] + .into(), + ..Default::default() + }; + let resolver = FakeResolver { + result: group_a, + ..Default::default() + }; + let mut a = event("$a", &["$old"], 6); + a.depth = 500; + (reads, resolver, vec![a]) + } + + #[test] + fn a_resolution_to_a_new_group_prunes() { + let (reads, resolver, events) = prunable(); + let (result, counts) = + new_state(&reads, &resolver, &events, &["$old"], &["$a", "$b"]).unwrap(); + assert_eq!(result.new_latest_event_ids, ids(&["$a"])); + assert_eq!(counts.potential_times_prune_extremities, 1); + assert_eq!(counts.times_pruned_extremities, 1); + } + + #[test] + fn the_state_dag_does_not_prune() { + let (reads, resolver, events) = prunable(); + let mut counts = PersistCounts::default(); + let result = block_on(get_new_state_after_events( + &reads, + &resolver, + ROOM, + SERVER, + &events, + &EventMap::new(), + &ids(&["$old"]), + ids(&["$a", "$b"]), + false, + 86_400_000 * 2, + &mut counts, + )) + .unwrap(); + assert_eq!(result.new_latest_event_ids, ids(&["$a", "$b"])); + assert_eq!(counts.potential_times_prune_extremities, 0); + } + + #[test] + fn the_room_version_comes_from_a_create_event_in_the_chunk() { + let (reads, resolver, mut events) = prunable(); + let mut create = event("$create", &[], 6); + create.event_type = Arc::from(CREATE); + create.state_key = Some(Arc::from("")); + create.create_room_version = Some(Arc::from("11")); + events.insert(0, create); + new_state(&reads, &resolver, &events, &["$old"], &["$a", "$b"]).unwrap(); + assert_eq!(*resolver.resolved.lock().unwrap(), vec![("11", 2)]); + + // A falsy version sends the reader to the database. + let resolver = FakeResolver::default(); + events[0].create_room_version = None; + new_state(&reads, &resolver, &events, &["$old"], &["$a", "$b"]).unwrap(); + assert_eq!(*resolver.resolved.lock().unwrap(), vec![("10", 2)]); + } + + #[test] + fn state_res_v1_falls_back() { + let (mut reads, resolver, events) = prunable(); + reads.room_version = Some(Arc::from("1")); + let err = new_state(&reads, &resolver, &events, &["$old"], &["$a", "$b"]).unwrap_err(); + assert!(err.downcast_ref::().is_some(), "{err:#}"); + assert!(resolver.checked.lock().unwrap().is_empty()); + } + + #[test] + fn deleted_groups_fail() { + let (reads, mut resolver, events) = prunable(); + resolver.deleted_groups = vec![7]; + let err = new_state(&reads, &resolver, &events, &["$old"], &["$a", "$b"]).unwrap_err(); + assert_eq!(err.to_string(), "State groups have been deleted: [7]"); + assert!(resolver.resolved.lock().unwrap().is_empty()); + } + + #[test] + fn bad_contexts_fail() { + let reads = FakeReads { + event_groups: [(Arc::from("$old"), 5)].into(), + ..Default::default() + }; + let mut rejected = event("$new", &["$old"], 6); + rejected.rejected = true; + let err = new_state( + &reads, + &FakeResolver::default(), + &[rejected], + &["$old"], + &["$x"], + ) + .unwrap_err(); + assert_eq!( + err.to_string(), + "Attempt to access state_group of rejected event" + ); + + let mut no_group = event("$new", &["$old"], 6); + no_group.state_group = None; + let err = new_state( + &reads, + &FakeResolver::default(), + &[no_group], + &["$old"], + &["$x"], + ) + .unwrap_err(); + assert_eq!( + err.to_string(), + "Context for new event $new has no state group" + ); + + // An extremity with no state group anywhere. + let err = new_state( + &reads, + &FakeResolver::default(), + &[event("$new", &["$old"], 6)], + &["$old"], + &["$unknown"], + ) + .unwrap_err(); + assert_eq!(err.to_string(), "no state group for $unknown"); + } + + #[test] + fn state_delta_deletes_and_inserts() { + let reads = FakeReads { + current_state: state(&[ + ("m.room.create", "", "$c"), + ("m.room.topic", "", "$t1"), + ("m.room.name", "", "$n"), + ("m.room.avatar", "", "$av"), + ]), + ..Default::default() + }; + let new = state(&[ + ("m.room.create", "", "$c"), + ("m.room.topic", "", "$t2"), + ("m.room.join_rules", "", "$j"), + ]); + let delta = block_on(calculate_state_delta(&reads, ROOM, &new)).unwrap(); + assert_eq!( + delta.to_delete, + vec![ + StateKey::new("m.room.avatar", ""), + StateKey::new("m.room.name", "") + ] + ); + assert_eq!( + delta.to_insert, + state(&[("m.room.topic", "", "$t2"), ("m.room.join_rules", "", "$j")]) + ); + assert!(!delta.no_longer_in_room); + } + + fn joined(reads: &FakeReads, events: &[NewEvent], delta: &DeltaState) -> bool { + block_on(is_server_still_joined(reads, ROOM, SERVER, events, delta)).unwrap() + } + + #[test] + fn still_joined_without_local_membership_changes() { + // Nobody local is joined, but no local membership changed. + let reads = FakeReads::default(); + let delta = DeltaState { + to_delete: vec![StateKey::new(MEMBER, "@x:remote")], + to_insert: state(&[(MEMBER, "@y:remote", "$y"), ("m.room.topic", "", "$t")]), + ..Default::default() + }; + assert!(joined(&reads, &[], &delta)); + assert!(reads.ignored.lock().unwrap().is_none()); + } + + #[test] + fn still_joined_through_a_local_join() { + let reads = FakeReads::default(); + let delta = DeltaState { + to_insert: state(&[(MEMBER, "@me:here", "$join")]), + ..Default::default() + }; + // In the chunk. + assert!(joined( + &reads, + &[member("$join", "@me:here", "join", 1)], + &delta + )); + // A leave in the chunk is read from the database too, where it is + // not found. + assert!(!joined( + &reads, + &[member("$join", "@me:here", "leave", 1)], + &delta + )); + + // In the database. + let reads = FakeReads { + memberships: [(Arc::from("$join"), Arc::from("join"))].into(), + ..Default::default() + }; + assert!(joined(&reads, &[], &delta)); + } + + #[test] + fn still_joined_through_another_local_user() { + let delta = DeltaState { + to_delete: vec![StateKey::new(MEMBER, "@gone:here")], + to_insert: state(&[(MEMBER, "@me:here", "$leave"), (MEMBER, "@x:remote", "$x")]), + ..Default::default() + }; + let reads = FakeReads { + memberships: [(Arc::from("$leave"), Arc::from("leave"))].into(), + local_joined: vec![Arc::from("@me:here"), Arc::from("@gone:here")], + ..Default::default() + }; + // The users whose membership changed do not count. + assert!(!joined(&reads, &[], &delta)); + let mut ignored = reads.ignored.lock().unwrap().clone().unwrap(); + ignored.sort(); + assert_eq!( + ignored, + vec![Arc::from("@gone:here"), Arc::from("@me:here")] + ); + + let reads = FakeReads { + local_joined: vec![Arc::from("@other:here")], + ..reads + }; + assert!(joined(&reads, &[], &delta)); + } + + fn orchestrate( + reads: &FakeReads, + resolver: &FakeResolver, + events: &[NewEvent], + ) -> (StateAndExtremities, PersistCounts) { + let mut counts = PersistCounts::default(); + let result = block_on(calculate_new_forward_extremities_and_state_delta( + reads, + resolver, + ROOM, + SERVER, + events, + &EventMap::new(), + 0, + &mut counts, + )) + .unwrap(); + (result, counts) + } + + #[test] + fn unchanged_extremities_change_nothing() { + let reads = FakeReads { + latest: ids(&["$old"]), + ..Default::default() + }; + let mut outlier = event("$out", &["$old"], 1); + outlier.outlier = true; + let (result, counts) = orchestrate(&reads, &FakeResolver::default(), &[outlier]); + assert!(result.new_forward_extremities.is_none()); + assert!(result.state_delta.is_none()); + assert_eq!(counts, PersistCounts::default()); + } + + #[test] + fn a_chain_of_messages_needs_no_state() { + let reads = FakeReads { + latest: ids(&["$old"]), + ..Default::default() + }; + let events = [event("$1", &["$old"], 1), event("$2", &["$1"], 1)]; + let (result, counts) = orchestrate(&reads, &FakeResolver::default(), &events); + assert_eq!(result.new_forward_extremities, Some(ids(&["$2"]))); + assert!(result.state_delta.is_none()); + assert_eq!(counts.state_delta, 0); + } + + #[test] + fn a_state_change_is_counted_and_applied() { + let reads = FakeReads { + latest: ids(&["$old"]), + event_groups: [(Arc::from("$old"), 5)].into(), + local_joined: vec![Arc::from("@me:here")], + ..Default::default() + }; + let mut topic = event("$t", &["$old"], 6); + topic.event_type = Arc::from("m.room.topic"); + topic.state_key = Some(Arc::from("")); + topic.state_group_deltas = + Arc::new([((5, 6), state(&[("m.room.topic", "", "$t")]))].into()); + let (result, counts) = orchestrate(&reads, &FakeResolver::default(), &[topic]); + assert_eq!(result.new_forward_extremities, Some(ids(&["$t"]))); + assert_eq!( + result.state_delta, + Some(DeltaState { + to_insert: state(&[("m.room.topic", "", "$t")]), + ..Default::default() + }) + ); + assert_eq!(counts.state_delta, 1); + assert_eq!(counts.state_delta_single_event, 1); + assert_eq!(counts.state_delta_reuse_delta, 1); + } + + #[test] + fn reuse_needs_prevs_equal_to_the_extremities() { + // Two extremities, merged by an event that only names one. + let reads = FakeReads { + latest: ids(&["$o1", "$o2"]), + event_groups: [(Arc::from("$o1"), 5), (Arc::from("$o2"), 5)].into(), + ..Default::default() + }; + let mut events = vec![event("$m", &["$o1"], 5), event("$n", &["$m", "$o2"], 5)]; + let (_, counts) = orchestrate(&reads, &FakeResolver::default(), &events); + assert_eq!(counts.state_delta, 1); + assert_eq!(counts.state_delta_single_event, 1); + assert_eq!(counts.state_delta_reuse_delta, 0); + + events[0].prev_event_ids = vec![Arc::from("$o2"), Arc::from("$o1")]; + let (_, counts) = orchestrate(&reads, &FakeResolver::default(), &events); + assert_eq!(counts.state_delta_reuse_delta, 1); + } + + #[test] + fn the_last_local_leave_leaves_the_room() { + let reads = FakeReads { + latest: ids(&["$old"]), + event_groups: [(Arc::from("$old"), 5)].into(), + groups: [(6, Arc::new(state(&[(MEMBER, "@me:here", "$leave")])))].into(), + current_state: state(&[(MEMBER, "@me:here", "$join")]), + local_joined: vec![Arc::from("@me:here")], + ..Default::default() + }; + let leave = member("$leave", "@me:here", "leave", 6); + let leave = NewEvent { + prev_event_ids: vec![Arc::from("$old")], + ..leave + }; + let (result, _) = orchestrate(&reads, &FakeResolver::default(), &[leave]); + let delta = result.state_delta.unwrap(); + assert_eq!(delta.to_insert, state(&[(MEMBER, "@me:here", "$leave")])); + assert!(delta.no_longer_in_room); + } + + fn state_dag_event(id: &str, prev_state: &[&str], group: i64) -> NewEvent { + NewEvent { + event_type: Arc::from("m.room.topic"), + state_key: Some(Arc::from("")), + prev_state_events: Some(prev_state.iter().map(|p| Arc::from(*p)).collect()), + ..event(id, prev_state, group) + } + } + + fn orchestrate_state_dag( + reads: &FakeReads, + resolver: &FakeResolver, + events: &[NewEvent], + ) -> (StateAndExtremities, PersistCounts) { + let mut counts = PersistCounts::default(); + let result = block_on(process_state_dag_forward_extremities_and_state_delta( + reads, + resolver, + ROOM, + SERVER, + events, + &EventMap::new(), + 0, + &mut counts, + )) + .unwrap(); + (result, counts) + } + + #[test] + fn state_dag_messages_only_move_the_room_dag() { + let reads = FakeReads { + latest: ids(&["$old"]), + state_dag: ids(&["$s"]), + ..Default::default() + }; + let (result, counts) = orchestrate_state_dag( + &reads, + &FakeResolver::default(), + &[event("$m", &["$old"], 1)], + ); + assert_eq!(result.new_forward_extremities, Some(ids(&["$m"]))); + assert_eq!(result.new_state_dag_extremities, Some(ids(&["$s"]))); + assert!(result.state_delta.is_none()); + assert_eq!(counts.state_delta, 0); + } + + #[test] + fn state_dag_state_follows_the_state_dag() { + let reads = FakeReads { + latest: ids(&["$old"]), + state_dag: ids(&["$s"]), + event_groups: [(Arc::from("$s"), 5)].into(), + groups: [(6, Arc::new(state(&[("m.room.topic", "", "$t")])))].into(), + current_state: state(&[("m.room.topic", "", "$s")]), + ..Default::default() + }; + let topic = NewEvent { + prev_event_ids: vec![Arc::from("$old")], + ..state_dag_event("$t", &["$s"], 6) + }; + let (result, counts) = orchestrate_state_dag(&reads, &FakeResolver::default(), &[topic]); + assert_eq!(result.new_forward_extremities, Some(ids(&["$t"]))); + assert_eq!(result.new_state_dag_extremities, Some(ids(&["$t"]))); + assert_eq!( + result.state_delta.unwrap().to_insert, + state(&[("m.room.topic", "", "$t")]) + ); + // The state DAG path does not count the state delta metrics. + assert_eq!(counts.state_delta, 0); + assert_eq!(counts.state_dag_forward_extremities[&1], 1); + } +} diff --git a/rust/src/state/mod.rs b/rust/src/state/mod.rs index bb242caadb..87b0fc7ca3 100644 --- a/rust/src/state/mod.rs +++ b/rust/src/state/mod.rs @@ -47,7 +47,7 @@ use pyo3::{ }; use crate::caches::{python::PyCacheKey, Cache}; -use crate::database::postgres::pool::PyConnectionPool; +use crate::database::postgres::pool::{ConnectionPool, PyConnectionPool}; use crate::deferred::create_deferred; use crate::event_auth::AuthFailure; use crate::events::Event; @@ -351,9 +351,9 @@ impl StateResStats { } /// What a resolution's deferred resolves to: the state and the stats. -struct ResolveOutput { - state: StateMap, - stats: StateResStats, +pub(crate) struct ResolveOutput { + pub(crate) state: StateMap, + pub(crate) stats: StateResStats, } impl<'py> IntoPyObject<'py> for ResolveOutput { @@ -381,15 +381,9 @@ enum StateGroupsResult { State(StateMap), } -/// Port of `_make_state_cache_entry`: how `new_state` relates to the input -/// groups, in the order they were given, as the Python breaks ties by its -/// dict's order. -fn make_state_cache_entry( - new_state: StateMap, - groups: &[(i64, Arc)], -) -> StateGroupsResult { - // First look for exact matches, by the set of event IDs as the Python - // does. +/// The first of the input groups whose state is `new_state`, compared by +/// the set of event IDs as `_make_state_cache_entry` does. +pub(crate) fn matching_group(new_state: &StateMap, groups: &[(i64, Arc)]) -> Option { let new_event_ids: HashSet<&Arc> = new_state.values().collect(); for (group, state) in groups { if new_event_ids.len() != state.len() { @@ -397,9 +391,22 @@ fn make_state_cache_entry( } let old_event_ids: HashSet<&Arc> = state.values().collect(); if new_event_ids == old_event_ids { - return StateGroupsResult::Group(*group); + return Some(*group); } } + None +} + +/// Port of `_make_state_cache_entry`: how `new_state` relates to the input +/// groups, in the order they were given, as the Python breaks ties by its +/// dict's order. +fn make_state_cache_entry( + new_state: StateMap, + groups: &[(i64, Arc)], +) -> StateGroupsResult { + if let Some(group) = matching_group(&new_state, groups) { + return StateGroupsResult::Group(group); + } // Failing that, the closest match. A delta cannot remove keys, so a // group with a key the new state lacks is no candidate. @@ -524,6 +531,11 @@ impl StateResolutionHandler { Arc::clone(&self.store.caches) } + /// The conflict cache, which the Rust event persister shares. + pub(crate) fn conflict_cache(&self) -> Arc { + Arc::clone(&self.cache) + } + /// Convert the state sets under the GIL and start a resolution: either /// the answer, if nothing conflicts, or the future that computes it. fn start( @@ -653,67 +665,95 @@ impl StateResolutionHandler { let caches = self.caches(); Ok(Started::Pending(Box::pin(async move { - let start = Instant::now(); - let store = PostgresStateResStore::begin(&pool, handle, Some(caches)) - .await - .map_err(resolve_error_to_py)?; - let mut connection_time = start.elapsed(); - let prepared = v2::prepare( - &store, + resolve_on_pool( + &pool, + handle, + caches, + &cache, &room_id, room_version, &state_sets, (unconflicted, conflicted), events, - Some(&cache), + stats, ) .await - .map_err(resolve_error_to_py)?; - - // The rest does no I/O, so give the connection back first. - let finish_start = Instant::now(); - let db = store.finish().await.map_err(resolve_error_to_py)?; - connection_time += finish_start.elapsed(); - let resolution = match prepared { - v2::Prepared::Done(resolution) => resolution, - v2::Prepared::Pending(pending) => pending - .run(Some(&cache)) - .await - .map_err(resolve_error_to_py)?, - }; - - let core = &resolution.stats; - let stats = StateResStats { - total_time: start.elapsed().as_secs_f64(), - convert_in_time: stats.convert_in_time, - state_map_cache_hits: stats.state_map_cache_hits, - db_time: db.db_time.as_secs_f64(), - parse_time: db.parse_time.as_secs_f64(), - cpu_time: (core.sort_time + core.auth_time + core.mainline_time).as_secs_f64(), - sort_time: core.sort_time.as_secs_f64(), - auth_time: core.auth_time.as_secs_f64(), - mainline_time: core.mainline_time.as_secs_f64(), - auth_chain_time: core.auth_chain_time.as_secs_f64(), - load_events_time: core.load_events_time.as_secs_f64(), - connection_time: connection_time.as_secs_f64(), - db_queries: db.queries, - db_rows: db.rows, - db_events: db.events_fetched, - event_cache_hits: db.events_cached, - chain_position_cache_hits: db.chain_positions_cached, - chain_link_cache_hits: db.chain_links_cached, - convert_out_time: 0.0, - conflicted_events: core.conflicted_events, - cache_hit: core.cache_hit, - }; - Ok(ResolveOutput { - state: resolution.merged(), - stats, - }) + .map_err(resolve_error_to_py) }))) } } +/// Resolve converted state sets with conflicts over the native pool: the +/// reads in one read-only transaction on a pooled connection, which goes +/// back to the pool before the algorithm runs. `stats` holds what the +/// caller recorded before. For callers in Rust, and behind +/// [`StateResolutionHandler`]'s deferreds. +#[allow(clippy::too_many_arguments)] +pub(crate) async fn resolve_on_pool( + pool: &ConnectionPool, + handle: tokio::runtime::Handle, + caches: Arc, + cache: &ConflictCache, + room_id: &str, + room_version: &'static RoomVersion, + state_sets: &[Arc], + separated: (StateMap, v2::Conflicted), + events: EventMap, + stats: StateResStats, +) -> anyhow::Result { + let start = Instant::now(); + let store = PostgresStateResStore::begin(pool, handle, Some(caches)).await?; + let mut connection_time = start.elapsed(); + let prepared = v2::prepare( + &store, + room_id, + room_version, + state_sets, + separated, + events, + Some(cache), + ) + .await?; + + // The rest does no I/O, so give the connection back first. + let finish_start = Instant::now(); + let db = store.finish().await?; + connection_time += finish_start.elapsed(); + let resolution = match prepared { + v2::Prepared::Done(resolution) => resolution, + v2::Prepared::Pending(pending) => pending.run(Some(cache)).await?, + }; + + let core = &resolution.stats; + let stats = StateResStats { + total_time: start.elapsed().as_secs_f64(), + convert_in_time: stats.convert_in_time, + state_map_cache_hits: stats.state_map_cache_hits, + db_time: db.db_time.as_secs_f64(), + parse_time: db.parse_time.as_secs_f64(), + cpu_time: (core.sort_time + core.auth_time + core.mainline_time).as_secs_f64(), + sort_time: core.sort_time.as_secs_f64(), + auth_time: core.auth_time.as_secs_f64(), + mainline_time: core.mainline_time.as_secs_f64(), + auth_chain_time: core.auth_chain_time.as_secs_f64(), + load_events_time: core.load_events_time.as_secs_f64(), + connection_time: connection_time.as_secs_f64(), + db_queries: db.queries, + db_rows: db.rows, + db_events: db.events_fetched, + event_cache_hits: db.events_cached, + chain_position_cache_hits: db.chain_positions_cached, + chain_link_cache_hits: db.chain_links_cached, + convert_out_time: 0.0, + conflicted_events: core.conflicted_events, + cache_hit: core.cache_hit, + }; + Ok(ResolveOutput { + state: resolution.merged(), + stats, + }) +} + /// State sets converted from Python, with the interner that built them (so /// the event map shares its strings) and what the conversion recorded. struct Converted { diff --git a/synapse/synapse_rust/persist.pyi b/synapse/synapse_rust/persist.pyi index 522c3c9980..22ac643262 100644 --- a/synapse/synapse_rust/persist.pyi +++ b/synapse/synapse_rust/persist.pyi @@ -169,6 +169,61 @@ class PersistReaderForTest: events_and_contexts: list[tuple[EventBase, EventContext]], now_ms: int, ) -> tuple[set[str], dict[str, Any]]: ... + def get_new_state_after_events( + self, + room_id: str, + server_name: str, + events_and_contexts: list[tuple[EventBase, EventContext]], + old_latest_event_ids: list[str], + new_latest_event_ids: list[str], + should_prune: bool, + now_ms: int, + ) -> tuple[ + dict[tuple[str, str], str] | None, + dict[tuple[str, str], str] | None, + set[str], + dict[str, Any], + ]: + """`(current_state, delta_ids, new_latest_event_ids, counts)`.""" + def calculate_state_delta( + self, room_id: str, current_state: dict[tuple[str, str], str] + ) -> tuple[list[tuple[str, str]], dict[tuple[str, str], str]]: + """`(to_delete, to_insert)`.""" + def is_server_still_joined( + self, + room_id: str, + server_name: str, + events_and_contexts: list[tuple[EventBase, EventContext]], + to_delete: list[tuple[str, str]], + to_insert: dict[tuple[str, str], str], + ) -> bool: ... + def calculate_new_forward_extremities_and_state_delta( + self, + room_id: str, + server_name: str, + events_and_contexts: list[tuple[EventBase, EventContext]], + now_ms: int, + ) -> tuple[ + set[str] | None, + tuple[list[tuple[str, str]], dict[tuple[str, str], str], bool] | None, + dict[str, Any], + ]: + """`(new_forward_extremities, delta, counts)`, the delta as + `(to_delete, to_insert, no_longer_in_room)`.""" + def process_state_dag_forward_extremities_and_state_delta( + self, + room_id: str, + server_name: str, + events_and_contexts: list[tuple[EventBase, EventContext]], + now_ms: int, + ) -> tuple[ + set[str] | None, + tuple[list[tuple[str, str]], dict[tuple[str, str], str], bool] | None, + set[str] | None, + dict[str, Any], + ]: + """`(new_forward_extremities, delta, new_state_dag_extremities, + counts)`.""" def persist_gil_stats() -> dict[str, tuple[int, float]]: """`{site: (calls, seconds held)}` for every time event persistence took diff --git a/tests/storage/test_persist_reader.py b/tests/storage/test_persist_reader.py index 3c9ff42af8..7d7d64eb0c 100644 --- a/tests/storage/test_persist_reader.py +++ b/tests/storage/test_persist_reader.py @@ -25,7 +25,7 @@ from unittest import SkipTest from twisted.internet.testing import MemoryReactor -from synapse.api.constants import EventTypes +from synapse.api.constants import EventTypes, Membership from synapse.api.errors import NotFoundError, SynapseError from synapse.api.room_versions import RoomVersions from synapse.events import EventBase @@ -35,8 +35,10 @@ from synapse.rest.client import login, room from synapse.server import HomeServer from synapse.storage.database import LoggingTransaction from synapse.storage.databases.main.event_federation import EventFederationStore +from synapse.storage.databases.main.events import DeltaState from synapse.storage.rust_pool import RustConnectionPool from synapse.synapse_rust.persist import PersistReaderForTest +from synapse.types import StateMap, create_requester from synapse.types.handlers import SLIDING_SYNC_DEFAULT_BUMP_EVENT_TYPES from synapse.util.clock import Clock @@ -955,3 +957,440 @@ class ExtremitiesTestCase(PersistReaderTestBase): ), {keep, being_persisted[0].event_id}, ) + + +_Delta = tuple[set[tuple[str, str]], dict[tuple[str, str], str], bool] + + +def _python_delta(delta: DeltaState | None) -> _Delta | None: + if delta is None: + return None + return (set(delta.to_delete), dict(delta.to_insert), delta.no_longer_in_room) + + +def _rust_delta( + delta: tuple[list[tuple[str, str]], dict[tuple[str, str], str], bool] | None, +) -> _Delta | None: + if delta is None: + return None + to_delete, to_insert, no_longer_in_room = delta + return (set(to_delete), to_insert, no_longer_in_room) + + +class StateDeltaTestCase(PersistReaderTestBase): + """The state delta in `rust/src/persist/state_delta.rs` against the + controller's Python methods, over the same database and the same new + events, which are built with real contexts but not persisted.""" + + def default_config(self) -> dict[str, Any]: + config = super().default_config() + config["experimental_features"]["msc4242_enabled"] = True + return config + + def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None: + super().prepare(reactor, clock, hs) + persistence = hs.get_storage_controllers().persistence + assert persistence is not None + self.persistence = persistence + + def _create( + self, + event_type: str, + content: dict[str, Any], + *, + state_key: str | None = None, + sender: str | None = None, + prev_event_ids: list[str] | None = None, + room_id: str | None = None, + ) -> tuple[EventBase, EventContext]: + """An event as the event creation handler builds it, with its state + group stored but the event not persisted.""" + event_dict: dict[str, Any] = { + "type": event_type, + "room_id": room_id or self.room_id, + "sender": sender or self.alice, + "content": content, + } + if state_key is not None: + event_dict["state_key"] = state_key + event, unpersisted = self.get_success( + self.hs.get_event_creation_handler().create_event( + create_requester(sender or self.alice), + event_dict, + prev_event_ids=prev_event_ids, + ) + ) + context = self.get_success(unpersisted.persist(event)) + return event, context + + def _topic( + self, topic: str, prev_event_ids: list[str] | None = None + ) -> tuple[EventBase, EventContext]: + return self._create( + EventTypes.Topic, + {"topic": topic}, + state_key="", + prev_event_ids=prev_event_ids, + ) + + def _leave( + self, user: str, prev_event_ids: list[str] | None = None + ) -> tuple[EventBase, EventContext]: + return self._create( + EventTypes.Member, + {"membership": Membership.LEAVE}, + state_key=user, + sender=user, + prev_event_ids=prev_event_ids, + ) + + def _persist(self, event_and_context: tuple[EventBase, EventContext]) -> str: + event, context = event_and_context + self.get_success(self.persistence.persist_event(event, context)) + return event.event_id + + def _latest(self) -> set[str]: + return set( + self.get_success(self.store.get_latest_event_ids_in_room(self.room_id)) + ) + + def _new_state( + self, + events: list[tuple[EventBase, EventContext]], + old: set[str], + new: set[str], + ) -> tuple[StateMap[str] | None, StateMap[str] | None, set[str], dict[str, Any]]: + current_state, delta_ids, new_latest = self.get_success( + self.persistence._get_new_state_after_events( + self.room_id, events, old, set(new) + ) + ) + rust_state, rust_delta, rust_latest, counts = ( + self.rust.get_new_state_after_events( + self.room_id, + self.hs.hostname, + events, + list(old), + list(new), + True, + self.clock.time_msec(), + ) + ) + self.assertEqual( + rust_state, None if current_state is None else dict(current_state) + ) + self.assertEqual(rust_delta, None if delta_ids is None else dict(delta_ids)) + self.assertEqual(rust_latest, new_latest) + return rust_state, rust_delta, rust_latest, counts + + def test_get_new_state_after_events(self) -> None: + latest = self._latest() + + # A message: the state groups do not change. + message = self._create(EventTypes.Message, {"msgtype": "m.text", "body": "x"}) + current_state, delta, _, _ = self._new_state( + [message], latest, {message[0].event_id} + ) + self.assertIsNone(current_state) + self.assertIsNone(delta) + + # A state event: its context knows the delta. + topic = self._topic("one") + current_state, delta, _, _ = self._new_state( + [topic], latest, {topic[0].event_id} + ) + self.assertIsNone(current_state) + self.assertEqual(delta, {(EventTypes.Topic, ""): topic[0].event_id}) + + # Without the delta, the new group's state. + topic[1].state_group_deltas = {} + current_state, delta, _, counts = self._new_state( + [topic], latest, {topic[0].event_id} + ) + assert current_state is not None + self.assertIsNone(delta) + self.assertEqual(current_state[(EventTypes.Topic, "")], topic[0].event_id) + self.assertEqual(counts["state_resolutions"], 0) + + def test_resolution_of_a_fork(self) -> None: + (parent,) = self._latest() + persisted = self._persist(self._topic("persisted", [parent])) + fork = self._topic("fork", [parent]) + current_state, delta, latest, counts = self._new_state( + [fork], {persisted}, {persisted, fork[0].event_id} + ) + assert current_state is not None + self.assertIsNone(delta) + self.assertIn( + current_state[(EventTypes.Topic, "")], {persisted, fork[0].event_id} + ) + self.assertEqual(latest, {persisted, fork[0].event_id}) + self.assertEqual(counts["state_resolutions"], 1) + + def _pending_deletions(self, groups: list[int]) -> dict[int, tuple[int, int]]: + rows = self.get_success( + self.hs.get_datastores().state_deletion.db_pool.simple_select_many_batch( + table="state_groups_pending_deletion", + column="state_group", + iterable=groups, + retcols=("state_group", "sequence_number", "insertion_ts"), + ) + ) + return {group: (sequence, ts) for group, sequence, ts in rows} + + def test_resolution_bumps_pending_deletions(self) -> None: + (parent,) = self._latest() + persisted_topic = self._topic("persisted", [parent]) + persisted = self._persist(persisted_topic) + fork = self._topic("fork", [parent]) + persisted_group = persisted_topic[1].state_group + fork_group = fork[1].state_group + assert persisted_group is not None and fork_group is not None + groups = [persisted_group, fork_group] + + self.get_success( + self.hs.get_datastores().state_deletion.mark_state_groups_as_pending_deletion( + groups + ) + ) + before = self._pending_deletions(groups) + self.assertEqual(before.keys(), set(groups)) + + # Only Rust runs, so the bump is its own. + self.reactor.advance(60) + self.rust.get_new_state_after_events( + self.room_id, + self.hs.hostname, + [fork], + [persisted], + [persisted, fork[0].event_id], + True, + self.clock.time_msec(), + ) + after = self._pending_deletions(groups) + for group in groups: + self.assertGreater(after[group][0], before[group][0]) + self.assertEqual(after[group][1], self.clock.time_msec()) + + def test_resolution_of_a_deleted_group(self) -> None: + (parent,) = self._latest() + persisted = self._persist(self._topic("persisted", [parent])) + fork = self._topic("fork", [parent]) + fork[1]._state_group = 1_000_000_000 + new = {persisted, fork[0].event_id} + + failure = self.get_failure( + self.persistence._get_new_state_after_events( + self.room_id, [fork], {persisted}, set(new) + ), + Exception, + ) + self.assertIn("deleted", str(failure.value)) + with self.assertRaisesRegex(Exception, "State groups have been deleted"): + self.rust.get_new_state_after_events( + self.room_id, + self.hs.hostname, + [fork], + [persisted], + list(new), + True, + self.clock.time_msec(), + ) + + def test_calculate_state_delta(self) -> None: + current = dict( + self.get_success(self.store.get_partial_current_state_ids(self.room_id)) + ) + new_state = dict(current) + del new_state[(EventTypes.JoinRules, "")] + new_state[(EventTypes.Name, "")] = "$changed" + new_state[(EventTypes.Topic, "")] = "$added" + + expected = self.get_success( + self.persistence._calculate_state_delta(self.room_id, new_state) + ) + to_delete, to_insert = self.rust.calculate_state_delta(self.room_id, new_state) + self.assertEqual(set(to_delete), set(expected.to_delete)) + self.assertEqual(to_insert, expected.to_insert) + self.assertEqual(set(to_delete), {(EventTypes.JoinRules, "")}) + self.assertEqual( + to_insert, + {(EventTypes.Name, ""): "$changed", (EventTypes.Topic, ""): "$added"}, + ) + + def _still_joined( + self, + events: list[tuple[EventBase, EventContext]], + to_delete: list[tuple[str, str]], + to_insert: dict[tuple[str, str], str], + ) -> bool: + expected = self.get_success( + self.persistence._is_server_still_joined( + self.room_id, events, DeltaState(to_delete, to_insert) + ) + ) + result = self.rust.is_server_still_joined( + self.room_id, self.hs.hostname, events, to_delete, to_insert + ) + self.assertEqual(result, expected) + return result + + def test_is_server_still_joined(self) -> None: + alice_key = (EventTypes.Member, self.alice) + bob_key = (EventTypes.Member, self.bob) + alice_leave = self._leave(self.alice) + bob_leave = self._leave(self.bob) + + # No local membership changes. + self.assertTrue( + self._still_joined( + [], [(EventTypes.Member, "@x:remote")], {(EventTypes.Topic, ""): "$t"} + ) + ) + # Alice leaves, Bob is still joined. + self.assertTrue( + self._still_joined([alice_leave], [], {alice_key: alice_leave[0].event_id}) + ) + # Both leave, or one leaves and the other's membership goes. + self.assertFalse( + self._still_joined( + [alice_leave, bob_leave], + [], + { + alice_key: alice_leave[0].event_id, + bob_key: bob_leave[0].event_id, + }, + ) + ) + self.assertFalse( + self._still_joined( + [alice_leave], [bob_key], {alice_key: alice_leave[0].event_id} + ) + ) + # A join in the database, which is Alice's current membership. + alice_join = self.get_success( + self.store.get_partial_current_state_ids(self.room_id) + )[alice_key] + self.assertTrue( + self._still_joined([alice_leave], [bob_key], {alice_key: alice_join}) + ) + + def test_still_joined_by_a_join_in_the_chunk(self) -> None: + alice_key = (EventTypes.Member, self.alice) + bob_key = (EventTypes.Member, self.bob) + # Alice changes her display name, a join. Bob's membership goes, so + # only the join in the chunk keeps the server in the room. + rejoin = self._create( + EventTypes.Member, + {"membership": Membership.JOIN, "displayname": "Alice"}, + state_key=self.alice, + ) + self.assertTrue( + self._still_joined([rejoin], [bob_key], {alice_key: rejoin[0].event_id}) + ) + # Without the event in the chunk, the server reads the database, + # where the join is not stored. + self.assertFalse( + self._still_joined([], [bob_key], {alice_key: rejoin[0].event_id}) + ) + + def _orchestrate( + self, events: list[tuple[EventBase, EventContext]] + ) -> tuple[set[str] | None, _Delta | None, dict[str, Any]]: + forward, delta = self.get_success( + self.persistence._calculate_new_forward_extremities_and_state_delta( + self.room_id, events + ) + ) + rust_forward, rust_delta, counts = ( + self.rust.calculate_new_forward_extremities_and_state_delta( + self.room_id, self.hs.hostname, events, self.clock.time_msec() + ) + ) + self.assertEqual(rust_forward, forward) + self.assertEqual(_rust_delta(rust_delta), _python_delta(delta)) + return rust_forward, _rust_delta(rust_delta), counts + + def test_calculate_new_forward_extremities_and_state_delta(self) -> None: + # A message only moves the extremity. + message = self._create(EventTypes.Message, {"msgtype": "m.text", "body": "x"}) + forward, delta, counts = self._orchestrate([message]) + self.assertEqual(forward, {message[0].event_id}) + self.assertIsNone(delta) + self.assertEqual(counts["state_delta"], 0) + + # A state event, whose delta the context knows. + topic = self._topic("one") + forward, delta, counts = self._orchestrate([topic]) + self.assertEqual( + delta, (set(), {(EventTypes.Topic, ""): topic[0].event_id}, False) + ) + self.assertEqual(counts["state_delta"], 1) + self.assertEqual(counts["state_delta_single_event"], 1) + self.assertEqual(counts["state_delta_reuse_delta"], 1) + + # The last local user leaves, so the server leaves the room. + self._persist(topic) + self._persist(self._leave(self.alice)) + bob_leave = self._leave(self.bob) + forward, delta, counts = self._orchestrate([bob_leave]) + self.assertEqual( + delta, + (set(), {(EventTypes.Member, self.bob): bob_leave[0].event_id}, True), + ) + + def test_orchestration_resolves_a_fork(self) -> None: + (parent,) = self._latest() + persisted = self._persist(self._topic("persisted", [parent])) + fork = self._topic("fork", [parent]) + forward, delta, counts = self._orchestrate([fork]) + self.assertEqual(forward, {persisted, fork[0].event_id}) + assert delta is not None + self.assertEqual(counts["state_resolutions"], 1) + self.assertEqual(counts["state_delta"], 1) + self.assertEqual(counts["state_delta_single_event"], 0) + + def test_state_dag_orchestration(self) -> None: + room_id = self.helper.create_room_as( + self.alice, + room_version=RoomVersions.MSC4242v12.identifier, + tok=self.alice_tok, + ) + + def compare( + events: list[tuple[EventBase, EventContext]], + ) -> tuple[set[str] | None, _Delta | None, set[str] | None, dict[str, Any]]: + forward, delta, state_dag = self.get_success( + self.persistence._process_state_dag_forward_extremities_and_state_delta( + room_id, + events, # type: ignore[arg-type] + ) + ) + rust_forward, rust_delta, rust_state_dag, counts = ( + self.rust.process_state_dag_forward_extremities_and_state_delta( + room_id, self.hs.hostname, events, self.clock.time_msec() + ) + ) + self.assertEqual(rust_forward, forward) + self.assertEqual(_rust_delta(rust_delta), _python_delta(delta)) + self.assertEqual(rust_state_dag, state_dag) + return rust_forward, _rust_delta(rust_delta), rust_state_dag, counts + + # A message moves only the room DAG. + message = self._create( + EventTypes.Message, {"msgtype": "m.text", "body": "x"}, room_id=room_id + ) + forward, delta, state_dag, _ = compare([message]) + self.assertEqual(forward, {message[0].event_id}) + self.assertIsNone(delta) + + # A state event moves both, and the state. + name = self._create( + EventTypes.Name, {"name": "dag"}, state_key="", room_id=room_id + ) + forward, delta, state_dag, counts = compare([name]) + self.assertEqual(state_dag, {name[0].event_id}) + self.assertEqual( + delta, (set(), {(EventTypes.Name, ""): name[0].event_id}, False) + ) + self.assertEqual(counts["state_dag_forward_extremities"], {1: 1})