diff --git a/rust/src/persist/chain_cover.rs b/rust/src/persist/chain_cover.rs index f6d6de45b0..da7392dca3 100644 --- a/rust/src/persist/chain_cover.rs +++ b/rust/src/persist/chain_cover.rs @@ -47,6 +47,8 @@ use anyhow::{anyhow, ensure}; use async_trait::async_trait; use futures::FutureExt; +use tokio_postgres::types::ToSql; + use super::reader::{self, AuthEventChain, ChainLink, ChainToCalculate}; use crate::{ database::postgres::{ @@ -54,6 +56,7 @@ use crate::{ write::{run_rolled_back_interaction, run_write_interaction, WriteTxn}, }, events::{constants::event_type::M_ROOM_CREATE, Event}, + storage::sync_store::SELECT_MANY_BATCH_SIZE, }; /// The sequence `event_chain_id_gen` allocates from. @@ -142,6 +145,34 @@ pub trait ChainCoverStore: Send + Sync { /// Allocate `n` new chain IDs. async fn next_chain_ids(&self, n: usize) -> anyhow::Result>; + + /// [`ChainCoverStore::events_with_chain_ids`] and + /// [`ChainCoverStore::chains_to_calculate`] together, in one round trip + /// where the store can manage it. + async fn chain_cover_inputs( + &self, + room_id: &str, + event_ids: &[Arc], + ) -> anyhow::Result<(HashSet>, Vec)> { + Ok(( + self.events_with_chain_ids(event_ids).await?, + self.chains_to_calculate(room_id).await?, + )) + } + + /// [`ChainCoverStore::max_sequence_numbers`] and + /// [`ChainCoverStore::chain_links`] together, in one round trip where + /// the store can manage it. + async fn max_sequence_numbers_and_links( + &self, + chain_ids: &[i64], + chains: Vec, + ) -> anyhow::Result<(HashMap, Vec)> { + Ok(( + self.max_sequence_numbers(chain_ids).await?, + self.chain_links(chains).await?, + )) + } } #[async_trait] @@ -215,6 +246,71 @@ impl ChainCoverStore for WriteTxn { .map(|row| Ok(row.try_get::<_, i64>(0)?)) .collect() } + + /// The first batch of events with the room's rows in one round trip, + /// then the rest of the events as [`reader::events_with_chain_ids`]. + async fn chain_cover_inputs( + &self, + room_id: &str, + event_ids: &[Arc], + ) -> anyhow::Result<(HashSet>, Vec)> { + let (first, rest) = event_ids.split_at(event_ids.len().min(SELECT_MANY_BATCH_SIZE)); + let ids: Vec<&str> = first.iter().map(|id| &**id).collect(); + let rooms: Vec<&str> = vec![room_id]; + let id_params: [&(dyn ToSql + Sync); 1] = [&ids]; + let room_params: [&(dyn ToSql + Sync); 1] = [&rooms]; + let mut results = self + .query_many(&[ + (reader::EVENTS_WITH_CHAIN_IDS_SQL, &id_params), + (reader::CHAINS_TO_CALCULATE_SQL, &room_params), + ]) + .await? + .into_iter(); + let mut found = reader::first_column_set(&results.next().unwrap_or_default())?; + let to_calculate = reader::parse_chains_to_calculate(&results.next().unwrap_or_default())?; + found.extend(reader::events_with_chain_ids(self, rest).await?); + Ok((found, to_calculate)) + } + + /// The sequence numbers with the first batch of chains' links in one + /// round trip, then the chains not reached as [`reader::chain_links`]. + async fn max_sequence_numbers_and_links( + &self, + chain_ids: &[i64], + chains: Vec, + ) -> anyhow::Result<(HashMap, Vec)> { + let (first, rest) = chains.split_at(chains.len().min(reader::CHAIN_LINKS_BATCH_SIZE)); + let sequence_params: [&(dyn ToSql + Sync); 1] = [&chain_ids]; + let link_params: [&(dyn ToSql + Sync); 1] = [&first]; + let mut queries: Vec<(&str, &[&(dyn ToSql + Sync)])> = Vec::new(); + if !chain_ids.is_empty() { + queries.push((reader::MAX_SEQUENCE_NUMBERS_SQL, &sequence_params)); + } + if !first.is_empty() { + queries.push((reader::CHAIN_LINKS_SQL, &link_params)); + } + let mut results = self.query_many(&queries).await?.into_iter(); + let max_sequence_numbers = if chain_ids.is_empty() { + HashMap::new() + } else { + reader::parse_max_sequence_numbers(&results.next().unwrap_or_default())? + }; + let mut links = if first.is_empty() { + Vec::new() + } else { + reader::parse_chain_links(&results.next().unwrap_or_default())? + }; + if !rest.is_empty() { + // A chain reached from the first batch needs no fetch of its own. + let reached: HashSet = links.iter().map(|link| link.origin_chain_id).collect(); + let rest = rest + .iter() + .copied() + .filter(|chain| !reached.contains(chain)); + links.extend(reader::chain_links(self, rest).await?); + } + Ok((max_sequence_numbers, links)) + } } /// `calculate_chain_cover_index_for_events`: the chain positions and links @@ -287,7 +383,7 @@ pub async fn calculate_chain_cover_index_for_events_txn( .iter() .map(|event| Arc::clone(&event.event_id)) .collect(); - let already_persisted = store.events_with_chain_ids(&ids).await?; + let (already_persisted, out_of_band) = store.chain_cover_inputs(room_id, &ids).await?; let state_events: Vec<&ChainCoverEvent> = state_events .iter() .copied() @@ -312,7 +408,7 @@ pub async fn calculate_chain_cover_index_for_events_txn( .insert(Arc::clone(&event.event_id), auth_events); } let persisting = ids_of(&state_events); - calculate_chain_cover_index(store, room_id, &persisting, graph).await + calculate_chain_cover_index(store, room_id, &persisting, graph, out_of_band).await } /// The type, state key and auth events of the events the step looks at. @@ -332,12 +428,14 @@ fn ids_of(events: &[&ChainCoverEvent]) -> HashSet> { } /// `_calculate_chain_cover_index` for the events of `persisting`, all in -/// `room_id`. +/// `room_id`. `out_of_band` are the room's `event_auth_chain_to_calculate` +/// rows. async fn calculate_chain_cover_index( store: &dyn ChainCoverStore, room_id: &str, persisting: &HashSet>, mut graph: AuthGraph, + out_of_band: Vec, ) -> anyhow::Result, NewEventChainLinks>> { // Map from event ID to chain ID/sequence number. let mut chain_map: HashMap, Position> = HashMap::new(); @@ -348,7 +446,7 @@ async fn calculate_chain_cover_index( // Any events of the room that still need handling. These should be out // of band memberships, whose auth chains we didn't have when we first // persisted them. - for row in store.chains_to_calculate(room_id).await? { + for row in out_of_band { let auth_events = store.auth_event_ids(&row.event_id).await?; to_calculate.insert(Arc::clone(&row.event_id)); graph @@ -419,7 +517,14 @@ async fn calculate_chain_cover_index( return Ok(HashMap::new()); } - let new_positions = allocate_chain_ids(store, &graph, &to_calculate, &chain_map).await?; + // The links of every chain seen so far. The chains allocated below are + // new, so have none yet. + let chains: HashSet = chain_map.values().map(|(chain_id, _)| *chain_id).collect(); + let tree = chain_tree(&graph, &to_calculate, &chain_map); + let (max_sequence_numbers, links) = store + .max_sequence_numbers_and_links(&tree.existing_chains, chains.into_iter().collect()) + .await?; + let new_positions = allocate_chain_ids(store, tree, &chain_map, max_sequence_numbers).await?; chain_map.extend(new_positions.iter().map(|(id, pos)| (Arc::clone(id), *pos))); let mut to_return: HashMap, NewEventChainLinks> = new_positions @@ -439,8 +544,7 @@ async fn calculate_chain_cover_index( // The existing links of every chain we have seen, so that new links // reachable through them are left out. let mut chain_links = LinkMap::default(); - let chains: HashSet = chain_map.values().map(|(chain_id, _)| *chain_id).collect(); - for link in store.chain_links(chains.into_iter().collect()).await? { + for link in links { chain_links.add_link( (link.origin_chain_id, link.origin_sequence_number), (link.target_chain_id, link.target_sequence_number), @@ -507,12 +611,21 @@ enum ChainId { /// allocated but not persisted. An event extends the chain of its first /// auth event with the same type and state key, if that chain has no /// event after it yet. Otherwise it starts a new chain. -async fn allocate_chain_ids( - store: &dyn ChainCoverStore, - graph: &AuthGraph, +/// The events to calculate as a tree: each with the auth event whose chain +/// it would extend, auth events first. +struct ChainTree<'g> { + tree: Vec<(Arc, Option<&'g Arc>)>, + /// The existing chains the tree may extend. + existing_chains: Vec, +} + +/// The first half of `_allocate_chain_ids`: which chain each event would +/// extend, before the chains' sequence numbers are read. +fn chain_tree<'g>( + graph: &'g AuthGraph, to_calculate: &HashSet>, chain_map: &HashMap, Position>, -) -> anyhow::Result, Position>> { +) -> ChainTree<'g> { // Each event and the auth event whose chain it would extend, auth // events first. This turns the auth graph into a tree. let mut existing_chains: HashSet = HashSet::new(); @@ -530,11 +643,22 @@ async fn allocate_chain_ids( } tree.push((event_id, parent)); } + ChainTree { + tree, + existing_chains: existing_chains.into_iter().collect(), + } +} - let existing_chains: Vec = existing_chains.into_iter().collect(); - let mut max_sequence_numbers: HashMap = store - .max_sequence_numbers(&existing_chains) - .await? +/// The second half of `_allocate_chain_ids`: the positions along the tree, +/// given `max_sequence_numbers` for its existing chains, with the new +/// chains' IDs from the sequence in one statement. +async fn allocate_chain_ids( + store: &dyn ChainCoverStore, + tree: ChainTree<'_>, + chain_map: &HashMap, Position>, + max_sequence_numbers: HashMap, +) -> anyhow::Result, Position>> { + let mut max_sequence_numbers: HashMap = max_sequence_numbers .into_iter() .map(|(chain_id, max)| (ChainId::Existing(chain_id), max)) .collect(); @@ -542,9 +666,9 @@ async fn allocate_chain_ids( // Allocate along the tree, with placeholders for new chains, so that the // new chain IDs come from the sequence in one statement. let mut new_chains = 0; - let mut new_positions: Vec<(Arc, (ChainId, i64))> = Vec::with_capacity(tree.len()); + let mut new_positions: Vec<(Arc, (ChainId, i64))> = Vec::with_capacity(tree.tree.len()); let mut allocated: HashMap, (ChainId, i64)> = HashMap::new(); - for (event_id, parent) in tree { + for (event_id, parent) in tree.tree { let mut position = None; if let Some(auth_id) = parent { let (chain_id, sequence_number) = match allocated.get(auth_id) { diff --git a/rust/src/persist/extremities.rs b/rust/src/persist/extremities.rs index 41adf27b6e..cc26c9226a 100644 --- a/rust/src/persist/extremities.rs +++ b/rust/src/persist/extremities.rs @@ -267,7 +267,8 @@ pub async fn calculate_new_extremities( // Drop the extremities behind soft failed or rejected prev events of // the new events. Otherwise they would dangle. - for prev in reads.prevs_before_rejected(&new_prevs, true).await? { + let prevs_to_check = not_accepted_in_chunk(&new_prevs, events); + for prev in reads.prevs_before_rejected(&prevs_to_check, true).await? { result.remove(&prev); } @@ -357,7 +358,8 @@ pub async fn calculate_new_state_dag_extremities( // Drop the extremities behind rejected prev state events of the new // events, so that they do not dangle. - for prev in reads.prevs_before_rejected(&new_prevs, false).await? { + let prevs_to_check = not_accepted_in_chunk(&new_prevs, events); + for prev in reads.prevs_before_rejected(&prevs_to_check, false).await? { result.remove(&prev); } @@ -374,6 +376,24 @@ pub async fn calculate_new_state_dag_extremities( Ok(result) } +/// The IDs of `event_ids` less the accepted events of the chunk. The read +/// of rejected and soft failed prev events comes before the chunk is +/// written, so for an event of the chunk it finds nothing, or an earlier +/// copy that is accepted too, and either way adds nothing. In a new room, +/// where every prev event is in the chunk, this saves the read. +fn not_accepted_in_chunk(event_ids: &[Arc], events: &[NewEvent]) -> Vec> { + let accepted: HashSet<&Arc> = events + .iter() + .filter(|event| !event.rejected && !event.soft_failed) + .map(|event| &event.event_id) + .collect(); + event_ids + .iter() + .filter(|id| !accepted.contains(id)) + .cloned() + .collect() +} + /// `_prune_extremities`: after state resolution, drop the extremities /// whose state group is not `resolved_state_group`, if they all look stuck. /// Either all of them go or none do, as state resolution over a subset diff --git a/rust/src/persist/python.rs b/rust/src/persist/python.rs index fb97f4a98d..dd36c5df9b 100644 --- a/rust/src/persist/python.rs +++ b/rust/src/persist/python.rs @@ -1120,7 +1120,7 @@ async fn update_current_state( let result = async { let (state, stats) = state_delta::calculate_current_state(&reader, &resolver, room_id).await?; - let delta = state_delta::calculate_state_delta(&reader, room_id, &state).await?; + let delta = state_delta::calculate_state_delta(&reader, room_id, &state, true).await?; let changes = sliding_sync::calculate_sliding_sync_table_changes( &reader, room_id, diff --git a/rust/src/persist/reader.rs b/rust/src/persist/reader.rs index 289834b40f..1311f39275 100644 --- a/rust/src/persist/reader.rs +++ b/rust/src/persist/reader.rs @@ -71,7 +71,39 @@ const PREVS_BATCH_SIZE: usize = 100; /// `chain_links` fetches the links of this many chains per statement, as /// `_get_chain_links` does. -const CHAIN_LINKS_BATCH_SIZE: usize = 1000; +pub const CHAIN_LINKS_BATCH_SIZE: usize = 1000; + +/// Which of the events (`$1`) have a chain position. +pub const EVENTS_WITH_CHAIN_IDS_SQL: &str = + "SELECT event_id FROM event_auth_chains WHERE event_id = ANY($1)"; + +/// The `event_auth_chain_to_calculate` rows of the rooms (`$1`). +pub const CHAINS_TO_CALCULATE_SQL: &str = "SELECT event_id, type, state_key + FROM event_auth_chain_to_calculate + WHERE room_id = ANY($1)"; + +/// The links from the chains (`$1`) and every chain reached through them. +pub const CHAIN_LINKS_SQL: &str = "WITH RECURSIVE links(chain_id) AS ( + SELECT + DISTINCT origin_chain_id + FROM event_auth_chain_links WHERE origin_chain_id = ANY($1) + UNION + SELECT + target_chain_id + FROM event_auth_chain_links + INNER JOIN links ON (chain_id = origin_chain_id) + ) + SELECT + origin_chain_id, origin_sequence_number, + target_chain_id, target_sequence_number + FROM links + INNER JOIN event_auth_chain_links ON (chain_id = origin_chain_id)"; + +/// The highest sequence number on each of the chains (`$1`). +pub const MAX_SEQUENCE_NUMBERS_SQL: &str = "SELECT chain_id, MAX(sequence_number) + FROM event_auth_chains + WHERE chain_id = ANY($1) + GROUP BY chain_id"; /// `EventIdMembership`: the user and membership of a membership event. #[derive(Debug, Clone, PartialEq, Eq)] @@ -483,12 +515,7 @@ pub async fn events_with_chain_ids( let mut found = HashSet::new(); for batch in event_ids.chunks(SELECT_MANY_BATCH_SIZE) { let ids: Vec<&str> = batch.iter().map(|id| &**id).collect(); - let rows = q - .query( - "SELECT event_id FROM event_auth_chains WHERE event_id = ANY($1)", - &[&ids], - ) - .await?; + let rows = q.query(EVENTS_WITH_CHAIN_IDS_SQL, &[&ids]).await?; found.extend(first_column_set(&rows)?); } Ok(found) @@ -503,13 +530,12 @@ pub async fn chains_to_calculate( if room_ids.is_empty() { return Ok(Vec::new()); } - let rows = q - .query( - "SELECT event_id, type, state_key FROM event_auth_chain_to_calculate - WHERE room_id = ANY($1)", - &[&room_ids], - ) - .await?; + let rows = q.query(CHAINS_TO_CALCULATE_SQL, &[&room_ids]).await?; + parse_chains_to_calculate(&rows) +} + +/// The rows of [`CHAINS_TO_CALCULATE_SQL`]. +pub fn parse_chains_to_calculate(rows: &[Row]) -> anyhow::Result> { rows.iter() .map(|row| { Ok(ChainToCalculate { @@ -583,33 +609,8 @@ pub async fn chain_links( for chain_id in &batch { to_fetch.remove(chain_id); } - let rows = q - .query( - "WITH RECURSIVE links(chain_id) AS ( - SELECT - DISTINCT origin_chain_id - FROM event_auth_chain_links WHERE origin_chain_id = ANY($1) - UNION - SELECT - target_chain_id - FROM event_auth_chain_links - INNER JOIN links ON (chain_id = origin_chain_id) - ) - SELECT - origin_chain_id, origin_sequence_number, - target_chain_id, target_sequence_number - FROM links - INNER JOIN event_auth_chain_links ON (chain_id = origin_chain_id)", - &[&batch], - ) - .await?; - for row in rows { - let link = ChainLink { - origin_chain_id: row.try_get(0)?, - origin_sequence_number: row.try_get(1)?, - target_chain_id: row.try_get(2)?, - target_sequence_number: row.try_get(3)?, - }; + let rows = q.query(CHAIN_LINKS_SQL, &[&batch]).await?; + for link in parse_chain_links(&rows)? { // A chain reached from this batch needs no fetch of its own. to_fetch.remove(&link.origin_chain_id); links.push(link); @@ -618,6 +619,20 @@ pub async fn chain_links( Ok(links) } +/// The rows of [`CHAIN_LINKS_SQL`]. +pub fn parse_chain_links(rows: &[Row]) -> anyhow::Result> { + rows.iter() + .map(|row| { + Ok(ChainLink { + origin_chain_id: row.try_get(0)?, + origin_sequence_number: row.try_get(1)?, + target_chain_id: row.try_get(2)?, + target_sequence_number: row.try_get(3)?, + }) + }) + .collect() +} + /// The highest sequence number allocated on each of the chains, for /// `_allocate_chain_ids`. pub async fn max_sequence_numbers( @@ -627,14 +642,12 @@ pub async fn max_sequence_numbers( if chain_ids.is_empty() { return Ok(HashMap::new()); } - let rows = q - .query( - "SELECT chain_id, MAX(sequence_number) FROM event_auth_chains - WHERE chain_id = ANY($1) - GROUP BY chain_id", - &[&chain_ids], - ) - .await?; + let rows = q.query(MAX_SEQUENCE_NUMBERS_SQL, &[&chain_ids]).await?; + parse_max_sequence_numbers(&rows) +} + +/// The rows of [`MAX_SEQUENCE_NUMBERS_SQL`]. +pub fn parse_max_sequence_numbers(rows: &[Row]) -> anyhow::Result> { rows.iter() .map(|row| Ok((row.try_get(0)?, row.try_get(1)?))) .collect() @@ -647,7 +660,7 @@ fn first_column_vec(rows: &[Row]) -> anyhow::Result>> { .collect() } -fn first_column_set(rows: &[Row]) -> anyhow::Result>> { +pub(crate) fn first_column_set(rows: &[Row]) -> anyhow::Result>> { rows.iter() .map(|row| Ok(Arc::from(row.try_get::<_, String>(0)?))) .collect() @@ -1249,9 +1262,9 @@ impl PyPersistReaderForTest { 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() + state_delta::calculate_state_delta(&reader, &room_id, ¤t_state, true).await + } + .boxed() })?; let (to_delete, to_insert, _) = delta_to_py(py, &delta)?; Ok((to_delete, to_insert)) diff --git a/rust/src/persist/sliding_sync.rs b/rust/src/persist/sliding_sync.rs index 26e82998d1..a0708391f0 100644 --- a/rust/src/persist/sliding_sync.rs +++ b/rust/src/persist/sliding_sync.rs @@ -47,7 +47,6 @@ use crate::{ events::Event, state::{StateKey, StateMap}, storage::events::RedactBehaviour, - types::state_filter::StateFilter, }; const MEMBER: &str = "m.room.member"; @@ -188,38 +187,78 @@ pub struct SlidingSyncTableChanges { /// The reads the step makes. #[async_trait] pub trait SlidingSyncReads: Sync { - /// See [`PersistReader::has_sliding_sync_joined_room`]. - async fn has_sliding_sync_joined_room(&self, room_id: &str) -> anyhow::Result; - - /// The stream ordering of the room's last event of one of - /// [`DEFAULT_BUMP_EVENT_TYPES`]. See - /// [`PersistReader::last_event_pos_in_room`]. - async fn last_bump_stream_ordering(&self, room_id: &str) -> anyhow::Result>; - - /// `get_partial_filtered_current_state_ids` for the relevant state. - async fn relevant_current_state_ids(&self, room_id: &str) -> anyhow::Result; + /// What the step reads about the room, all in one round trip, as none + /// of it needs the rest. + async fn joined_room_reads(&self, room_id: &str) -> anyhow::Result; /// `get_events`: the state events found, redacted where they have been /// and leaving out the rejected ones. async fn state_events(&self, event_ids: &[Arc]) -> anyhow::Result>; } +/// What [`SlidingSyncReads::joined_room_reads`] reads. +#[derive(Debug, Default)] +pub struct JoinedRoomReads { + /// Whether the room has a `sliding_sync_joined_rooms` row. See + /// [`PersistReader::has_sliding_sync_joined_room`]. + pub has_joined_room_row: bool, + /// The stream ordering of the room's last event of one of + /// [`DEFAULT_BUMP_EVENT_TYPES`]. See + /// [`PersistReader::last_event_pos_in_room`]. + pub last_bump_stream_ordering: Option, + /// `get_partial_filtered_current_state_ids` for the relevant state. + pub relevant_current_state_ids: StateMap, +} + #[async_trait] impl SlidingSyncReads for PersistReader<'_> { - async fn has_sliding_sync_joined_room(&self, room_id: &str) -> anyhow::Result { - PersistReader::has_sliding_sync_joined_room(self, room_id).await - } - - async fn last_bump_stream_ordering(&self, room_id: &str) -> anyhow::Result> { - let last = self - .last_event_pos_in_room(room_id, Some(&DEFAULT_BUMP_EVENT_TYPES)) - .await?; - Ok(last.map(|(_, pos)| pos.stream)) - } - - async fn relevant_current_state_ids(&self, room_id: &str) -> anyhow::Result { - let filter = StateFilter::from_types(RELEVANT_STATE_TYPES.map(|t| (t, Some("")))); - (**self).current_state_ids(room_id, &filter).await + async fn joined_room_reads(&self, room_id: &str) -> anyhow::Result { + let conn = self.borrow().await?; + let bump_types: Vec<&str> = DEFAULT_BUMP_EVENT_TYPES.to_vec(); + let relevant_types: Vec<&str> = RELEVANT_STATE_TYPES.to_vec(); + // Three statements on one connection, polled together, so that they + // go out in one round trip. + let (row, bump, state) = futures::future::try_join3( + conn.query( + "SELECT room_id FROM sliding_sync_joined_rooms WHERE room_id = $1", + &[&room_id], + ), + conn.query( + "SELECT stream_ordering FROM events + LEFT JOIN rejections USING (event_id) + WHERE room_id = $1 + AND type = ANY($2) + AND NOT outlier + AND rejections.event_id IS NULL + ORDER BY stream_ordering DESC + LIMIT 1", + &[&room_id, &bump_types], + ), + conn.query( + "SELECT type, state_key, event_id FROM current_state_events + WHERE room_id = $1 AND type = ANY($2) AND state_key = ''", + &[&room_id, &relevant_types], + ), + ) + .await?; + let mut relevant_current_state_ids = StateMap::with_capacity(state.len()); + for row in &state { + let event_type: String = row.try_get(0)?; + let state_key: String = row.try_get(1)?; + let event_id: String = row.try_get(2)?; + relevant_current_state_ids.insert( + StateKey(Arc::from(event_type), Arc::from(state_key)), + Arc::from(event_id), + ); + } + Ok(JoinedRoomReads { + has_joined_room_row: !row.is_empty(), + last_bump_stream_ordering: match bump.first() { + Some(row) => Some(row.try_get(0)?), + None => None, + }, + relevant_current_state_ids, + }) } async fn state_events(&self, event_ids: &[Arc]) -> anyhow::Result> { @@ -269,6 +308,13 @@ pub async fn calculate_sliding_sync_table_changes( .map(|event| (&*event.event_id, event)) .collect(); + // What both parts below may need, read once. + let joined = if !delta.no_longer_in_room || !delta.to_insert.is_empty() { + reads.joined_room_reads(room_id).await? + } else { + JoinedRoomReads::default() + }; + // A local user's membership leaves the current state only on a state // reset out of the room. changes.to_delete_membership_snapshots = delta @@ -308,7 +354,7 @@ pub async fn calculate_sliding_sync_table_changes( if !snapshots.is_empty() { // The current state as the delta leaves it. - let mut state_ids = reads.relevant_current_state_ids(room_id).await?; + let mut state_ids = joined.relevant_current_state_ids.clone(); for key in &delta.to_delete { state_ids.remove(key); } @@ -342,15 +388,14 @@ pub async fn calculate_sliding_sync_table_changes( let mut state_ids = StateMap::new(); // A missing row is inserted in full, so that a row that exists can // be relied on. - if !reads.has_sliding_sync_joined_room(room_id).await? { + if !joined.has_joined_room_row { // A remote room that was just joined can have a backfilled // last bump event, whose negative stream ordering cannot be // compared. The API uses the membership's position instead. - changes.joined_room_bump_stamp_to_fully_insert = reads - .last_bump_stream_ordering(room_id) - .await? + changes.joined_room_bump_stamp_to_fully_insert = joined + .last_bump_stream_ordering .filter(|stream| *stream > 0); - state_ids = reads.relevant_current_state_ids(room_id).await?; + state_ids = joined.relevant_current_state_ids; } for (key, event_id) in &delta.to_insert { if is_relevant(key) { @@ -521,16 +566,12 @@ mod tests { #[async_trait] impl SlidingSyncReads for FakeReads { - async fn has_sliding_sync_joined_room(&self, _room_id: &str) -> anyhow::Result { - Ok(self.joined_room_row) - } - - async fn last_bump_stream_ordering(&self, _room_id: &str) -> anyhow::Result> { - Ok(self.last_bump) - } - - async fn relevant_current_state_ids(&self, _room_id: &str) -> anyhow::Result { - Ok(self.current_state.clone()) + async fn joined_room_reads(&self, _room_id: &str) -> anyhow::Result { + Ok(JoinedRoomReads { + has_joined_room_row: self.joined_room_row, + last_bump_stream_ordering: self.last_bump, + relevant_current_state_ids: self.current_state.clone(), + }) } async fn state_events( diff --git a/rust/src/persist/state_delta.rs b/rust/src/persist/state_delta.rs index 0a12d84a68..75dfb84681 100644 --- a/rust/src/persist/state_delta.rs +++ b/rust/src/persist/state_delta.rs @@ -617,8 +617,15 @@ pub async fn calculate_state_delta( reads: &dyn StateDeltaReads, room_id: &str, current_state: &StateMap, + room_has_events: bool, ) -> anyhow::Result { - let existing_state = reads.current_state_ids(room_id).await?; + // A room with no forward extremities has no events, so no current + // state: the chunk creates it. Python reads the empty state. + let existing_state = if room_has_events { + reads.current_state_ids(room_id).await? + } else { + StateMap::new() + }; let mut to_delete: Vec = existing_state .keys() @@ -709,6 +716,7 @@ async fn delta_from_new_state( server_name: &str, events: &[NewEvent], new_state: NewState, + room_has_events: bool, ) -> anyhow::Result> { let mut delta = if let Some(delta_ids) = new_state.delta_ids { // A known delta only adds or replaces state. @@ -717,7 +725,7 @@ async fn delta_from_new_state( ..DeltaState::default() } } else if let Some(current_state) = &new_state.current_state { - calculate_state_delta(reads, room_id, current_state).await? + calculate_state_delta(reads, room_id, current_state, room_has_events).await? } else { return Ok(None); }; @@ -807,7 +815,15 @@ pub async fn calculate_new_forward_extremities_and_state_delta( 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?; + let state_delta = delta_from_new_state( + reads, + room_id, + server_name, + events, + new_state, + !latest_event_ids.is_empty(), + ) + .await?; Ok(StateAndExtremities { new_forward_extremities: Some(new_forward_extremities), @@ -871,7 +887,15 @@ pub async fn process_state_dag_forward_extremities_and_state_delta( .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?; + let state_delta = delta_from_new_state( + reads, + room_id, + server_name, + events, + new_state, + !existing_room_dag.is_empty(), + ) + .await?; Ok(StateAndExtremities { new_forward_extremities: Some(new_room_dag), @@ -1386,7 +1410,7 @@ mod tests { ("m.room.topic", "", "$t2"), ("m.room.join_rules", "", "$j"), ]); - let delta = block_on(calculate_state_delta(&reads, ROOM, &new)).unwrap(); + let delta = block_on(calculate_state_delta(&reads, ROOM, &new, true)).unwrap(); assert_eq!( delta.to_delete, vec![ diff --git a/synapse/storage/databases/main/rust_persist.py b/synapse/storage/databases/main/rust_persist.py index 1c386bc125..72314ece67 100644 --- a/synapse/storage/databases/main/rust_persist.py +++ b/synapse/storage/databases/main/rust_persist.py @@ -206,4 +206,6 @@ def _finish_txn_stream_ids( id_gen._db.runInteraction, "MultiWriterIdGenerator._update_table", id_gen._update_stream_positions_table_txn, + # One statement: no transaction around it. + db_autocommit=True, ) diff --git a/synapse/storage/databases/main/state_deltas.py b/synapse/storage/databases/main/state_deltas.py index 1e71c9ed7b..6b28493ef0 100644 --- a/synapse/storage/databases/main/state_deltas.py +++ b/synapse/storage/databases/main/state_deltas.py @@ -130,58 +130,53 @@ class StateDeltasStore(SQLBaseStore): def get_current_state_deltas_txn( txn: LoggingTransaction, ) -> tuple[int, list[StateDelta]]: - # First we group state deltas by `stream_id` and calculate which - # groups can be returned without exceeding the provided `limit`. - sql_grouped = """ - SELECT stream_id, COUNT(*) AS c - FROM current_state_delta_stream - WHERE stream_id > ? AND stream_id <= ? - GROUP BY stream_id - ORDER BY stream_id - LIMIT ? - """ + # Group the deltas by `stream_id`, up to `limit` + 1 groups, and + # return the first group and the groups after it that fit in + # `limit` rows all together, so that a group is never split. The + # first group is always returned, so that the caller always + # makes progress. One statement, so that it runs in autocommit. group_limit = limit + 1 - txn.execute(sql_grouped, (prev_stream_id, max_stream_id, group_limit)) - grouped_rows = txn.fetchall() - - if not grouped_rows: + sql = """ + WITH groups AS ( + SELECT stream_id, COUNT(*) AS c + FROM current_state_delta_stream + WHERE stream_id > ? AND stream_id <= ? + GROUP BY stream_id + ORDER BY stream_id + LIMIT ? + ), budget AS ( + SELECT + stream_id, + CAST(SUM(c) OVER (ORDER BY stream_id) AS BIGINT) AS included_rows, + COUNT(*) OVER () AS group_count, + MIN(stream_id) OVER () AS first_stream_id + FROM groups + ) + SELECT + d.stream_id, d.room_id, d.type, d.state_key, d.event_id, + d.prev_event_id, b.group_count + FROM budget AS b + INNER JOIN current_state_delta_stream AS d + ON d.stream_id = b.stream_id + WHERE b.included_rows <= ? OR b.stream_id = b.first_stream_id + ORDER BY d.stream_id ASC + """ + txn.execute(sql, (prev_stream_id, max_stream_id, group_limit, limit)) + rows = txn.fetchall() + if not rows: # Nothing to return in the range; we are up to date through max_stream_id. return max_stream_id, [] - # Always retrieve the first group, at the bare minimum. This ensures the - # caller always makes progress, even if a single group exceeds `limit`. - fetch_upto_stream_id, included_rows = grouped_rows[0] + fetch_upto_stream_id = rows[-1][0] + group_count = rows[0][6] + included_groups = len({row[0] for row in rows}) - # Determine which other groups we can retrieve at the same time, - # without blowing the budget. - included_all_groups = True - for stream_id, count in grouped_rows[1:]: - if included_rows + count > limit: - included_all_groups = False - break - included_rows += count - fetch_upto_stream_id = stream_id - - # If we retrieved fewer groups than the limit *and* we didn't hit the - # `LIMIT ?` cap on the grouping query, we know we've caught up with - # the stream. + # If we retrieved every group *and* fewer groups than the `LIMIT` + # cap on the grouping query, we know we've caught up with the + # stream. caught_up_with_stream = ( - included_all_groups and len(grouped_rows) < group_limit + included_groups == group_count and group_count < group_limit ) - - # At this point we should have advanced, or bailed out early above. - assert fetch_upto_stream_id != prev_stream_id - - # 2) Fetch the actual rows for only the included stream_id groups. - sql_rows = """ - SELECT stream_id, room_id, type, state_key, event_id, prev_event_id - FROM current_state_delta_stream - WHERE ? < stream_id AND stream_id <= ? - ORDER BY stream_id ASC - """ - txn.execute(sql_rows, (prev_stream_id, fetch_upto_stream_id)) - rows = txn.fetchall() - clipped_stream_id = ( max_stream_id if caught_up_with_stream else fetch_upto_stream_id ) @@ -199,7 +194,9 @@ class StateDeltasStore(SQLBaseStore): ] return await self.db_pool.runInteraction( - "get_current_state_deltas", get_current_state_deltas_txn + "get_current_state_deltas", + get_current_state_deltas_txn, + db_autocommit=True, ) def _get_max_stream_id_in_current_state_deltas_txn( diff --git a/synapse/storage/util/id_generators.py b/synapse/storage/util/id_generators.py index e215f7d4b7..227b2cd0fe 100644 --- a/synapse/storage/util/id_generators.py +++ b/synapse/storage/util/id_generators.py @@ -403,7 +403,9 @@ class MultiWriterIdGenerator(AbstractStreamIdGenerator): # Cast safety: this corresponds to the types returned by the query above. rows.extend( (instance, self._return_factor * stream_id) - for instance, stream_id in cast(Iterable[tuple[str | None, int]], cur) + for instance, stream_id in cast( + Iterable[tuple[str | None, int]], cur + ) ) # Sort by stream_id (ascending, lowest -> highest) so that we handle @@ -511,6 +513,8 @@ class MultiWriterIdGenerator(AbstractStreamIdGenerator): self._db.runInteraction, "MultiWriterIdGenerator._update_table", self._update_stream_positions_table_txn, + # One statement: no transaction around it. + db_autocommit=True, ) return self._return_factor * next_id @@ -558,6 +562,8 @@ class MultiWriterIdGenerator(AbstractStreamIdGenerator): self._db.runInteraction, "MultiWriterIdGenerator._update_table", self._update_stream_positions_table_txn, + # One statement: no transaction around it. + db_autocommit=True, ) return [self._return_factor * next_id for next_id in next_ids] diff --git a/tests/storage/test_rust_persist_txn.py b/tests/storage/test_rust_persist_txn.py index 4c5a37c43e..e0ae089740 100644 --- a/tests/storage/test_rust_persist_txn.py +++ b/tests/storage/test_rust_persist_txn.py @@ -446,7 +446,9 @@ class _TwinTestBase(unittest.HomeserverTestCase): return call - def run_as_background_process(desc: str, func: Any, *args: Any) -> Any: + def run_as_background_process( + desc: str, func: Any, *args: Any, **kwargs: Any + ) -> Any: # Only the `stream_positions` updates. The invalidations start # others, which Python runs after its transaction unrecorded. if desc == "MultiWriterIdGenerator._update_table": @@ -458,7 +460,7 @@ class _TwinTestBase(unittest.HomeserverTestCase): {}, ) ) - return original_run(desc, func, *args) + return original_run(desc, func, *args, **kwargs) original_run = self.hs.run_as_background_process with ExitStack() as stack: