diff --git a/rust/src/database/postgres/write.rs b/rust/src/database/postgres/write.rs index 714ff4d82d..c4f993cac0 100644 --- a/rust/src/database/postgres/write.rs +++ b/rust/src/database/postgres/write.rs @@ -351,7 +351,7 @@ macro_rules! owned_params { /// How to turn the error of a queued statement into the one its caller /// wants, e.g. a constraint violation into a domain error. -pub type MapError = fn(WriteError) -> WriteError; +pub type MapError = Box WriteError + Send + Sync>; /// A statement queued by one of the `*_later` methods. struct Queued { @@ -460,8 +460,8 @@ impl WriteTxn { /// requests after it fail only because it aborted the transaction. async fn send<'s, T, F>( &'s self, - sql: Option<&str>, - last: impl FnOnce(&'s PooledConnection, Option) -> F, + sqls: &[&str], + last: impl FnOnce(&'s PooledConnection, Vec) -> F, context: impl FnOnce() -> String, ) -> Result where @@ -471,7 +471,11 @@ impl WriteTxn { let queued = self.take_queued(); let mut begin = self.take_begin(); - let sqls: Vec<&str> = queued.iter().map(|q| &*q.sql).chain(sql).collect(); + let sqls: Vec<&str> = queued + .iter() + .map(|q| &*q.sql) + .chain(sqls.iter().copied()) + .collect(); // Each distinct statement once: a batch queues one row per statement. let mut by_sql: HashMap<&str, Option> = HashMap::new(); for sql in &sqls { @@ -510,7 +514,7 @@ impl WriteTxn { .map(|sql| by_sql[sql].clone().expect("cached or prepared above")); let queued_statements: Vec = statements.by_ref().take(queued.len()).collect(); - let last_statement = statements.next(); + let last_statements: Vec = statements.collect(); // tokio-postgres encodes a request's parameters when it is first // polled, and if that fails it sends nothing, while the requests @@ -528,7 +532,7 @@ impl WriteTxn { q.params.iter().map(|p| p.as_ref() as &(dyn ToSql + Sync)), ) })); - let requests = join(writes, last(conn, last_statement)); + let requests = join(writes, last(conn, last_statements)); let (written, result) = match begin { None => requests.await, Some(begin) => { @@ -541,7 +545,7 @@ impl WriteTxn { } }; - for (q, outcome) in queued.iter().zip(written) { + for (q, outcome) in queued.into_iter().zip(written) { if let Err(e) = outcome { let error = self.statement_failed(&describe(&q.sql), e); return Err(match q.map_error { @@ -560,16 +564,43 @@ impl WriteTxn { params: &[&(dyn ToSql + Sync)], ) -> Result, WriteError> { self.send( - Some(sql), - |conn, statement| async move { - conn.query(&statement.expect("has a statement"), params) - .await - }, + &[sql], + |conn, statements| async move { conn.query(&statements[0], params).await }, || describe(sql), ) .await } + /// Run several statements that return rows, after the queued writes, all + /// in one round trip: their rows, in order. For reads that need nothing + /// from each other. + pub async fn query_many( + &self, + queries: &[(&str, &[&(dyn ToSql + Sync)])], + ) -> Result>, WriteError> { + let sqls: Vec<&str> = queries.iter().map(|(sql, _)| *sql).collect(); + self.send( + &sqls, + |conn, statements| async move { + let results = join_in_order( + statements + .iter() + .zip(queries) + .map(|(statement, (_, params))| conn.query(statement, params)), + ) + .await; + results.into_iter().collect() + }, + || { + sqls.iter() + .map(|sql| describe(sql)) + .collect::>() + .join("; ") + }, + ) + .await + } + /// Run a statement, after the queued writes, returning the number of rows /// it changed. pub async fn execute( @@ -578,11 +609,8 @@ impl WriteTxn { params: &[&(dyn ToSql + Sync)], ) -> Result { self.send( - Some(sql), - |conn, statement| async move { - conn.execute(&statement.expect("has a statement"), params) - .await - }, + &[sql], + |conn, statements| async move { conn.execute(&statements[0], params).await }, || describe(sql), ) .await @@ -601,11 +629,11 @@ impl WriteTxn { return Ok(0); } self.send( - Some(sql), - |conn, statement| async move { - let statement = statement.expect("has a statement"); + &[sql], + |conn, statements| async move { + let statement = &statements[0]; let counts = - join_in_order(rows.iter().map(|params| conn.execute(&statement, params))).await; + join_in_order(rows.iter().map(|params| conn.execute(statement, params))).await; counts.into_iter().sum::>() }, || describe(sql), @@ -618,7 +646,7 @@ impl WriteTxn { if self.queued.lock().expect("not poisoned").is_empty() { return Ok(()); } - self.send(None, |_, _| async { Ok(()) }, String::new).await + self.send(&[], |_, _| async { Ok(()) }, String::new).await } /// Queue a write whose result nobody reads. It goes out ahead of the @@ -786,7 +814,7 @@ impl WriteTxn { return Ok(()); } self.send( - None, + &[], |conn, _| conn.batch_execute("COMMIT"), || "committing".to_owned(), ) @@ -1267,6 +1295,48 @@ mod tests { teardown(&rt, &pool, &table); } + /// Several reads go out together, behind the queued writes, and come + /// back in order. + #[test] + fn query_many_reads_in_order_behind_queued_writes() { + let Some((rt, pool, table)) = setup() else { + return; + }; + rt.block_on(async { + let seen = run_write_interaction(&pool, "test", |txn, _| { + let table = table.clone(); + async move { + txn.insert_many_later( + &table, + vec![ + OwnedColumn::new("id", "text", vec!["a".to_owned(), "b".to_owned()]), + OwnedColumn::new("n", "bigint", vec![1i64, 2]), + ], + None, + )?; + let id = "b"; + let by_id: [&(dyn ToSql + Sync); 1] = [&id]; + let results = txn + .query_many(&[ + (&format!("SELECT n FROM {table} ORDER BY id"), &[]), + (&format!("SELECT n FROM {table} WHERE id = $1"), &by_id), + ("SELECT 42::bigint", &[]), + ]) + .await?; + Ok(results + .iter() + .map(|rows| rows.iter().map(|row| row.get::<_, i64>(0)).collect()) + .collect::>>()) + } + .boxed() + }) + .await + .unwrap(); + assert_eq!(seen, vec![vec![1, 2], vec![2], vec![42]]); + }); + teardown(&rt, &pool, &table); + } + /// A queued write that fails surfaces at the next statement, as itself /// (the next statement only fails because it aborted the transaction), /// mapped as asked. Nothing is committed. @@ -1286,12 +1356,12 @@ mod tests { txn.execute_later_mapping_error( format!("INSERT INTO {table} (id, n) VALUES ('a', 2)"), vec![], - |error| match error { + Box::new(|error| match error { WriteError::Constraint { .. } => { WriteError::Other(anyhow::anyhow!("mapped duplicate")) } other => other, - }, + }), ); txn.query(&format!("SELECT id FROM {table}"), &[]).await?; Ok(()) diff --git a/rust/src/persist/prepare.rs b/rust/src/persist/prepare.rs index dcf70c33bc..165d5bbdfa 100644 --- a/rust/src/persist/prepare.rs +++ b/rust/src/persist/prepare.rs @@ -51,10 +51,7 @@ use super::{ sticky, }; use crate::{ - clock::Clock, - database::postgres::pool::ConnectionPool, - events::Event, - state::EventMap, + clock::Clock, database::postgres::pool::ConnectionPool, events::Event, state::EventMap, }; /// A chunk of one room's events, in the shapes the steps read. diff --git a/rust/src/persist/txn/edges.rs b/rust/src/persist/txn/edges.rs index 8f92f96b0b..34abaa58a4 100644 --- a/rust/src/persist/txn/edges.rs +++ b/rust/src/persist/txn/edges.rs @@ -120,6 +120,11 @@ pub async fn update_backward_extremities( .filter(|(e, _)| !e.outlier) .flat_map(|(_, prevs)| prevs.iter().map(String::as_str)) .collect(); + // A prev event the chunk itself persists is in `events` by now, as the + // queued inserts go ahead of the read, so needs no read. + for event in events.iter().filter(|e| !e.outlier) { + potential.remove(event.event.event_id()); + } if potential.is_empty() { return Ok(()); } diff --git a/rust/src/persist/txn/events.rs b/rust/src/persist/txn/events.rs index 688e761309..6908ac13d1 100644 --- a/rust/src/persist/txn/events.rs +++ b/rust/src/persist/txn/events.rs @@ -23,6 +23,7 @@ use std::{ use anyhow::{anyhow, Context}; use log::{debug, info}; +use tokio_postgres::Row; use super::{ edges::{update_backward_extremities, BackwardEvent}, @@ -44,7 +45,7 @@ use crate::{ /// `_update_room_depths_txn`: move the room's front in the events stream /// cache to the chunk's highest stream ordering, unless it was backfilled, /// and lower its `min_depth` to the deepest new event if that is lower. -pub async fn update_room_depths( +pub fn update_room_depths( txn: &WriteTxn, room_id: &str, events: &[Persisting<'_>], @@ -69,28 +70,32 @@ pub async fn update_room_depths( ); } - // `_update_min_depth_for_room_txn`. - let rows = txn - .query( - "SELECT min_depth FROM room_depth WHERE room_id = $1", - &[&room_id], - ) - .await?; - let min_depth: Option = match rows.first() { - Some(row) => row.try_get(0).map_err(anyhow::Error::new)?, - None => None, - }; - if min_depth.is_some_and(|min_depth| depth_update >= min_depth) { - return Ok(()); - } - txn.upsert_many_later( - "room_depth", - vec![OwnedColumn::copy("room_id", "text", &[room_id])], - vec![OwnedColumn::copy("min_depth", "bigint", &[depth_update])], - )?; + // `_update_min_depth_for_room_txn`: Python reads the depth and writes + // only a lower one. The condition does the same without the read. + txn.execute_later( + "INSERT INTO room_depth (room_id, min_depth) VALUES ($1, $2) + ON CONFLICT (room_id) DO UPDATE SET min_depth = EXCLUDED.min_depth + WHERE room_depth.min_depth IS NULL OR room_depth.min_depth > EXCLUDED.min_depth", + crate::owned_params![room_id, depth_update], + ); Ok(()) } +/// Which of the chunk's events (`$1`) we have, and whether as outliers. +/// For [`update_outliers`]. +pub const OUTLIERS_SQL: &str = "SELECT event_id, outlier FROM events WHERE event_id = ANY($1)"; + +/// The redactions in the room (`$1`) of the chunk's events (`$2`), redaction +/// events included, as `_add_to_cache` counts them. For +/// [`apply_existing_redactions`]. +pub const EXISTING_REDACTIONS_SQL: &str = "SELECT redactions.event_id, redactions.redacts, \ + redaction_events.sender, redactions.recheck + FROM redactions + INNER JOIN events AS redaction_events + ON redactions.event_id = redaction_events.event_id + AND redaction_events.room_id = $1 + WHERE redactions.redacts = ANY($2)"; + /// `_update_outliers_txn`: drop the events we have already persisted. An /// event we have as an outlier, now persisted as a non-outlier, stops being /// one: it gets its state group, a row in `ex_outlier_stream`, and its @@ -104,15 +109,9 @@ pub async fn update_outliers<'a>( txn: &WriteTxn, input: &TxnInput<'_>, events: Vec>, + rows: Vec, effects: &mut PersistEffects, ) -> Result>, WriteError> { - let event_ids: Vec<&str> = events.iter().map(|e| e.event().event_id()).collect(); - let rows = txn - .query( - "SELECT event_id, outlier FROM events WHERE event_id = ANY($1)", - &[&event_ids], - ) - .await?; let mut have_persisted: HashMap = HashMap::with_capacity(rows.len()); for row in rows { have_persisted.insert( @@ -202,34 +201,22 @@ struct RedactionInfo<'a> { /// `_apply_existing_redaction_txn`: store redacted the events of the chunk /// that an earlier event or another event of the chunk redacts. Redaction /// events are never redacted here, so that circular redactions work. -pub async fn apply_existing_redactions<'a>( - txn: &WriteTxn, - room_id: &str, +/// +/// `rows` are [`EXISTING_REDACTIONS_SQL`]'s, which may cover events since +/// dropped. Also returns the events of the chunk a redaction names, applied +/// or not, redaction events among them, which `_add_to_cache` leaves out of +/// the event cache. +pub fn apply_existing_redactions<'a>( room_version: &RoomVersion, events: Vec>, -) -> Result>, WriteError> { + rows: Vec, +) -> Result<(Vec>, HashSet), WriteError> { + let chunk_ids: HashSet<&str> = events.iter().map(|e| e.event().event_id()).collect(); let event_ids: HashSet<&str> = events .iter() .filter(|e| e.event().type_() != event_type::M_ROOM_REDACTION) .map(|e| e.event().event_id()) .collect(); - if event_ids.is_empty() { - return Ok(events); - } - - let ids: Vec<&str> = event_ids.iter().copied().collect(); - let rows = txn - .query( - "SELECT redactions.event_id, redactions.redacts, redaction_events.sender, - redactions.recheck - FROM redactions - INNER JOIN events AS redaction_events - ON redactions.event_id = redaction_events.event_id - AND redaction_events.room_id = $1 - WHERE redactions.redacts = ANY($2)", - &[&room_id, &ids], - ) - .await?; let mut redactions = Vec::with_capacity(rows.len()); for row in &rows { redactions.push(RedactionInfo { @@ -248,7 +235,7 @@ pub async fn apply_existing_redactions<'a>( continue; }; if event_.type_() == event_type::M_ROOM_REDACTION - && event_ids.contains(redacts) + && chunk_ids.contains(redacts) && !event.context().is_rejected() { redactions.push(RedactionInfo { @@ -263,12 +250,17 @@ pub async fn apply_existing_redactions<'a>( } } + let redacted_targets: HashSet = redactions + .iter() + .filter(|r| chunk_ids.contains(r.redacts)) + .map(|r| r.redacts.to_owned()) + .collect(); let events_by_id: HashMap<&str, &Persisting<'a>> = events.iter().map(|e| (e.event().event_id(), e)).collect(); // From each redacted event to the event that redacts it. let mut redacted_by: HashMap = HashMap::new(); for redaction in &redactions { - if redacted_by.contains_key(redaction.redacts) { + if redacted_by.contains_key(redaction.redacts) || !event_ids.contains(redaction.redacts) { continue; } let event = events_by_id @@ -294,7 +286,7 @@ pub async fn apply_existing_redactions<'a>( } drop(events_by_id); - events + let events = events .into_iter() .map(|mut event| { if let Some(by) = redacted_by.get(event.event().event_id()) { @@ -307,7 +299,8 @@ pub async fn apply_existing_redactions<'a>( } Ok(event) }) - .collect() + .collect::, WriteError>>()?; + Ok((events, redacted_targets)) } fn domain_error(e: impl std::fmt::Display) -> WriteError { diff --git a/rust/src/persist/txn/metadata.rs b/rust/src/persist/txn/metadata.rs index 4c41020b2e..bf9dd8375c 100644 --- a/rust/src/persist/txn/metadata.rs +++ b/rust/src/persist/txn/metadata.rs @@ -17,8 +17,10 @@ //! the events themselves are stored. Push actions, redactions, relations //! and threads, search, retention, labels, expiry and room members. //! -//! The event cache prefill at its end, `_add_to_cache`, is chosen here and -//! done by Python, which holds the events. +//! The event cache prefill at its end, `_add_to_cache`, is chosen here, from +//! what the transaction knows, and done by Python, which holds the events. + +use std::collections::HashSet; use anyhow::anyhow; use serde_json::{Map, Value}; @@ -46,11 +48,13 @@ const LABELS: &str = "org.matrix.labels"; const SELF_DESTRUCT_AFTER: &str = "org.matrix.self_destruct_after"; /// `_update_metadata_tables_txn` over `events`, the events left once the -/// rejected ones are stored. `input.events` is the whole chunk. +/// rejected ones are stored. `input.events` is the whole chunk. `redacted` +/// are the events of the chunk a redaction names. pub async fn update_metadata_tables( txn: &WriteTxn, input: &TxnInput<'_>, events: &[Persisting<'_>], + redacted: &HashSet, effects: &mut PersistEffects, ) -> Result<(), WriteError> { set_push_actions(txn, input, events).await?; @@ -88,37 +92,30 @@ pub async fn update_metadata_tables( } store_room_members(txn, input, events).await?; - add_to_cache(txn, events, effects).await + add_to_cache(events, redacted, effects); + Ok(()) } /// `_add_to_cache`: the events to put in the event cache once the /// transaction commits, those neither rejected nor redacted. Python puts /// them in, from the events it holds. -async fn add_to_cache( - txn: &WriteTxn, +/// +/// Python reads the rejections and redactions back. Here a new event's +/// rejection is in its context, and the redactions of the chunk's events in +/// the room were read for `_apply_existing_redaction_txn`. (Python would +/// also leave out an event named by a redaction in another room, which +/// nothing applies.) +fn add_to_cache( events: &[Persisting<'_>], + redacted: &HashSet, effects: &mut PersistEffects, -) -> Result<(), WriteError> { - let event_ids: Vec<&str> = events.iter().map(|e| e.event().event_id()).collect(); - let rows = txn - .query( - "SELECT e.event_id, r.redacts, rej.event_id - FROM events AS e - LEFT JOIN rejections AS rej USING (event_id) - LEFT JOIN redactions AS r ON e.event_id = r.redacts - WHERE e.event_id = ANY($1)", - &[&event_ids], - ) - .await?; - for row in rows { - let event_id: String = row.try_get(0).map_err(anyhow::Error::new)?; - let redacts: Option<&str> = row.try_get(1).map_err(anyhow::Error::new)?; - let rejects: Option<&str> = row.try_get(2).map_err(anyhow::Error::new)?; - if redacts.is_none() && rejects.is_none() { - effects.add_event_cache_prefill(event_id); +) { + for persisting in events { + let event_id = persisting.event().event_id(); + if !persisting.context().is_rejected() && !redacted.contains(event_id) { + effects.add_event_cache_prefill(event_id.to_owned()); } } - Ok(()) } /// `_set_push_actions_for_event_and_users_txn`: write the push actions of diff --git a/rust/src/persist/txn/mod.rs b/rust/src/persist/txn/mod.rs index 8c020dcabe..d1bdfdfdad 100644 --- a/rust/src/persist/txn/mod.rs +++ b/rust/src/persist/txn/mod.rs @@ -45,6 +45,7 @@ use std::collections::HashMap; use anyhow::anyhow; use pyo3::PyErr; use serde_json::Value; +use tokio_postgres::types::ToSql; use self::streams::TxnStream; use super::{ @@ -260,29 +261,27 @@ async fn persist_events_txn_steps<'a>( // The chunk creates the room: store it, which claims the room ID. // The row's lock holds off a delete until we commit, as the share // lock below does. - let inserted = txn - .execute( - "INSERT INTO rooms (room_id, room_version, is_public, creator, has_auth_chain_index) - VALUES ($1, $2, $3, $4, TRUE)", - &[ - &room_id, - &create.event.room_version().identifier, - &is_public, - &create.event.sender(), - ], - ) - .await; - match inserted { - Err(WriteError::Constraint { sqlstate, .. }) if sqlstate == "23505" => { - return Err(WriteError::Other(anyhow::Error::new(RoomIdCollision( - room_id.to_owned(), - )))); - } - other => other?, - }; - } else { - // Protect against a race with deleting the room: the share lock holds - // off a delete until we commit. + let collided_room_id = room_id.to_owned(); + txn.execute_later_mapping_error( + "INSERT INTO rooms (room_id, room_version, is_public, creator, has_auth_chain_index) + VALUES ($1, $2, $3, $4, TRUE)", + crate::owned_params![ + room_id, + create.event.room_version().identifier, + is_public, + create.event.sender() + ], + Box::new(move |error| match error { + WriteError::Constraint { sqlstate, .. } if sqlstate == "23505" => { + WriteError::Other(anyhow::Error::new(RoomIdCollision(collided_room_id))) + } + error => error, + }), + ); + } + + let (Some(first), Some(last)) = (input.events.first(), input.events.last()) else { + // Python checks the room before it finds nothing to persist. let rows = txn .query( "SELECT room_version FROM rooms WHERE room_id = $1 FOR SHARE", @@ -292,9 +291,6 @@ async fn persist_events_txn_steps<'a>( if rows.is_empty() { return Err(anyhow!("Room does not exist {room_id}").into()); } - } - - let (Some(first), Some(last)) = (input.events.first(), input.events.last()) else { return Err(anyhow!("no events to persist in {room_id}").into()); }; let min_stream_order = first.stream_ordering; @@ -315,12 +311,48 @@ async fn persist_events_txn_steps<'a>( }) .collect(); - events::update_room_depths(txn, room_id, &events, effects).await?; + events::update_room_depths(txn, room_id, &events, effects)?; - let events = events::update_outliers(txn, input, events, effects).await?; + // What the chunk must read before its writes, in one round trip, behind + // the writes queued so far: whether the room is still there (its share + // lock holds off a delete until we commit; a room the chunk creates is + // locked by its insert), which events we have, their redactions, and + // some cache stream IDs for the invalidations to come. + let event_ids: Vec<&str> = events.iter().map(|e| e.event().event_id()).collect(); + let room_params: [&(dyn ToSql + Sync); 1] = [&room_id]; + let outlier_params: [&(dyn ToSql + Sync); 1] = [&event_ids]; + let redaction_params: [&(dyn ToSql + Sync); 2] = [&room_id, &event_ids]; + let cache_sequence = input.cache_stream.sequence(); + let nextval_params: [&(dyn ToSql + Sync); 2] = [&cache_sequence, &streams::PREFETCH]; + let mut queries: Vec<(&str, &[&(dyn ToSql + Sync)])> = vec![ + (events::OUTLIERS_SQL, &outlier_params), + (events::EXISTING_REDACTIONS_SQL, &redaction_params), + (streams::NEXTVAL_SQL, &nextval_params), + ]; + if new_room.is_none() { + queries.insert( + 0, + ( + "SELECT room_version FROM rooms WHERE room_id = $1 FOR SHARE", + &room_params, + ), + ); + } + let mut results = txn.query_many(&queries).await?.into_iter(); + if new_room.is_none() && results.next().is_some_and(|rows| rows.is_empty()) { + return Err(anyhow!("Room does not exist {room_id}").into()); + } + let outlier_rows = results.next().unwrap_or_default(); + let redaction_rows = results.next().unwrap_or_default(); + input + .cache_stream + .add_prefetched(&results.next().unwrap_or_default())?; + + let events = events::update_outliers(txn, input, events, outlier_rows, effects).await?; // From here on, only events we have not seen before. - let events = events::apply_existing_redactions(txn, room_id, room_version, events).await?; + let (events, redacted) = + events::apply_existing_redactions(room_version, events, redaction_rows)?; let newly_redacted = events::newly_redacted_event_ids(txn, room_id, &events).await?; // Redacted events stop being sticky. A sticky event de-outliered in // this chunk does not: its row has this chunk's stream ordering, but its @@ -353,7 +385,7 @@ async fn persist_events_txn_steps<'a>( // From here on, only events that were not rejected. let events = events::store_rejections(txn, input, events).await?; - metadata::update_metadata_tables(txn, input, &events, effects).await?; + metadata::update_metadata_tables(txn, input, &events, &redacted, effects).await?; // Last, as it reads the memberships the metadata step stored. if let Some(delta) = &state.state_delta { diff --git a/rust/src/persist/txn/streams.rs b/rust/src/persist/txn/streams.rs index 0cac138562..feaa63dd31 100644 --- a/rust/src/persist/txn/streams.rs +++ b/rust/src/persist/txn/streams.rs @@ -34,7 +34,11 @@ use crate::{ }; /// `_load_next_mult_id_txn`'s statement: `$2` IDs from the sequence `$1`. -const NEXTVAL_SQL: &str = "SELECT nextval($1::text) FROM generate_series(1, $2::bigint)"; +pub const NEXTVAL_SQL: &str = "SELECT nextval($1::text) FROM generate_series(1, $2::bigint)"; + +/// How many IDs [`TxnStream::add_prefetched`] is asked for: enough for a +/// chunk's cache invalidations, as a rule. The rest are finished unused. +pub const PREFETCH: i64 = 8; /// A stream's IDs taken by one persist, over all its attempts. An ID the /// committed attempt took is handed to Python with the effects. The rest @@ -46,6 +50,9 @@ pub struct TxnStream { sequence: &'static str, /// The IDs allocated and neither finished nor handed to Python. pending: Mutex>, + /// Allocated ahead, for [`TxnStream::next_ids`] to hand out before it + /// asks the sequence. Also in `pending`. + spare: Mutex>, } impl TxnStream { @@ -54,6 +61,7 @@ impl TxnStream { tracker, sequence, pending: Mutex::new(Vec::new()), + spare: Mutex::new(Vec::new()), } } @@ -87,6 +95,24 @@ impl TxnStream { self.pending.lock().unwrap_or_else(|e| e.into_inner()) } + fn lock_spare(&self) -> std::sync::MutexGuard<'_, Vec> { + self.spare.lock().unwrap_or_else(|e| e.into_inner()) + } + + /// The Postgres sequence the IDs come from, for a [`NEXTVAL_SQL`] sent + /// with other reads, whose rows go to [`TxnStream::add_prefetched`]. + pub fn sequence(&self) -> &'static str { + self.sequence + } + + /// Keep the IDs `rows` of a [`NEXTVAL_SQL`] allocated, for + /// [`TxnStream::next_ids`] to hand out without a round trip of its own. + pub fn add_prefetched(&self, rows: &[Row]) -> anyhow::Result<()> { + let ids = Fetch::begin(self).allocated_rows(rows)?; + self.lock_spare().extend(ids); + Ok(()) + } + /// `_load_next_id_txn`: the next ID on the stream, unfinished. pub async fn next_id(&self, txn: &WriteTxn) -> Result { let ids = self.next_ids(txn, 1).await?; @@ -98,10 +124,18 @@ impl TxnStream { /// `_load_next_mult_id_txn`: the next `n` IDs on the stream, /// unfinished. pub async fn next_ids(&self, txn: &WriteTxn, n: usize) -> Result, WriteError> { - let n = i64::try_from(n).map_err(anyhow::Error::new)?; - let fetch = Fetch::begin(self); - let rows = txn.query(NEXTVAL_SQL, &[&self.sequence, &n]).await?; - Ok(fetch.allocated_rows(&rows)?) + let mut ids: Vec = { + let mut spare = self.lock_spare(); + let taken = n.min(spare.len()); + spare.drain(..taken).collect() + }; + if ids.len() < n { + let more = i64::try_from(n - ids.len()).map_err(anyhow::Error::new)?; + let fetch = Fetch::begin(self); + let rows = txn.query(NEXTVAL_SQL, &[&self.sequence, &more]).await?; + ids.extend(fetch.allocated_rows(&rows)?); + } + Ok(ids) } /// As [`Self::next_ids`], but outside a transaction, as `get_next_mult` diff --git a/tests/rest/client/bench_createroom.py b/tests/rest/client/bench_createroom.py index fc49bffc5f..402650a0a2 100644 --- a/tests/rest/client/bench_createroom.py +++ b/tests/rest/client/bench_createroom.py @@ -92,6 +92,13 @@ class CreateRoomBench(HomeserverTestCase): room.register_servlets, ] + def default_config(self) -> dict[str, Any]: + config = super().default_config() + # The test default of one entry makes every read after persisting a + # database fetch, which no server does. + config["event_cache_size"] = 10000 + return config + def setUp(self) -> None: if not os.environ.get("SYNAPSE_BENCH_CREATEROOM"): raise SkipTest("set SYNAPSE_BENCH_CREATEROOM to run the benchmark")