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<DbStats>`.
- `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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01ULdieJyQoHNzeHq42H2M1p

rust/src/state/mod.rs

@@ -136,59 +136,37 @@ pub type StateMap = HashMap<StateKey, Arc<str>>;
 
 /// 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<StateResEventInner>);
 
 struct StateResEventInner {
-    source: EventSource,
+    event: Event,
     auth_event_ids: Box<[Arc<str>]>,
 }
 
-enum EventSource {
-    Python(Py<Event>),
-    Rust(Event),
-}
-
 impl StateResEvent {
-    /// Wrap an event from Python.
-    pub fn from_py(event: Py<Event>) -> anyhow::Result<Self> {
-        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<Self> {
+    pub fn new(event: Event) -> anyhow::Result<Self> {
         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<str>] {
         &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<Py<Event>> {
-        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<EventMap> {
     let mut map = EventMap::new();
     for (event_id, event) in events.into_iter().flat_map(|m| m.iter()) {
-        let event = event.cast_into::<Event>()?.unbind();
+        let event = event.cast::<Event>()?.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,

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<PooledConnection>,
     handle: Handle,
-    queries: AtomicU64,
-    rows: AtomicU64,
-    events_fetched: AtomicU64,
-    db_nanos: AtomicU64,
-    parse_nanos: AtomicU64,
+    db_stats: Mutex<DbStats>,
 }
 
 /// 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)
     }
 

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<Prepared> {
-    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<Arc<str>>,
         base_state: StateMap,
-    ) -> anyhow::Result<(StateMap, CpuStats)> {
-        let mut cpu = CpuStats::default();
-
+        stats: &mut ResolveStats,
+    ) -> anyhow::Result<StateMap> {
         // 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

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}));