From a7ffaa582a87f7edc2f3d3aaa22393842a2f2946 Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Mon, 19 Oct 2026 11:14:33 +0000 Subject: [PATCH] Trace Rust state resolution, event loads and HTTP requests - `state.resolve` around a state resolution on the database, tagged `room_id`, `state_sets` and, once done, the resolution's database and auth chain time. It covers the handler's resolutions and the persister's, under `persist.state_res`. A room without a chain cover index is tagged `fallback`, not an error: the caller resolves its state in Python. - `events.load` around an event load, tagged with the number of events asked for, `reload`, and on success how many this load read from the database (`fetched`) and how many it did not (`cached`, which counts events another load was already reading). Its `db.*` spans nest under it. - `outgoing-client-request` around a request of the Rust HTTP client (MAS introspection), with the tags Python's `SimpleHttpClient` span has: `span.kind`, `http.method`, `http.url` (without credentials or query) and `http.status_code`. As in Python, an error status is the caller's to judge, so only a failure to get a response is marked as an error, before the conversion to a Python exception. A trial test loads an uncached event with tracing on and checks that `events.load` is under the caller's Python span, with the event fetched and its reads under it. Rust tests check `state.resolve` on a failed resolution and the HTTP client span against a local server. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01W4zxygfJC93SsugxskhoJG --- rust/src/http_client.rs | 174 ++++++++++++++++++++++++----- rust/src/state/mod.rs | 52 +++++++++ rust/src/state/store.rs | 70 ++++++++++++ rust/src/storage/python.rs | 20 ++++ tests/logging/test_rust_tracing.py | 62 ++++++++++ 5 files changed, 350 insertions(+), 28 deletions(-) diff --git a/rust/src/http_client.rs b/rust/src/http_client.rs index 802f93b2d0..673112bc61 100644 --- a/rust/src/http_client.rs +++ b/rust/src/http_client.rs @@ -18,9 +18,11 @@ use anyhow::Context; use http_body_util::BodyExt; use pyo3::prelude::*; use reqwest::RequestBuilder; +use tracing::Instrument; use crate::deferred::create_deferred; use crate::errors::HttpResponseException; +use crate::logging::opentracing::record_error; use crate::runtime::RustRuntime; /// Called when registering modules with python. @@ -104,35 +106,151 @@ impl HttpClient { builder: RequestBuilder, response_limit: usize, ) -> PyResult> { + let client = self.client.clone(); create_deferred(py, &self.runtime, async move { - let response = builder.send().await.context("sending request")?; - - let status = response.status(); - - // A light-weight way to read the response up until the `response_limit`. We - // want to avoid allocating a giant response object on the server above our - // expected `response_limit` to avoid out-of-memory DOS problems. - let body = reqwest::Body::from(response); - let limited_body = http_body_util::Limited::new(body, response_limit); - let collected = limited_body - .collect() - .await - .map_err(anyhow::Error::from_boxed) - .with_context(|| { - format!( - "Response body exceeded response limit ({} bytes)", - response_limit - ) - })?; - let bytes: bytes::Bytes = collected.to_bytes(); - - if !status.is_success() { - return Err(HttpResponseException::new(status, bytes)); - } - - // Because of the `pyo3` `bytes` feature, we can pass this back to Python - // land efficiently - Ok(bytes) + let request = builder.build().context("building request")?; + traced_request(client, request, response_limit).await }) } } + +/// Send `request` in Python's `SimpleHttpClient` span, created in the +/// task so that the caller's Python span is its parent. +async fn traced_request( + client: reqwest::Client, + request: reqwest::Request, + response_limit: usize, +) -> PyResult { + // Without credentials or the query, which may carry secrets. + let mut url = request.url().clone(); + let _ = url.set_password(None); + let _ = url.set_username(""); + url.set_query(None); + let span = tracing::info_span!( + "outgoing-client-request", + span.kind = "client", + http.method = request.method().as_str(), + http.url = url.as_str(), + http.status_code = tracing::field::Empty, + error = tracing::field::Empty, + ); + let result = read_response(client, request, response_limit) + .instrument(span.clone()) + .await; + match result { + Ok(bytes) => Ok(bytes), + // As Python's span, an error status is not a failure of the + // request: the caller decides. + Err(ReadError::Status(status, bytes)) => Err(HttpResponseException::new(status, bytes)), + Err(ReadError::Other(err)) => { + // Before the conversion to a Python error, which needs the GIL. + record_error(&span, &format!("{err:#}")); + Err(err.into()) + } + } +} + +/// Why [`read_response`] failed. +enum ReadError { + /// The server answered with an error status. + Status(reqwest::StatusCode, bytes::Bytes), + Other(anyhow::Error), +} + +impl From for ReadError { + fn from(err: anyhow::Error) -> Self { + ReadError::Other(err) + } +} + +/// Send `request` and read its response, up to `response_limit` bytes. +async fn read_response( + client: reqwest::Client, + request: reqwest::Request, + response_limit: usize, +) -> Result { + let response = client.execute(request).await.context("sending request")?; + + let status = response.status(); + tracing::Span::current().record("http.status_code", status.as_u16()); + + // A light-weight way to read the response up until the `response_limit`. We + // want to avoid allocating a giant response object on the server above our + // expected `response_limit` to avoid out-of-memory DOS problems. + let body = reqwest::Body::from(response); + let limited_body = http_body_util::Limited::new(body, response_limit); + let collected = limited_body + .collect() + .await + .map_err(anyhow::Error::from_boxed) + .with_context(|| { + format!( + "Response body exceeded response limit ({} bytes)", + response_limit + ) + })?; + let bytes: bytes::Bytes = collected.to_bytes(); + + if !status.is_success() { + return Err(ReadError::Status(status, bytes)); + } + + // Because of the `pyo3` `bytes` feature, we can pass this back to Python + // land efficiently + Ok(bytes) +} + +#[cfg(test)] +mod tests { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + use super::*; + use crate::logging::opentracing::layer::testing::{tag, traced_until_dropped}; + use crate::logging::opentracing::thrift::TagValue; + use crate::logging::opentracing::SamplerConfig; + + /// The request is Python's `outgoing-client-request` span, with its + /// status, and a URL without credentials or query. + #[tokio::test] + async fn traced_request_span() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut buf = [0; 4096]; + let _ = socket.read(&mut buf).await.unwrap(); + socket + .write_all(b"HTTP/1.1 200 OK\r\ncontent-length: 2\r\nconnection: close\r\n\r\nok") + .await + .unwrap(); + }); + + let (guard, sink) = traced_until_dropped(SamplerConfig::Const(true)); + let client = reqwest::Client::builder().no_proxy().build().unwrap(); + let request = client + .get(format!( + "http://user:secret@127.0.0.1:{port}/introspect?token=abc" + )) + .build() + .unwrap(); + let body = traced_request(client, request, 1024).await.unwrap(); + drop(guard); + server.await.unwrap(); + + assert_eq!(&body[..], b"ok"); + let spans = sink.take(); + let [span] = &spans[..] else { + panic!("expected one span: {spans:#?}") + }; + assert_eq!(span.operation, "outgoing-client-request"); + assert_eq!(tag(span, "http.method"), Some(&TagValue::Str("GET".into()))); + assert_eq!( + tag(span, "http.url"), + Some(&TagValue::Str( + format!("http://127.0.0.1:{port}/introspect").into() + )) + ); + assert_eq!(tag(span, "http.status_code"), Some(&TagValue::Long(200))); + assert_eq!(tag(span, "error"), None); + } +} diff --git a/rust/src/state/mod.rs b/rust/src/state/mod.rs index a5c05eeabd..378f145f04 100644 --- a/rust/src/state/mod.rs +++ b/rust/src/state/mod.rs @@ -45,12 +45,14 @@ use pyo3::{ prelude::*, types::{PyDict, PyModule, PyTuple}, }; +use tracing::Instrument; use crate::caches::{python::PyCacheKey, Cache}; use crate::database::postgres::pool::{ConnectionPool, PyConnectionPool}; use crate::deferred::create_deferred; use crate::event_auth::AuthFailure; use crate::events::Event; +use crate::logging::opentracing::record_error; use crate::room_versions::{RoomVersion, StateResolutionVersions}; use crate::runtime::RustRuntime; use crate::storage::{caches::StoreCaches, store::Store}; @@ -700,6 +702,56 @@ pub(crate) async fn resolve_on_pool( separated: (StateMap, v2::Conflicted), events: EventMap, stats: StateResStats, +) -> anyhow::Result { + let span = tracing::info_span!( + "state.resolve", + room_id = room_id, + state_sets = state_sets.len() as u64, + db_ms = tracing::field::Empty, + auth_chain_ms = tracing::field::Empty, + fallback = tracing::field::Empty, + error = tracing::field::Empty, + ); + let output = resolve_on_pool_in_span( + pool, + handle, + caches, + cache, + room_id, + room_version, + state_sets, + separated, + events, + stats, + ) + .instrument(span.clone()) + .await; + match &output { + Ok(output) => { + span.record("db_ms", output.stats.db_time * 1000.0); + span.record("auth_chain_ms", output.stats.auth_chain_time * 1000.0); + } + // Expected: the caller resolves the state in Python instead. + Err(err) if err.downcast_ref::().is_some() => { + span.record("fallback", true); + } + Err(err) => record_error(&span, err), + } + output +} + +#[allow(clippy::too_many_arguments)] +async fn resolve_on_pool_in_span( + pool: &ConnectionPool, + handle: tokio::runtime::Handle, + caches: Arc, + cache: &ConflictCache, + room_id: &str, + room_version: &'static RoomVersion, + state_sets: &[Arc], + separated: (StateMap, v2::Conflicted), + events: EventMap, + stats: StateResStats, ) -> anyhow::Result { let start = Instant::now(); let store = PostgresStateResStore::begin(pool, handle, Some(caches)).await?; diff --git a/rust/src/state/store.rs b/rust/src/state/store.rs index 7935e6c5a7..cf25c1db62 100644 --- a/rust/src/state/store.rs +++ b/rust/src/state/store.rs @@ -576,4 +576,74 @@ mod tracing_tests { ] ); } + + /// A resolution that fails reports `state.resolve` marked failed, with + /// its reads under it. + #[test] + fn traced_resolve_failure() { + use crate::caches::registry::CacheRegistry; + use crate::room_versions::RoomVersion; + use crate::state::{conflict_cache_for_test, resolve_on_pool, v2, StateKey, StateMap}; + use crate::storage::caches::StoreCaches; + + let Some((rt, pool)) = test_pool() else { + return; + }; + let registry = CacheRegistry::new(); + let caches = Arc::new(StoreCaches::new(®istry, 10)); + let cache = conflict_cache_for_test(10); + let key = StateKey::new("m.room.name", ""); + let state_sets: Vec> = ["$a", "$b"] + .into_iter() + .map(|event_id| Arc::new(StateMap::from([(key.clone(), Arc::from(event_id))]))) + .collect(); + let separated = v2::separate(&state_sets); + + let (result, spans) = traced(SamplerConfig::Const(true), || { + let result = rt.block_on(resolve_on_pool( + &pool, + rt.handle().clone(), + caches, + &cache, + "!unknown:test", + &RoomVersion::V11, + &state_sets, + separated, + Default::default(), + Default::default(), + )); + // The failed store rolls back on a task of its own, which holds + // the spans open: wait for it to give the connection back. + rt.block_on(async { + 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; + } + }); + result + }); + assert!(result.is_err()); + + let resolve = spans + .iter() + .find(|span| span.operation == "state.resolve") + .expect("a state.resolve span"); + assert_eq!( + tag(resolve, "room_id"), + Some(&TagValue::Str("!unknown:test".into())) + ); + assert_eq!(tag(resolve, "state_sets"), Some(&TagValue::Long(2))); + assert_eq!(tag(resolve, "error"), Some(&TagValue::Bool(true))); + assert!(spans + .iter() + .any(|span| span.operation == "db.connection" && span.parent_id == resolve.span_id)); + } } diff --git a/rust/src/storage/python.rs b/rust/src/storage/python.rs index 7da9e28837..b4e6fe8f59 100644 --- a/rust/src/storage/python.rs +++ b/rust/src/storage/python.rs @@ -56,7 +56,9 @@ use pyo3::{ }, Bound, IntoPyObject, Py, PyAny, PyClassInitializer, PyErr, PyResult, Python, }; +use tracing::Instrument; +use crate::logging::opentracing::record_error; use crate::{ caches::{ python::{PyCacheKey, RustCache}, @@ -213,6 +215,14 @@ impl RustStore { let ids: Vec> = event_ids.into_iter().map(Arc::from).collect(); Ok(async move { let mut stats = EventLoadStats::default(); + let span = tracing::info_span!( + "events.load", + events = ids.len() as u64, + reload, + cached = tracing::field::Empty, + fetched = tracing::field::Empty, + error = tracing::field::Empty, + ); let events = async { let querier = PoolQuerier::checkout(&pool).await?; let loader = EventLoader::new(&querier, Some(&store.caches.events)); @@ -222,7 +232,17 @@ impl RustStore { loader.load(&ids, &mut stats).await } } + .instrument(span.clone()) .await; + match &events { + Ok(_) => { + // Events another caller's load read count as cached: + // this load did not read them. + span.record("cached", stats.cached); + span.record("fetched", stats.fetched); + } + Err(err) => record_error(&span, err), + } events .map(|events| LoadedEvents { events, diff --git a/tests/logging/test_rust_tracing.py b/tests/logging/test_rust_tracing.py index 6ef95358aa..1485112f68 100644 --- a/tests/logging/test_rust_tracing.py +++ b/tests/logging/test_rust_tracing.py @@ -676,3 +676,65 @@ class SyncTracingTestCase(_AgentTestCase): # Encoded on the reactor, under the request's Python span. self.assertEqual(encode.traceIdLow, entries.traceIdLow) self.assertIn("SyncRestServlet", self._python_ancestors(encode, python_by_id)) + + +class EventLoadTracingTestCase(_AgentTestCase): + """An event load in Rust reports `events.load` under the caller's Python + span, with the event fetched from the database. + + Needs the native Rust Postgres driver: + + SYNAPSE_TEST_RUST_DRIVER=1 POSTGRES=1 run-trial.sh tests.logging.test_rust_tracing + """ + + servlets = [ + admin.register_servlets, + room.register_servlets, + login.register_servlets, + ] + + def make_homeserver(self, reactor: MemoryReactor, clock: Clock) -> HomeServer: + hs = self.setup_test_homeserver() + if hs.get_datastores().main._rust_event_loader is None: + raise SkipTest("needs the native Rust Postgres driver") + return hs + + def prepare_before_tracing(self) -> None: + self.user = self.register_user("user", "pass") + self.token = self.login("user", "pass") + room_id = self.helper.create_room_as(self.user, tok=self.token) + self.event_id = self.helper.send(room_id, "hello", tok=self.token)["event_id"] + + def test_event_load_span(self) -> None: + store = self.hs.get_datastores().main + # Out of the cache, so that the load reads the database. + store._invalidate_local_get_event_cache(self.event_id) + + tracer = opentracing.global_tracer() + with LoggingContext(name="ctx", server_name="test_server"): + with start_active_span("load", tracer=tracer) as scope: + self.get_success(store.get_event(self.event_id)) + + spans = self._received_spans() + trace_id = _signed(_context(scope.span).trace_id) + loads = [ + s + for s in spans + if s.operationName == "events.load" and s.traceIdLow == trace_id + ] + self.assertEqual(len(loads), 1, [s.operationName for s in spans]) + (load,) = loads + # Under the caller's span, through the store's own traced calls. + ancestors = self._python_ancestors(load, self._python_spans_by_id()) + self.assertEqual(ancestors[-1], "load", ancestors) + tags = {t.key: t.vLong for t in load.tags if t.vType == ttypes.TagType.LONG} + self.assertEqual(tags["fetched"], 1) + self.assertEqual(tags["events"], 1) + + # Its reads nest under it. + self.assertTrue( + any( + s.operationName == "db.connection" and s.parentSpanId == load.spanId + for s in spans + ) + )