From 2f0faa5e481d8e782aa4f7889d121e25f2c99c95 Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Mon, 19 Oct 2026 09:48:32 +0000 Subject: [PATCH] Trace Rust reads on the native driver MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The read paths now report the same spans as the write path: - `PoolQuerier` (event fetch, stream ID allocation): `db.connection` from checkout until it drops, and a `db.query` per statement, tagged `db.prepared` when the statement missed the connection's cache. - The sync store's borrowed connections: `db.connection` from checkout until the drop, and for a transaction a `db.txn` (`db.txn_desc` of `sync` or `sync_write`, `db.txn_id`, `db.readonly`) with a `commit` log, or a `db.rollback` child when a dropped transaction rolls back. Each statement is a `db.query`, with `db.begin` when the pending BEGIN went with it. - The state resolution store: `db.connection` → `db.txn` (`state_res`, read only) → a `db.query` each for the BEGIN, every read and the COMMIT. As on the write path and in Python, `db.connection` starts once the connection is in hand, and the wait for it is a `checked out` log on the caller's span. A failed `db.query` is marked with the `error` tag and log. `db.query` spans take the held connection or transaction span as an explicit parent. The tests collect spans on a current-thread runtime: the registry releases a span's enter through the thread's default subscriber, so a span entered on a worker thread would never close under a test's thread-local subscriber. Production installs the layer globally. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01W4zxygfJC93SsugxskhoJG --- rust/src/database/postgres/write.rs | 7 +- rust/src/logging/opentracing/layer.rs | 6 + rust/src/logging/opentracing/mod.rs | 47 ++++++ rust/src/state/store.rs | 152 ++++++++++++++--- rust/src/storage/pool_querier.rs | 131 +++++++++++++-- rust/src/storage/sync_store.rs | 233 +++++++++++++++++++++++--- 6 files changed, 524 insertions(+), 52 deletions(-) diff --git a/rust/src/database/postgres/write.rs b/rust/src/database/postgres/write.rs index 8c41817461..245a06d9f4 100644 --- a/rust/src/database/postgres/write.rs +++ b/rust/src/database/postgres/write.rs @@ -1089,7 +1089,7 @@ where let mut elapsed = Duration::ZERO; let mut attempt = 1; // One ID for every attempt, as Python names the transaction once. - let txn_id = NEXT_TXN_ID.fetch_add(1, Ordering::Relaxed); + let txn_id = next_txn_id(); // Python's `db.{desc}` span: the whole interaction, every attempt and // the waits for a connection included. let span = tracing::info_span!( @@ -1122,6 +1122,11 @@ where /// Numbers transactions for their `db.txn_id` tag, as Python numbers its own. static NEXT_TXN_ID: AtomicU64 = AtomicU64::new(0); +/// A number for a transaction's `db.txn_id` tag, `{desc}-{number:x}`. +pub(crate) fn next_txn_id() -> u64 { + NEXT_TXN_ID.fetch_add(1, Ordering::Relaxed) +} + /// What each write interaction run in a [`with_txn_timings`] scope took, /// as `(name, seconds)`, in the order they finished. pub type TxnTimings = Vec<(String, f64)>; diff --git a/rust/src/logging/opentracing/layer.rs b/rust/src/logging/opentracing/layer.rs index d702791b01..7c2bbea23d 100644 --- a/rust/src/logging/opentracing/layer.rs +++ b/rust/src/logging/opentracing/layer.rs @@ -737,6 +737,12 @@ pub(crate) mod testing { /// Run `f` with a [`JaegerLayer`] as this thread's subscriber, returning /// the spans it reported in the order they closed. + /// + /// Spans must only be entered on this thread: run async code on a + /// current-thread runtime. The registry releases the reference an enter + /// takes through the thread's default subscriber, so a span entered on a + /// tokio worker thread never closes. (In production the layer is the + /// global default, which every thread sees.) pub(crate) fn traced( sampler: SamplerConfig, f: impl FnOnce() -> R, diff --git a/rust/src/logging/opentracing/mod.rs b/rust/src/logging/opentracing/mod.rs index 9b7782dc63..bc62138d9a 100644 --- a/rust/src/logging/opentracing/mod.rs +++ b/rust/src/logging/opentracing/mod.rs @@ -145,6 +145,53 @@ pub fn record_error(span: &tracing::Span, error: &dyn std::fmt::Display) { tracing::info!(parent: span, event = "error", message = %error); } +/// A `db.query` span for one round trip running `sql`, a child of `parent` +/// (or of the current span, if `None`). `db.statement` is only filled in +/// when the span is sampled. `db.queued`, `db.prepared` and `db.begin` are +/// left for the caller to record. +pub fn db_query_span(parent: Option<&tracing::Span>, sql: &str) -> tracing::Span { + let span = match parent { + Some(parent) => tracing::info_span!( + parent: parent, + "db.query", + db.type = "sql", + db.statement = tracing::field::Empty, + db.queued = tracing::field::Empty, + db.prepared = tracing::field::Empty, + db.begin = tracing::field::Empty, + error = tracing::field::Empty, + ), + None => tracing::info_span!( + "db.query", + db.type = "sql", + db.statement = tracing::field::Empty, + db.queued = tracing::field::Empty, + db.prepared = tracing::field::Empty, + db.begin = tracing::field::Empty, + error = tracing::field::Empty, + ), + }; + if !sql.is_empty() && layer::is_sampled(&span) { + span.record("db.statement", one_line_sql(sql).as_str()); + } + span +} + +/// Run `request` in the `db.query` span `span`, marking the span failed +/// if the request fails. +pub async fn traced_query( + span: tracing::Span, + request: impl std::future::Future>, +) -> Result { + use tracing::Instrument; + + let result = request.instrument(span.clone()).await; + if let Err(error) = &result { + record_error(&span, error); + } + result +} + /// `sql` on one line, for the `db.statement` tag: each line trimmed, blank /// lines dropped, joined with spaces. Python's `_make_sql_one_line`. pub fn one_line_sql(sql: &str) -> String { diff --git a/rust/src/state/store.rs b/rust/src/state/store.rs index 26f8b5b0df..7935e6c5a7 100644 --- a/rust/src/state/store.rs +++ b/rust/src/state/store.rs @@ -33,9 +33,12 @@ use pyo3::{ }; use tokio::runtime::Handle; use tokio_postgres::{types::ToSql, Row}; +use tracing::Instrument; use super::{auth_chain, EventMap}; use crate::database::postgres::pool::{ConnectionPool, PooledConnection, PyConnectionPool}; +use crate::database::postgres::write::next_txn_id; +use crate::logging::opentracing::{db_query_span, traced_query}; use crate::storage::caches::StoreCaches; use crate::storage::events::{EventLoadStats, EventLoader, Querier}; @@ -118,6 +121,11 @@ pub struct DbStats { /// connection goes back to the pool. pub struct PostgresStateResStore { conn: Option, + /// `db.connection`, from the checkout until the store is done. Only + /// held, so that it lasts as long as the store. + _connection_span: tracing::Span, + /// `db.txn`, the read-only transaction, a child of `connection_span`. + txn_span: tracing::Span, handle: Handle, /// Events, chain positions and chain links loaded earlier, consulted /// before the database. @@ -134,12 +142,30 @@ impl PostgresStateResStore { handle: Handle, caches: Option>, ) -> anyhow::Result { + let waiting = Instant::now(); let conn = pool .get() .await .map_err(|e| anyhow::Error::new(e).context("checking out a database connection"))?; + // The wait goes on the caller's span, as on the write path: Python's + // `db.connection` starts once the connection is in hand. + tracing::info!( + message = "checked out", + wait_ms = waiting.elapsed().as_secs_f64() * 1000.0 + ); + let connection_span = tracing::info_span!("db.connection"); + let txn_id = next_txn_id(); + let txn_span = tracing::info_span!( + parent: &connection_span, + "db.txn", + db.txn_desc = "state_res", + db.txn_id = %format_args!("state_res-{txn_id:x}"), + db.readonly = true, + ); let store = PostgresStateResStore { conn: Some(conn), + _connection_span: connection_span, + txn_span, handle, caches, db_stats: Mutex::default(), @@ -147,11 +173,13 @@ impl PostgresStateResStore { // If this fails, dropping the store sends a ROLLBACK, which is // harmless outside a transaction. let start = Instant::now(); - store - .client() - .batch_execute("BEGIN ISOLATION LEVEL REPEATABLE READ READ ONLY") - .await - .context("BEGIN")?; + const BEGIN: &str = "BEGIN ISOLATION LEVEL REPEATABLE READ READ ONLY"; + traced_query( + db_query_span(Some(&store.txn_span), BEGIN), + store.client().batch_execute(BEGIN), + ) + .await + .context("BEGIN")?; store.record(start, 0); Ok(store) } @@ -160,7 +188,13 @@ impl PostgresStateResStore { pub async fn finish(mut self) -> anyhow::Result { let stats = self.stats(); if let Some(conn) = self.conn.take() { - conn.batch_execute("COMMIT").await.context("COMMIT")?; + tracing::info!(parent: &self.txn_span, message = "commit"); + traced_query( + db_query_span(Some(&self.txn_span), "COMMIT"), + conn.batch_execute("COMMIT"), + ) + .await + .context("COMMIT")?; } Ok(stats) } @@ -200,14 +234,24 @@ impl PostgresStateResStore { ) -> anyhow::Result> { let start = Instant::now(); let client = self.client(); - let statement = client - .prepare_cached(sql, false) - .await - .with_context(|| format!("preparing {}", sql.trim()))?; - let rows = client - .query(&statement, params) - .await - .with_context(|| format!("running {}", sql.trim()))?; + let span = db_query_span(Some(&self.txn_span), sql); + let rows = traced_query(span, async { + let statement = match client.cached(sql) { + Some(statement) => statement, + None => { + tracing::Span::current().record("db.prepared", true); + client + .prepare_cached(sql, false) + .await + .with_context(|| format!("preparing {}", sql.trim()))? + } + }; + client + .query(&statement, params) + .await + .with_context(|| format!("running {}", sql.trim())) + }) + .await?; self.record(start, rows.len()); if log::log_enabled!(log::Level::Debug) { let head: String = sql.split_whitespace().take(6).collect::>().join(" "); @@ -226,12 +270,16 @@ impl Drop for PostgresStateResStore { if let Some(conn) = self.conn.take() { // Roll back before the connection goes back to the pool, so the // next user doesn't find it inside our transaction. - self.handle.spawn(async move { - if let Err(e) = conn.batch_execute("ROLLBACK").await { - warn!("Failed to roll back state res transaction: {e}"); + let span = tracing::info_span!(parent: &self.txn_span, "db.rollback"); + self.handle.spawn( + async move { + if let Err(e) = conn.batch_execute("ROLLBACK").await { + warn!("Failed to roll back state res transaction: {e}"); + } + drop(conn); } - drop(conn); - }); + .instrument(span), + ); } } } @@ -463,3 +511,69 @@ pub(super) fn register_test_functions(py: Python<'_>, m: &Bound<'_, PyModule>) - m.add_function(wrap_pyfunction!(load_events_for_test, m)?)?; Ok(()) } + +#[cfg(test)] +mod tracing_tests { + //! Against a real server: these run when `SYNAPSE_TEST_POSTGRES_DSN` is + //! set, and otherwise do nothing. + + use super::*; + use crate::logging::opentracing::layer::testing::{tag, traced}; + use crate::logging::opentracing::thrift::TagValue; + use crate::logging::opentracing::SamplerConfig; + use crate::storage::pool_querier::tests::test_pool; + + /// `db.connection` → `db.txn` → a `db.query` each for the BEGIN, the + /// read and the COMMIT. + #[test] + fn traced_store() { + let Some((rt, pool)) = test_pool() else { + return; + }; + let (_, spans) = traced(SamplerConfig::Const(true), || { + rt.block_on(async { + let store = PostgresStateResStore::begin(&pool, rt.handle().clone(), None) + .await + .unwrap(); + store.query("SELECT 1", &[]).await.unwrap(); + store.finish().await.unwrap(); + }) + }); + + let operations: Vec<&str> = spans.iter().map(|span| &*span.operation).collect(); + assert_eq!( + operations, + [ + "db.query", + "db.query", + "db.query", + "db.txn", + "db.connection" + ] + ); + let (queries, rest) = spans.split_at(3); + let (txn, connection) = (&rest[0], &rest[1]); + assert_eq!(txn.parent_id, connection.span_id); + assert_eq!( + tag(txn, "db.txn_desc"), + Some(&TagValue::Str("state_res".into())) + ); + let statements: Vec<_> = queries + .iter() + .map(|query| { + assert_eq!(query.parent_id, txn.span_id); + tag(query, "db.statement").cloned() + }) + .collect(); + assert_eq!( + statements, + [ + Some(TagValue::Str( + "BEGIN ISOLATION LEVEL REPEATABLE READ READ ONLY".into() + )), + Some(TagValue::Str("SELECT 1".into())), + Some(TagValue::Str("COMMIT".into())), + ] + ); + } +} diff --git a/rust/src/storage/pool_querier.rs b/rust/src/storage/pool_querier.rs index 14554c33ee..fe6b707dcd 100644 --- a/rust/src/storage/pool_querier.rs +++ b/rust/src/storage/pool_querier.rs @@ -19,6 +19,7 @@ use async_trait::async_trait; use tokio_postgres::{types::ToSql, Row}; use crate::database::postgres::pool::{ConnectionPool, PooledConnection}; +use crate::logging::opentracing::{db_query_span, traced_query}; use crate::storage::events::Querier; /// Runs each statement on its own, in autocommit, on one connection checked @@ -29,31 +30,141 @@ use crate::storage::events::Querier; /// waiting for a connection could stall until the checkout times out: the /// Rust readers that join it hold pooled connections while they wait, and /// enough of them leave none for the load. +/// +/// Traced as a `db.connection` span from the checkout until the drop, with a +/// `db.query` child per statement. pub struct PoolQuerier { conn: PooledConnection, + span: tracing::Span, } impl PoolQuerier { pub async fn checkout(pool: &ConnectionPool) -> anyhow::Result { + let waiting = std::time::Instant::now(); let conn = pool .get() .await .map_err(|e| anyhow::anyhow!("checking out a database connection: {e}"))?; - Ok(PoolQuerier { conn }) + // The wait goes on the caller's span, as on the write path: Python's + // `db.connection` starts once the connection is in hand. + tracing::info!( + message = "checked out", + wait_ms = waiting.elapsed().as_secs_f64() * 1000.0 + ); + // Never entered, so it lasts until it closes: when the querier and + // the queries it parents are done. + let span = tracing::info_span!("db.connection"); + Ok(PoolQuerier { conn, span }) } } #[async_trait] impl Querier for PoolQuerier { async fn query(&self, sql: &str, params: &[&(dyn ToSql + Sync)]) -> anyhow::Result> { - let statement = self - .conn - .prepare_cached(sql, false) - .await - .with_context(|| format!("preparing {}", sql.trim()))?; - self.conn - .query(&statement, params) - .await - .with_context(|| format!("running {}", sql.trim())) + let span = db_query_span(Some(&self.span), sql); + traced_query(span, async { + // A statement missing from the cache costs a prepare round trip. + let statement = match self.conn.cached(sql) { + Some(statement) => statement, + None => { + tracing::Span::current().record("db.prepared", true); + self.conn + .prepare_cached(sql, false) + .await + .with_context(|| format!("preparing {}", sql.trim()))? + } + }; + self.conn + .query(&statement, params) + .await + .with_context(|| format!("running {}", sql.trim())) + }) + .await + } +} + +#[cfg(test)] +pub(crate) mod tests { + //! Against a real server: these run when `SYNAPSE_TEST_POSTGRES_DSN` is + //! set (e.g. to `host=postgres user=postgres password=postgres`), and + //! otherwise do nothing. + + use tokio::runtime::Runtime; + + use super::*; + use crate::database::postgres::pool::{create_pool, PoolConfig}; + use crate::logging::opentracing::layer::testing::{tag, traced}; + use crate::logging::opentracing::thrift::TagValue; + use crate::logging::opentracing::SamplerConfig; + + pub(crate) fn test_pool() -> Option<(Runtime, ConnectionPool)> { + let Ok(dsn) = std::env::var("SYNAPSE_TEST_POSTGRES_DSN") else { + eprintln!("skipping: set SYNAPSE_TEST_POSTGRES_DSN to run"); + return None; + }; + // Current-thread, so that spans are only entered on the thread whose + // subscriber `traced` sets. + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let pool = create_pool( + &dsn, + rt.handle().clone(), + PoolConfig { + max_size: 2, + ..PoolConfig::default() + }, + ) + .unwrap(); + Some((rt, pool)) + } + + /// A `db.connection` span from checkout to drop, with a `db.query` per + /// statement saying whether it needed a prepare. + #[test] + fn traced_queries() { + let Some((rt, pool)) = test_pool() else { + return; + }; + let (_, spans) = traced(SamplerConfig::Const(true), || { + rt.block_on(async { + let querier = PoolQuerier::checkout(&pool).await.unwrap(); + // A statement no other test prepares, then again cached. + let sql = "SELECT 1 AS pool_querier_traced\n WHERE true"; + querier.query(sql, &[]).await.unwrap(); + querier.query(sql, &[]).await.unwrap(); + querier + .query("SELECT nonsense FROM", &[]) + .await + .unwrap_err(); + }) + }); + + let [first, second, failed, connection] = &spans[..] else { + panic!("expected four spans: {spans:#?}") + }; + assert_eq!(connection.operation, "db.connection"); + for query in [first, second, failed] { + assert_eq!(query.operation, "db.query"); + assert_eq!(query.parent_id, connection.span_id); + } + assert_eq!( + tag(first, "db.statement"), + Some(&TagValue::Str( + "SELECT 1 AS pool_querier_traced WHERE true".into() + )) + ); + assert_eq!(tag(first, "db.prepared"), Some(&TagValue::Bool(true))); + assert_eq!(tag(second, "db.prepared"), None); + // Starts once the connection is in hand, and lasts until the querier + // is dropped, after its queries. + assert!(connection.start <= first.start); + assert!( + connection.start + connection.duration >= failed.start + failed.duration, + "{connection:?} ends before {failed:?}" + ); + assert_eq!(tag(failed, "error"), Some(&TagValue::Bool(true))); + assert_eq!(tag(first, "error"), None); } } diff --git a/rust/src/storage/sync_store.rs b/rust/src/storage/sync_store.rs index fea550d398..04223ba5bd 100644 --- a/rust/src/storage/sync_store.rs +++ b/rust/src/storage/sync_store.rs @@ -48,11 +48,14 @@ use log::{debug, warn}; use pyo3::{prelude::*, types::PyModule}; use tokio::runtime::Handle; use tokio_postgres::{types::ToSql, Row}; +use tracing::Instrument; use crate::caches::lock_ignoring_poison; use crate::config::SynapseHomeServerConfig; use crate::database::postgres::pool::{ConnectionPool, PooledConnection, PyConnectionPool}; +use crate::database::postgres::write::next_txn_id; use crate::events::Event; +use crate::logging::opentracing::{db_query_span, traced_query}; use crate::state::{StateKey, StateMap}; use crate::storage::caches::StoreCaches; use crate::storage::events::{EventLoadStats, EventLoader, Querier, RedactBehaviour}; @@ -231,6 +234,10 @@ pub struct SyncReader<'a> { pub struct Borrowed<'r, 'a> { reader: &'r SyncReader<'a>, conn: Option, + /// `db.connection`, from the checkout until the drop. + connection_span: tracing::Span, + /// `db.txn` once a transaction is begun, a child of `connection_span`. + txn_span: Option, 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 @@ -265,6 +272,16 @@ impl<'a> SyncReader<'a> { // Keep the pool's error as the cause, so that a caller can tell // a lost connection, which is worth retrying. .map_err(|e| anyhow::Error::new(e).context("checking out a database connection"))?; + let waiting = start; + // The wait goes on the caller's span, as on the write path: Python's + // `db.connection` starts once the connection is in hand. + tracing::info!( + message = "checked out", + wait_ms = waiting.elapsed().as_secs_f64() * 1000.0 + ); + // Never entered, so it lasts until it closes: when the borrow and + // the spans it parents are done. + let connection_span = tracing::info_span!("db.connection"); self.update_stats(|stats| { stats.checkouts += 1; stats.checkout_time += start.elapsed(); @@ -272,6 +289,8 @@ impl<'a> SyncReader<'a> { Ok(Borrowed { reader: self, conn: Some(conn), + connection_span, + txn_span: None, in_transaction: false, pending_begin: Mutex::new(None), begin_failed: AtomicBool::new(false), @@ -285,6 +304,7 @@ impl<'a> SyncReader<'a> { /// [`Borrowed::commit`]. pub(crate) async fn begin(&self) -> anyhow::Result> { let mut borrowed = self.borrow().await?; + borrowed.begin_txn_span("sync", true); // 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. @@ -303,6 +323,7 @@ impl<'a> SyncReader<'a> { /// with [`Borrowed::commit`]. pub(crate) async fn begin_write(&self) -> anyhow::Result> { let mut borrowed = self.borrow().await?; + borrowed.begin_txn_span("sync_write", false); // 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. @@ -2252,6 +2273,23 @@ impl Borrowed<'_, '_> { self.conn.as_ref().expect("held until dropped") } + /// Open the `db.txn` span for the transaction about to begin. + fn begin_txn_span(&mut self, desc: &'static str, readonly: bool) { + let txn_id = next_txn_id(); + self.txn_span = Some(tracing::info_span!( + parent: &self.connection_span, + "db.txn", + db.txn_desc = desc, + db.txn_id = %format_args!("{desc}-{txn_id:x}"), + db.readonly = readonly, + )); + } + + /// The parent of this connection's `db.query` spans. + fn query_parent(&self) -> &tracing::Span { + self.txn_span.as_ref().unwrap_or(&self.connection_span) + } + /// 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. @@ -2260,6 +2298,8 @@ impl Borrowed<'_, '_> { request: impl std::future::Future>, ) -> anyhow::Result> { let begin = lock_ignoring_poison(&self.pending_begin).take(); + // Called inside the statement's `db.query` span. + tracing::Span::current().record("db.begin", begin.is_some()); let Some(begin) = begin else { return Ok(request.await); }; @@ -2280,19 +2320,25 @@ impl Borrowed<'_, '_> { ) -> anyhow::Result> { let start = Instant::now(); let conn = self.conn(); + let span = db_query_span(Some(self.query_parent()), sql); // 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 + let rows = traced_query(span, async { + let cached = conn.cached(sql); + tracing::Span::current().record("db.prepared", cached.is_none()); + match cached { + 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()))?; + .with_context(|| format!("running {}", sql.trim())) + }) + .await?; self.record(start, rows.len()); if log::log_enabled!(log::Level::Debug) { let head: String = sql.split_whitespace().take(6).collect::>().join(" "); @@ -2312,10 +2358,13 @@ impl Borrowed<'_, '_> { sql: &str, ) -> anyhow::Result>> { let start = Instant::now(); - let messages = self - .after_begin(self.conn().simple_query(sql)) - .await? - .with_context(|| format!("running {}", sql.trim()))?; + let span = db_query_span(Some(self.query_parent()), sql); + let messages = traced_query(span, async { + self.after_begin(self.conn().simple_query(sql)) + .await? + .with_context(|| format!("running {}", sql.trim())) + }) + .await?; let mut results: Vec> = Vec::new(); let mut current = Vec::new(); let mut rows = 0; @@ -2336,9 +2385,13 @@ impl Borrowed<'_, '_> { /// Run statements that return nothing. pub(crate) async fn execute(&self, sql: &str) -> anyhow::Result<()> { let start = Instant::now(); - self.after_begin(self.conn().batch_execute(sql)) - .await? - .with_context(|| format!("running {sql}"))?; + let span = db_query_span(Some(self.query_parent()), sql); + traced_query(span, async { + self.after_begin(self.conn().batch_execute(sql)) + .await? + .with_context(|| format!("running {sql}")) + }) + .await?; self.record(start, 0); Ok(()) } @@ -2359,6 +2412,9 @@ impl Borrowed<'_, '_> { self.in_transaction = false; } if self.in_transaction { + if let Some(txn_span) = &self.txn_span { + tracing::info!(parent: txn_span, message = "commit"); + } self.execute("COMMIT").await?; self.in_transaction = false; } @@ -2383,12 +2439,16 @@ impl Drop for Borrowed<'_, '_> { // do that on, the connection is dropped rather than returned. match Handle::try_current() { Ok(handle) => { - handle.spawn(async move { - if let Err(e) = conn.batch_execute("ROLLBACK").await { - warn!("Failed to roll back a sync transaction: {e}"); + let span = tracing::info_span!(parent: self.query_parent(), "db.rollback"); + handle.spawn( + async move { + if let Err(e) = conn.batch_execute("ROLLBACK").await { + warn!("Failed to roll back a sync transaction: {e}"); + } + drop(conn); } - drop(conn); - }); + .instrument(span), + ); } Err(_) => { drop(PooledConnection::take(conn)); @@ -3020,3 +3080,132 @@ fn events_to_py(py: Python<'_>, events: Vec) -> PyResult>> pub fn register_test_classes(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::() } + +#[cfg(test)] +mod tracing_tests { + //! Against a real server: these run when `SYNAPSE_TEST_POSTGRES_DSN` is + //! set, and otherwise do nothing. + + use super::*; + use crate::caches::registry::CacheRegistry; + use crate::logging::opentracing::layer::testing::{tag, traced}; + use crate::logging::opentracing::thrift::{FinishedSpan, TagValue}; + use crate::logging::opentracing::SamplerConfig; + use crate::storage::caches::StoreCaches; + use crate::storage::pool_querier::tests::test_pool; + + fn store() -> SyncStore { + let registry = CacheRegistry::new(); + SyncStore::new( + Arc::new(StoreCaches::new(®istry, 10)), + SyncStoreConfig { + retention_enabled: false, + retention_default_min_lifetime: None, + retention_default_max_lifetime: None, + }, + ) + } + + fn named<'a>(spans: &'a [FinishedSpan], name: &str) -> Vec<&'a FinishedSpan> { + spans.iter().filter(|span| span.operation == name).collect() + } + + /// A committed transaction: `db.connection` → `db.txn` → a `db.query` + /// carrying the BEGIN, then the COMMIT's. + #[test] + fn traced_transaction() { + let Some((rt, pool)) = test_pool() else { + return; + }; + let store = store(); + let (_, spans) = traced(SamplerConfig::Const(true), || { + rt.block_on(async { + let reader = store.reader(pool.clone()); + let txn = reader.begin().await.unwrap(); + txn.query("SELECT 1", &[]).await.unwrap(); + txn.commit().await.unwrap(); + }) + }); + + let [connection] = named(&spans, "db.connection")[..] else { + panic!("expected one db.connection: {spans:#?}") + }; + let [txn] = named(&spans, "db.txn")[..] else { + panic!("expected one db.txn: {spans:#?}") + }; + assert_eq!(txn.parent_id, connection.span_id); + assert_eq!(tag(txn, "db.txn_desc"), Some(&TagValue::Str("sync".into()))); + assert_eq!(tag(txn, "db.readonly"), Some(&TagValue::Bool(true))); + assert!(matches!( + tag(txn, "db.txn_id"), + Some(TagValue::Str(id)) if id.starts_with("sync-") + )); + // The connection span lasts from checkout until the borrow is done. + assert!(connection.start + connection.duration >= txn.start + txn.duration); + assert_eq!(txn.logs.len(), 1, "the commit log"); + + let queries = named(&spans, "db.query"); + let [select, commit] = &queries[..] else { + panic!("expected two db.query spans: {spans:#?}") + }; + for query in [select, commit] { + assert_eq!(query.parent_id, txn.span_id); + } + assert_eq!(tag(select, "db.begin"), Some(&TagValue::Bool(true))); + assert_eq!(tag(commit, "db.begin"), Some(&TagValue::Bool(false))); + assert_eq!( + tag(commit, "db.statement"), + Some(&TagValue::Str("COMMIT".into())) + ); + } + + /// A transaction dropped without committing rolls back in a + /// `db.rollback` span; a one-shot query has no `db.txn`. + #[test] + fn traced_rollback_and_one_shot() { + let Some((rt, pool)) = test_pool() else { + return; + }; + let store = store(); + let (_, spans) = traced(SamplerConfig::Const(true), || { + rt.block_on(async { + let reader = store.reader(pool.clone()); + let txn = reader.begin_write().await.unwrap(); + txn.query("SELECT 1", &[]).await.unwrap(); + drop(txn); + reader.query_one_shot("SELECT 2", &[]).await.unwrap(); + // Wait for the rollback, on its own task, to give its + // connection back. + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5); + loop { + let status = pool.status(); + if status.available == status.size { + break; + } + assert!( + std::time::Instant::now() < deadline, + "rollback never finished" + ); + tokio::time::sleep(std::time::Duration::from_millis(5)).await; + } + }) + }); + + let [txn] = named(&spans, "db.txn")[..] else { + panic!("expected one db.txn: {spans:#?}") + }; + assert_eq!(tag(txn, "db.readonly"), Some(&TagValue::Bool(false))); + let [rollback] = named(&spans, "db.rollback")[..] else { + panic!("expected one db.rollback: {spans:#?}") + }; + assert_eq!(rollback.parent_id, txn.span_id); + + let connections = named(&spans, "db.connection"); + assert_eq!(connections.len(), 2); + let one_shot = named(&spans, "db.query") + .into_iter() + .find(|query| tag(query, "db.statement") == Some(&TagValue::Str("SELECT 2".into()))) + .expect("the one-shot query"); + assert!(connections.iter().any(|c| c.span_id == one_shot.parent_id)); + } +}