From b5988fdc64f89a9b6b05c694eba89674bcedbc04 Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Tue, 13 Oct 2026 15:32:24 +0000 Subject: [PATCH] Simplify the state res types - `StateResEvent` wraps a plain `Event`. An event handed in from Python is a shallow copy sharing the Python object's data (its internal metadata included), so the `EventSource::{Python, Rust}` enum and the separate constructors go. - `CpuStats` folds into `ResolveStats`, and the store's five atomic counters become one `Mutex`. - `prepare` takes the state already split by `separate`, as the handler needs; the three-line `prepare` wrapper that called `separate` itself is gone and `prepare_separated` takes its name. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01ULdieJyQoHNzeHq42H2M1p --- rust/src/state/mod.rs | 46 ++++++---------------- rust/src/state/store.rs | 42 ++++++-------------- rust/src/state/v2.rs | 79 +++++++++++--------------------------- rust/src/state/v2/tests.rs | 5 +-- 4 files changed, 49 insertions(+), 123 deletions(-) diff --git a/rust/src/state/mod.rs b/rust/src/state/mod.rs index c4509518cd..62333dda87 100644 --- a/rust/src/state/mod.rs +++ b/rust/src/state/mod.rs @@ -136,59 +136,37 @@ pub type StateMap = HashMap>; /// An event taking part in state resolution. /// -/// Either an event handed in from Python (shared with the Python object, read -/// without the GIL as `Event` is frozen) or one loaded from the database by -/// Rust. The auth event IDs are computed once, as the algorithm reads them -/// many times. +/// An event handed in from Python is a shallow copy sharing the Python +/// object's data, including its internal metadata. The auth event IDs are +/// computed once, as the algorithm reads them many times. #[derive(Clone)] pub struct StateResEvent(Arc); struct StateResEventInner { - source: EventSource, + event: Event, auth_event_ids: Box<[Arc]>, } -enum EventSource { - Python(Py), - Rust(Event), -} - impl StateResEvent { - /// Wrap an event from Python. - pub fn from_py(event: Py) -> anyhow::Result { - let auth_event_ids = auth_event_ids_of(event.get())?; - Ok(StateResEvent(Arc::new(StateResEventInner { - source: EventSource::Python(event), - auth_event_ids, - }))) - } - - /// Wrap an event loaded by Rust. - pub fn from_rust(event: Event) -> anyhow::Result { + pub fn new(event: Event) -> anyhow::Result { let auth_event_ids = auth_event_ids_of(&event)?; Ok(StateResEvent(Arc::new(StateResEventInner { - source: EventSource::Rust(event), + event, auth_event_ids, }))) } pub fn event(&self) -> &Event { - match &self.0.source { - EventSource::Python(event) => event.get(), - EventSource::Rust(event) => event, - } + &self.0.event } pub fn auth_event_ids(&self) -> &[Arc] { &self.0.auth_event_ids } - /// The event as a Python object, creating one for an event loaded by Rust. + /// The event as a new Python object sharing this event's data. pub fn to_py(&self, py: Python<'_>) -> PyResult> { - match &self.0.source { - EventSource::Python(event) => Ok(event.clone_ref(py)), - EventSource::Rust(event) => Py::new(py, event.shallow_clone()), - } + Py::new(py, self.0.event.shallow_clone()) } } @@ -277,10 +255,10 @@ fn event_map_from_py( ) -> PyResult { let mut map = EventMap::new(); for (event_id, event) in events.into_iter().flat_map(|m| m.iter()) { - let event = event.cast_into::()?.unbind(); + let event = event.cast::()?.get().shallow_clone(); map.insert( interner.intern(event_id.extract()?), - StateResEvent::from_py(event).map_err(|e| PyValueError::new_err(format!("{e:#}")))?, + StateResEvent::new(event).map_err(|e| PyValueError::new_err(format!("{e:#}")))?, ); } Ok(map) @@ -478,7 +456,7 @@ impl StateResolutionHandler { let store = PostgresStateResStore::begin(&pool, handle) .await .map_err(resolve_error_to_py)?; - let prepared = v2::prepare_separated( + let prepared = v2::prepare( &store, &room_id, room_version, diff --git a/rust/src/state/store.rs b/rust/src/state/store.rs index a304409761..3b253521e6 100644 --- a/rust/src/state/store.rs +++ b/rust/src/state/store.rs @@ -20,10 +20,7 @@ use std::{ collections::HashSet, - sync::{ - atomic::{AtomicU64, Ordering}, - Arc, - }, + sync::{Arc, Mutex}, time::{Duration, Instant}, }; @@ -118,11 +115,7 @@ pub struct DbStats { pub struct PostgresStateResStore { conn: Option, handle: Handle, - queries: AtomicU64, - rows: AtomicU64, - events_fetched: AtomicU64, - db_nanos: AtomicU64, - parse_nanos: AtomicU64, + db_stats: Mutex, } /// How many event IDs to send per event-loading statement. @@ -148,11 +141,7 @@ impl PostgresStateResStore { let store = PostgresStateResStore { conn: Some(conn), handle, - queries: AtomicU64::new(0), - rows: AtomicU64::new(0), - events_fetched: AtomicU64::new(0), - db_nanos: AtomicU64::new(0), - parse_nanos: AtomicU64::new(0), + db_stats: Mutex::default(), }; // If this fails, dropping the store sends a ROLLBACK, which is // harmless outside a transaction. @@ -176,24 +165,19 @@ impl PostgresStateResStore { } pub fn stats(&self) -> DbStats { - DbStats { - queries: self.queries.load(Ordering::Relaxed), - rows: self.rows.load(Ordering::Relaxed), - events_fetched: self.events_fetched.load(Ordering::Relaxed), - db_time: Duration::from_nanos(self.db_nanos.load(Ordering::Relaxed)), - parse_time: Duration::from_nanos(self.parse_nanos.load(Ordering::Relaxed)), - } + *self.db_stats.lock().expect("not poisoned") } fn client(&self) -> &PooledConnection { self.conn.as_ref().expect("connection is held until finish") } + /// Record a statement started at `start` that returned `rows` rows. fn record(&self, start: Instant, rows: usize) { - self.queries.fetch_add(1, Ordering::Relaxed); - self.rows.fetch_add(rows as u64, Ordering::Relaxed); - self.db_nanos - .fetch_add(start.elapsed().as_nanos() as u64, Ordering::Relaxed); + let mut stats = self.db_stats.lock().expect("not poisoned"); + stats.queries += 1; + stats.rows += rows as u64; + stats.db_time += start.elapsed(); } /// Run a statement through the connection's prepared statement cache. @@ -242,14 +226,12 @@ impl StateResStore for PostgresStateResStore { let start = Instant::now(); for row in rows { if let Some((event_id, event)) = event_from_row(&row)? { - events.insert(event_id, StateResEvent::from_rust(event)?); + events.insert(event_id, StateResEvent::new(event)?); } } - self.parse_nanos - .fetch_add(start.elapsed().as_nanos() as u64, Ordering::Relaxed); + self.db_stats.lock().expect("not poisoned").parse_time += start.elapsed(); } - self.events_fetched - .fetch_add(events.len() as u64, Ordering::Relaxed); + self.db_stats.lock().expect("not poisoned").events_fetched += events.len() as u64; Ok(events) } diff --git a/rust/src/state/v2.rs b/rust/src/state/v2.rs index 60de1fac10..a6e55d254d 100644 --- a/rust/src/state/v2.rs +++ b/rust/src/state/v2.rs @@ -16,9 +16,9 @@ //! //! The work is split in two: //! -//! - [`prepare`] separates the conflicted state, computes the auth chain -//! difference, and loads every event the algorithm can touch (the -//! "prefetch closure", see [`prefetch`]). +//! - [`separate`] splits the state into unconflicted and conflicted keys, and +//! [`prepare`] computes the auth chain difference and loads every event the +//! algorithm can touch (the "prefetch closure", see [`prefetch`]). //! - [`PendingResolution::run`] does no I/O. It runs the power sort, the //! iterative auth checks and the mainline sort over the loaded events (see //! [`Core::resolve_conflicted_set`]). @@ -84,7 +84,8 @@ impl Resolution { } } -/// Where the time went in a resolution. +/// Where the time went in a resolution. The times are wall time in each stage +/// of the core, including time yielded to other tasks. #[derive(Debug, Default, Clone, Copy)] pub struct ResolveStats { /// The resolution of the conflicted set came from the cache. @@ -104,9 +105,9 @@ pub struct ResolveStats { /// `event_map` holds events that are not yet persisted (e.g. those being /// persisted now). Everything else is read from `store`. /// -/// This is [`prepare`] followed by [`PendingResolution::run`]. Callers that -/// hold a database connection for the store should call those separately and -/// release the connection in between, as `run` does no I/O. +/// This is [`separate`], [`prepare`] and [`PendingResolution::run`]. Callers +/// that hold a database connection for the store should call those separately +/// and release the connection in between, as `run` does no I/O. pub async fn resolve( store: &dyn StateResStore, room_id: &str, @@ -120,6 +121,7 @@ pub async fn resolve( room_id, room_version, state_sets, + separate(state_sets), event_map, conflict_cache, ) @@ -170,14 +172,10 @@ impl PendingResolution { room_version, events: &event_map, }; - let (resolved, cpu) = core - .resolve_conflicted_set(&full_conflicted_set, base_state) + let resolved = core + .resolve_conflicted_set(&full_conflicted_set, base_state, &mut stats) .await?; - stats.sort_time = cpu.sort_time; - stats.auth_time = cpu.auth_time; - stats.mainline_time = cpu.mainline_time; - let resolved = Arc::new(resolved); if let (Some(cache), Some(key)) = (conflict_cache, cache_key) { cache.insert(key, Arc::clone(&resolved)); @@ -191,32 +189,11 @@ impl PendingResolution { } } -/// Everything in a resolution that reads the database: separate the -/// conflicted state, compute the auth chain difference, check the conflict -/// cache and load the events the algorithm can read (see [`prefetch`]). +/// Everything in a resolution that reads the database: compute the auth chain +/// difference, check the conflict cache and load the events the algorithm can +/// read (see [`prefetch`]). Takes the state already split by [`separate`], so +/// a caller can answer a conflict-free resolution without a store. pub async fn prepare( - store: &dyn StateResStore, - room_id: &str, - room_version: &'static RoomVersion, - state_sets: &[StateMap], - event_map: EventMap, - conflict_cache: Option<&ConflictCache>, -) -> anyhow::Result { - let separated = separate(state_sets); - prepare_separated( - store, - room_id, - room_version, - state_sets, - separated, - event_map, - conflict_cache, - ) - .await -} - -/// [`prepare`], with the state already split by [`separate`]. -pub async fn prepare_separated( store: &dyn StateResStore, room_id: &str, room_version: &'static RoomVersion, @@ -743,15 +720,6 @@ impl Yielder { } } -/// Time spent in each stage of the core, including time yielded to other -/// tasks. -#[derive(Debug, Default, Clone, Copy)] -pub struct CpuStats { - pub sort_time: Duration, - pub auth_time: Duration, - pub mainline_time: Duration, -} - /// The part of the algorithm that runs over a fully loaded event map. pub struct Core<'a> { pub room_id: &'a str, @@ -787,14 +755,13 @@ impl<'a> Core<'a> { /// `_resolve_conflicted_set`. /// /// Returns the base state with the conflicted events that pass auth - /// applied over it. + /// applied over it, recording the time each stage took in `stats`. pub async fn resolve_conflicted_set( &self, full_conflicted_set: &HashSet>, base_state: StateMap, - ) -> anyhow::Result<(StateMap, CpuStats)> { - let mut cpu = CpuStats::default(); - + stats: &mut ResolveStats, + ) -> anyhow::Result { // Sort the power events (kicks, bans, etc). let start = Instant::now(); let mut yielder = Yielder::default(); @@ -808,7 +775,7 @@ impl<'a> Core<'a> { let sorted_power_events = self .reverse_topological_power_sort(&power_events, full_conflicted_set) .await?; - cpu.sort_time = start.elapsed(); + stats.sort_time = start.elapsed(); debug!("sorted {} power events", sorted_power_events.len()); // Auth each one in turn. @@ -816,7 +783,7 @@ impl<'a> Core<'a> { let resolved_state = self .iterative_auth_checks(&sorted_power_events, base_state) .await?; - cpu.auth_time = start.elapsed(); + stats.auth_time = start.elapsed(); // Sort the remaining events by the mainline of the resolved power // levels. @@ -830,15 +797,15 @@ impl<'a> Core<'a> { debug!("sorting {} remaining events", leftover_events.len()); let pl = lookup(&resolved_state, POWER_LEVELS_KEY).cloned(); let leftover_events = self.mainline_sort(leftover_events, pl.as_deref()).await?; - cpu.mainline_time = start.elapsed(); + stats.mainline_time = start.elapsed(); let start = Instant::now(); let resolved_state = self .iterative_auth_checks(&leftover_events, resolved_state) .await?; - cpu.auth_time += start.elapsed(); + stats.auth_time += start.elapsed(); - Ok((resolved_state, cpu)) + Ok(resolved_state) } /// Port of `_reverse_topological_power_sort`: the events and their auth diff --git a/rust/src/state/v2/tests.rs b/rust/src/state/v2/tests.rs index e0267f9845..eceee72e07 100644 --- a/rust/src/state/v2/tests.rs +++ b/rust/src/state/v2/tests.rs @@ -92,7 +92,7 @@ impl FakeEvent { if let Some(state_key) = self.state_key { event["state_key"] = json!(state_key); } - StateResEvent::from_rust(make_event(&RoomVersion::V1, event)).unwrap() + StateResEvent::new(make_event(&RoomVersion::V1, event)).unwrap() } fn auth_types(&self) -> HashSet<(&str, &str)> { @@ -821,8 +821,7 @@ fn power_level_for_sender() { if !room_version.implicit_room_creator { create_content["creator"] = json!(ALICE); } - let make = - |partial: Value| StateResEvent::from_rust(make_event(room_version, partial)).unwrap(); + let make = |partial: Value| StateResEvent::new(make_event(room_version, partial)).unwrap(); let create = make(json!({"room_id": ROOM_ID, "sender": ALICE, "type": CREATE, "state_key": "", "content": create_content}));