Send BEGIN with the first statement in the Rust sync store

The sync store's transactions (`SyncReader::begin` and `begin_write`),
which the Rust sync and sliding sync handlers and the persister's reader
use, still sent `BEGIN` on its own and waited for it.

As for `WriteTxn`, the `BEGIN` now goes out with the first statement:
with its execution when the connection has it cached, otherwise with its
prepare. If the `BEGIN` fails, the connection is discarded. A
transaction that runs nothing sends neither `BEGIN` nor `COMMIT`.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01PVLsNh6akt4ifEizY4D1ta

rust/src/storage/sync_store.rs

@@ -34,7 +34,10 @@
 
 use std::{
     collections::{HashMap, HashSet},
-    sync::{Arc, Mutex},
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Arc, Mutex,
+    },
     time::{Duration, Instant},
 };
 
@@ -229,6 +232,13 @@ pub struct Borrowed<'r, 'a> {
     reader: &'r SyncReader<'a>,
     conn: Option<PooledConnection>,
     in_transaction: bool,
+    /// The `BEGIN` of a transaction no statement has sent yet. The first
+    /// statement sends it in front of its own request, so that they share a
+    /// round trip.
+    pending_begin: Mutex<Option<&'static str>>,
+    /// The `BEGIN` failed, so the statement sent behind it ran outside the
+    /// transaction: the connection is discarded, not reused.
+    begin_failed: AtomicBool,
 }
 
 impl<'a> SyncReader<'a> {
@@ -263,6 +273,8 @@ impl<'a> SyncReader<'a> {
             reader: self,
             conn: Some(conn),
             in_transaction: false,
+            pending_begin: Mutex::new(None),
+            begin_failed: AtomicBool::new(false),
         })
     }
 
@@ -273,13 +285,13 @@ impl<'a> SyncReader<'a> {
     /// [`Borrowed::commit`].
     pub(crate) async fn begin(&self) -> anyhow::Result<Borrowed<'_, 'a>> {
         let mut borrowed = self.borrow().await?;
-        // Set before awaiting: if we are dropped after BEGIN is sent but
-        // before it returns, the transaction is open and must be rolled back
-        // before the connection goes back to the pool.
+        // Set before anything is sent: if we are dropped after the BEGIN goes
+        // out but before it returns, the transaction is open and must be
+        // rolled back before the connection goes back to the pool.
         borrowed.in_transaction = true;
-        borrowed
-            .execute("BEGIN ISOLATION LEVEL REPEATABLE READ READ ONLY")
-            .await?;
+        // Sent with the first statement.
+        *lock_ignoring_poison(&borrowed.pending_begin) =
+            Some("BEGIN ISOLATION LEVEL REPEATABLE READ READ ONLY");
         self.update_stats(|stats| stats.transactions += 1);
         Ok(borrowed)
     }
@@ -291,13 +303,13 @@ impl<'a> SyncReader<'a> {
     /// with [`Borrowed::commit`].
     pub(crate) async fn begin_write(&self) -> anyhow::Result<Borrowed<'_, 'a>> {
         let mut borrowed = self.borrow().await?;
-        // Set before awaiting: if we are dropped after BEGIN is sent but
-        // before it returns, the transaction is open and must be rolled back.
-        // A ROLLBACK outside one only warns.
+        // Set before anything is sent: if we are dropped after the BEGIN goes
+        // out but before it returns, the transaction is open and must be
+        // rolled back. A ROLLBACK outside one only warns.
         borrowed.in_transaction = true;
-        borrowed
-            .execute("BEGIN ISOLATION LEVEL REPEATABLE READ")
-            .await?;
+        // Sent with the first statement.
+        *lock_ignoring_poison(&borrowed.pending_begin) =
+            Some("BEGIN ISOLATION LEVEL REPEATABLE READ");
         self.update_stats(|stats| stats.transactions += 1);
         Ok(borrowed)
     }
@@ -2240,6 +2252,26 @@ impl Borrowed<'_, '_> {
         self.conn.as_ref().expect("held until dropped")
     }
 
+    /// Await `request`, sending the transaction's `BEGIN` in front of it if
+    /// no statement has sent it yet. tokio-postgres sends a request when its
+    /// future is first polled, and `join` polls the `BEGIN` first.
+    async fn after_begin<T>(
+        &self,
+        request: impl std::future::Future<Output = Result<T, tokio_postgres::Error>>,
+    ) -> anyhow::Result<Result<T, tokio_postgres::Error>> {
+        let begin = lock_ignoring_poison(&self.pending_begin).take();
+        let Some(begin) = begin else {
+            return Ok(request.await);
+        };
+        let (begun, result) =
+            futures::future::join(self.conn().batch_execute(begin), request).await;
+        if let Err(e) = begun {
+            self.begin_failed.store(true, Ordering::Relaxed);
+            return Err(anyhow::Error::new(e).context(format!("running {begin}")));
+        }
+        Ok(result)
+    }
+
     /// Run a statement through the connection's prepared statement cache.
     pub async fn query(
         &self,
@@ -2248,14 +2280,19 @@ impl Borrowed<'_, '_> {
     ) -> anyhow::Result<Vec<Row>> {
         let start = Instant::now();
         let conn = self.conn();
-        let statement = conn
-            .prepare_cached(sql, false)
-            .await
-            .with_context(|| format!("preparing {}", sql.trim()))?;
-        let rows = conn
-            .query(&statement, params)
-            .await
-            .with_context(|| format!("running {}", sql.trim()))?;
+        // A pending `BEGIN` goes with the statement's first request: its
+        // execution if it is cached, otherwise its prepare.
+        let rows = match conn.cached(sql) {
+            Some(statement) => self.after_begin(conn.query(&statement, params)).await?,
+            None => {
+                let statement = self
+                    .after_begin(conn.prepare_cached(sql, false))
+                    .await?
+                    .with_context(|| format!("preparing {}", sql.trim()))?;
+                conn.query(&statement, params).await
+            }
+        }
+        .with_context(|| format!("running {}", sql.trim()))?;
         self.record(start, rows.len());
         if log::log_enabled!(log::Level::Debug) {
             let head: String = sql.split_whitespace().take(6).collect::<Vec<_>>().join(" ");
@@ -2276,9 +2313,8 @@ impl Borrowed<'_, '_> {
     ) -> anyhow::Result<Vec<Vec<tokio_postgres::SimpleQueryRow>>> {
         let start = Instant::now();
         let messages = self
-            .conn()
-            .simple_query(sql)
-            .await
+            .after_begin(self.conn().simple_query(sql))
+            .await?
             .with_context(|| format!("running {}", sql.trim()))?;
         let mut results: Vec<Vec<tokio_postgres::SimpleQueryRow>> = Vec::new();
         let mut current = Vec::new();
@@ -2300,9 +2336,8 @@ impl Borrowed<'_, '_> {
     /// Run statements that return nothing.
     pub(crate) async fn execute(&self, sql: &str) -> anyhow::Result<()> {
         let start = Instant::now();
-        self.conn()
-            .batch_execute(sql)
-            .await
+        self.after_begin(self.conn().batch_execute(sql))
+            .await?
             .with_context(|| format!("running {sql}"))?;
         self.record(start, 0);
         Ok(())
@@ -2319,6 +2354,10 @@ impl Borrowed<'_, '_> {
     /// Commit the transaction opened by [`SyncReader::begin`] and return
     /// the connection to the pool.
     pub async fn commit(mut self) -> anyhow::Result<()> {
+        if lock_ignoring_poison(&self.pending_begin).take().is_some() {
+            // No statement ran, so there is no transaction to commit.
+            self.in_transaction = false;
+        }
         if self.in_transaction {
             self.execute("COMMIT").await?;
             self.in_transaction = false;
@@ -2330,7 +2369,13 @@ impl Borrowed<'_, '_> {
 impl Drop for Borrowed<'_, '_> {
     fn drop(&mut self) {
         let Some(conn) = self.conn.take() else { return };
-        if !self.in_transaction {
+        if self.begin_failed.load(Ordering::Relaxed) {
+            // The session state is unknown: discard the connection.
+            drop(PooledConnection::take(conn));
+            return;
+        }
+        if !self.in_transaction || lock_ignoring_poison(&self.pending_begin).is_some() {
+            // No transaction was opened.
             return;
         }
         // Roll back before the connection goes back to the pool, so the next
@@ -2422,6 +2467,36 @@ impl PySyncStoreForTest {
         }
     }
 
+    /// What the first statement of a [`SyncReader::begin`] and of a
+    /// [`SyncReader::begin_write`] transaction sees: `(isolation level,
+    /// read only)` of each. A transaction that runs nothing also commits.
+    fn transaction_settings(&self, py: Python<'_>) -> PyResult<Vec<(String, String)>> {
+        self.run(py, |reader| {
+            async move {
+                let mut settings = Vec::new();
+                for write in [false, true] {
+                    let txn = if write {
+                        reader.begin_write().await?
+                    } else {
+                        reader.begin().await?
+                    };
+                    let rows = txn
+                        .query(
+                            "SELECT current_setting('transaction_isolation'),
+                                current_setting('transaction_read_only')",
+                            &[],
+                        )
+                        .await?;
+                    settings.push((rows[0].try_get(0)?, rows[0].try_get(1)?));
+                    txn.commit().await?;
+                }
+                reader.begin().await?.commit().await?;
+                Ok(settings)
+            }
+            .boxed()
+        })
+    }
+
     /// [`crate::visibility::filter_events_for_server`]. `None` where Python
     /// would wait for the full state.
     #[pyo3(signature = (target_server_name, local_server_name, events, *, redact, filter_out_erased_senders, filter_out_remote_partial_state_events))]

synapse/synapse_rust/sync.pyi

@@ -374,6 +374,9 @@ class SyncStoreForTest:
     Rust sync handler's caches. For `tests/storage/test_sync_store.py`."""
 
     def __init__(self, pool: ConnectionPool, handler: SyncHandler) -> None: ...
+    def transaction_settings(self) -> list[tuple[str, str]]:
+        """The `(transaction_isolation, transaction_read_only)` the first
+        statement of a read-only and of a writing transaction sees."""
     def filter_events_for_server(
         self,
         target_server_name: str,

tests/storage/test_sync_store.py

@@ -101,6 +101,14 @@ class SyncStoreTestCase(unittest.HomeserverTestCase):
         self.room_id = self.helper.create_room_as(self.alice, tok=self.alice_tok)
         self.helper.join(self.room_id, self.bob, tok=self.bob_tok)
 
+    def test_transactions_begin_with_their_first_statement(self) -> None:
+        """A transaction's `BEGIN` goes out with its first statement, which runs
+        inside the transaction."""
+        self.assertEqual(
+            self.rust.transaction_settings(),
+            [("repeatable read", "on"), ("repeatable read", "off")],
+        )
+
     def test_rooms_for_user(self) -> None:
         other = self.helper.create_room_as(self.bob, tok=self.bob_tok)
         self.helper.invite(other, self.bob, self.alice, tok=self.bob_tok)